Repository navigation
feat: 🎸 add vit_tp_exp: DeiT-S tensor-parallel example, cross-checked… - #30
Merged
Merged
Conversation
… L1-L4 against official DeiT - Port the DeiT-S model verbatim from timm 0.3.2 plus the facebookresearch/deit recipe (train/eval loops, transforms, RASampler, MetricLogger) onto the tinyexp Cfg/Exp layout; attn.qkv split (D1) with bidirectional weight converters, checkpoints keep the official fused-qkv layout and load official DeiT checkpoints directly. - TPAccelerator on torch.distributed.tensor.parallel: Megatron-style colwise/rowwise plan sharding attention on head boundaries; numerical equivalence vs single-rank asserted in CI (gloo), nccl+AMP on GPU. - Full official recipe wiring (300ep AdamW/cosine, mixup+cutmix, repeated augmentation, EMA off at tp>1), mode=bench throughput/memory check, per-epoch checkpoints with resume, optional Redis raw-byte cache for slow shared filesystems (D7, bit-identical samples). - Cross-check record in docs/vit_tp.md: L1 numerics (CI allclose); L2 official-ckpt eval 79.82/79.81 top-1 at tp=1/tp=2 (official 79.8); L3 full 300-epoch training from scratch on 8xH200 (rjob + redis cache): 79.84 top-1 / 94.99 top-5 in 17h56m (paper 79.8 / ~95.0); L4 TP-vs-DDP bench table. - Engine hardening, each with a regression test: DDPAccelerator.reduce_sum stages cpu tensors through the device for nccl; store_and_run_exp resolves the canonical importable twin of a python -m __main__ exp class so ray ships it by reference. Co-Authored-By: Claude Code <[email protected]>
Avoid mutating same-device tensors during all-reduce while preserving CPU metric support.
Document portable ImageNet and DDP commands directly in the example and remove platform-specific helper scripts.
Reuse timm VisionTransformer, Block, and Attention while keeping split QKV projections for tensor parallel sharding. Update documentation and checkpoint loading accordingly.
Reuse timm's DeiT-compatible RepeatAugSampler instead of maintaining a duplicate implementation. Split train and validation worker settings, enable persistent workers when applicable, and remove stale sampler files and lint rules.
Keep the previous positive-repeat validation after switching to timm's sampler, preventing invalid configurations from silently producing no training samples.
Keep one focused ViT test module with model equivalence, QKV conversion, Redis gating, and TP numerical checks. Remove process-oriented smoke tests and reuse timm's transformer block directly.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
… L1-L4 against official DeiT