Skip to content

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

Merged
zengarden merged 7 commits into
mainfrom
feat/vit-tp-exp
Sep 19, 2026
Merged

zengarden merged 7 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 7 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 <[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.
@zengarden
zengarden merged commit 6a490dc into main Sep 19, 2026
8 checks passed
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