feat: 🎸 add vit_tp_exp: DeiT-S tensor-parallel example, cross-checked… - #30
Open
zengarden wants to merge 4 commits into
Open
feat: 🎸 add vit_tp_exp: DeiT-S tensor-parallel example, cross-checked…#30zengarden wants to merge 4 commits into
zengarden wants to merge 4 commits into
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 <noreply@anthropic.com>
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.
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