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:
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user