Skip to content

feat: 🎸 add vit_tp_exp: DeiT-S tensor-parallel example, cross-checked… - #30

Open
zengarden wants to merge 4 commits into
mainfrom
feat/vit-tp-exp
Open

feat: 🎸 add vit_tp_exp: DeiT-S tensor-parallel example, cross-checked…#30
zengarden wants to merge 4 commits into
mainfrom
feat/vit-tp-exp

Conversation

@zengarden

Copy link
Copy Markdown
Member

… 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.

Zeming LI and others added 4 commits September 11, 2026 23:14
… 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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant