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>
This commit is contained in:
2026-06-15 19:34:40 +08:00
parent 264660527f
commit 7090b475fb
4 changed files with 207 additions and 52 deletions

View File

@@ -102,6 +102,9 @@ fn ddp_matches_single_gpu_and_params_consistent() {
max_grad_norm: 1.0,
log_every: 1_000_000, // silence per-step logging in the test
seed: 7,
eval_every: 0,
eval_batches: 0,
ckpt_path: None,
};
// Single-GPU baseline (world=1) over the global batch.
@@ -121,13 +124,13 @@ fn ddp_matches_single_gpu_and_params_consistent() {
let ctx = DdpContext::init(rank, world, id, dev);
let device = Device::Cuda(dev);
let model = build_model(cfg, device);
let losses = train_rank(&ctx, &model, device, corpus, &dcfg);
let res = train_rank(&ctx, &model, device, corpus, None, &dcfg);
let host = model
.params()
.iter()
.map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
.collect::<Vec<_>>();
(losses, host)
(res.losses, host)
})
})
.collect();
@@ -224,10 +227,13 @@ fn ddp_throughput_scaling() {
max_grad_norm: 1.0,
log_every: 1_000_000,
seed: 1,
eval_every: 0,
eval_batches: 0,
ckpt_path: None,
};
let total_tokens = (steps * dcfg.batch_size * seq_len) as f64;
let t = Instant::now();
let _ = launch(&devices, &corpus, &dcfg, move |device| {
let _ = launch(&devices, &corpus, None, &dcfg, move |device| {
build_model(cfg, device)
});
let secs = t.elapsed().as_secs_f64();