Commit Graph

4 Commits

Author SHA1 Message Date
7090b475fb train: bring DDP trainer to parity with bin/train (val + checkpoint + cache + arch)
The T8 DDP path now matches the single-GPU `bin/train`: CLI-tunable arch
(scaling-ladder rung), the cached token-id stream (`load_cached`), held-out
val-loss eval + best-val checkpointing, and LR warmup→cosine. Rank 0 owns the
val corpus and runs the no-grad eval / writes the best checkpoint (params are
bit-identical across ranks). The eval/checkpoint logic is reused from
`xtrain-train` (`eval_loss`, `checkpoint::save`) rather than duplicated.

- DdpConfig gains eval_every / eval_batches / ckpt_path.
- train_rank takes `valid: Option<&Corpus>` and returns DdpResult
  (losses + evals + best_val); launch threads the val corpus to rank 0 only.
- bin/train_ddp reworked to the bin/train CLI (positional tokenizer/corpus +
  --dim/--heads/--head-dim/--layers/--ffn/--steps/--batch/--seq/--max-lr/
  --val-tokens/--eval-every/--ckpt), reusing the u16 cache.
- DDP correctness test updated to the new signatures (semantics unchanged).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 19:34:40 +08:00
ad82e8bf92 dist: lengthen scaling bench so NCCL init amortizes
30-step bench charged the one-time NCCL init + 4 model builds (present at world=4,
absent at world=1) against the wall clock, understating steady-state scaling
(in-loop tok/s already showed ~53k at 4 GPUs). Bump to 150 steps.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 17:18:23 +08:00
818f76a18f dist: drop unused import; relax DDP-vs-single-GPU param tolerance
dash5 verify: loss trajectory matches single-GPU to max_rel 1.16e-7 and
cross-rank params are bit-identical (0.0), but DDP-vs-single-GPU per-param rel
diff is ~2.8e-3 after 20 AdamW steps — expected, since the two differ only in
gradient summation order (fp add isn't associative) and that rounding compounds.
Bump check (c) 1e-3 -> 1e-2 (a/b stay tight). Also remove an unused DType import.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 17:17:31 +08:00
cf5e3987df dist: multi-rank launcher + ddp acceptance test
bin/train_ddp: spawn one thread per visible GPU (CUDA_VISIBLE_DEVICES selects
the set), NCCL all-reduce gradients each step, train the tiny transformer on
TinyStories; doubles as the throughput driver (prints global tok/s). no_cuda
build keeps a stub main.

tests/ddp_correctness: (1) 2-rank DDP vs single-GPU over the same synthetic data
-> loss trajectory max_rel < 1e-3, cross-rank params bit-identical (==0.0), DDP
vs single-GPU params rel < 1e-3; (2) 1/2/4-GPU throughput table on a fixed
per-GPU workload. Gated #[cfg(not(no_cuda))], auto-skips with < 2 GPUs.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 17:15:41 +08:00