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>

crates/xtrain-distributed/src/ddp.rs

@@ -12,6 +12,7 @@
 //! is exactly the single-GPU batch in the same order, so the all-reduced grad sum
 //! equals the single-GPU summed grad.
 
+use std::path::PathBuf;
 use std::thread;
 use std::time::Instant;
 
@@ -19,8 +20,10 @@ use xtrain_autodiff::tape::Var;
 use xtrain_model::{Config, TinyTransformer, ids_tensor};
 use xtrain_optim::GpuAdamW;
 use xtrain_tensor::Device;
+use xtrain_train::checkpoint;
 use xtrain_train::clip::clip_grad_norm_gpu;
 use xtrain_train::data::Corpus;
+use xtrain_train::eval_loss;
 use xtrain_train::schedule::LrSchedule;
 
 use crate::{DdpContext, get_unique_id};
@@ -38,20 +41,43 @@ pub struct DdpConfig {
     pub max_grad_norm: f32,
     pub log_every: usize,
     pub seed: u64,
+    /// Evaluate held-out val loss every `eval_every` steps (0 = never). Only rank
+    /// 0 holds the `valid` corpus and runs the eval (no grad), mirroring
+    /// `xtrain_train::TrainConfig`. The best-val model is checkpointed by rank 0
+    /// (every rank's params are identical, so rank 0's are the model's).
+    pub eval_every: usize,
+    pub eval_batches: usize,
+    /// Best-val checkpoint path (written by rank 0 when val improves). When unset,
+    /// or when `eval_every == 0`, no checkpoint is written.
+    pub ckpt_path: Option<PathBuf>,
+}
+
+/// Outcome of a DDP run on this rank: per-step mean-loss trace plus, when
+/// `eval_every > 0`, the (step, val_loss) eval points and the best val loss
+/// (eval/best are only populated on rank 0, which owns the `valid` corpus).
+pub struct DdpResult {
+    pub losses: Vec<f32>,
+    pub evals: Vec<(usize, f32)>,
+    pub best_val: Option<f32>,
 }
 
 /// Run `cfg.steps` DDP steps on this rank's `model`/`corpus`, using `ctx` for the
 /// gradient all-reduce. Returns this rank's per-step mean-loss trace (the mean
 /// over the GLOBAL batch — every rank computes the same value because losses are
-/// all-reduced alongside the grads). The optimizer step is identical on every
-/// rank, so the parameters stay in lockstep.
+/// all-reduced alongside the grads) plus eval/best-val (rank 0 only). The
+/// optimizer step is identical on every rank, so the parameters stay in lockstep.
+///
+/// `valid` is the held-out corpus for periodic val-loss eval. Only rank 0 needs
+/// it (it runs the no-grad eval and writes the best-val checkpoint); pass `None`
+/// on the other ranks (or when `cfg.eval_every == 0`).
 pub fn train_rank(
     ctx: &DdpContext,
     model: &TinyTransformer,
     device: Device,
     corpus: &Corpus,
+    valid: Option<&Corpus>,
     cfg: &DdpConfig,
-) -> Vec<f32> {
+) -> DdpResult {
     assert_eq!(
         cfg.batch_size % ctx.world,
         0,
@@ -63,12 +89,17 @@ pub fn train_rank(
     let mut opt = GpuAdamW::new(cfg.weight_decay);
     let mut rng = cfg.seed;
     let mut losses = Vec::with_capacity(cfg.steps);
+    let mut evals = Vec::new();
+    let mut best_val: Option<f32> = None;
     // Each rank reaches the global batch mean as (Σ_global / world) · (1/b_local),
     // where b_local = batch_size / world (see DdpContext::all_reduce_average_grads).
     let batch_local = cfg.batch_size / ctx.world;
     let inv_batch_local = 1.0 / batch_local as f32;
     let start = Instant::now();
     let mut tokens_seen: u64 = 0;
+    // Rank 0 owns the held-out eval + best-val checkpoint (params are identical
+    // across ranks, so rank 0's are the model). Other ranks never touch `valid`.
+    let do_eval = ctx.rank == 0 && cfg.eval_every > 0 && valid.is_some();
 
     for step in 0..cfg.steps {
         let lr = cfg.schedule.lr(step);
@@ -114,19 +145,52 @@ pub fn train_rank(
                 cfg.steps, ctx.world
             );
         }
+
+        // Periodic held-out eval + best-val checkpoint (rank 0 only). Mirrors the
+        // single-GPU `xtrain_train::train` loop, reusing its `eval_loss` /
+        // `checkpoint::save` so single-GPU and DDP share one eval/ckpt path. Other
+        // ranks have nothing to do here (params are identical across ranks).
+        if do_eval && ((step + 1) % cfg.eval_every == 0 || step == cfg.steps - 1) {
+            let v = valid.unwrap();
+            let vl = eval_loss(model, device, v, cfg.seq_len, cfg.eval_batches);
+            evals.push((step, vl));
+            let improved = best_val.map(|b| vl < b).unwrap_or(true);
+            println!(
+                "  [rank0] eval @ step {step}: val loss {vl:.4}{}",
+                if improved { " (best)" } else { "" }
+            );
+            if improved {
+                best_val = Some(vl);
+                if let Some(path) = &cfg.ckpt_path {
+                    checkpoint::save(path, &params).expect("best checkpoint save");
+                }
+            }
+        }
+    }
+    DdpResult {
+        losses,
+        evals,
+        best_val,
     }
-    losses
 }
 
 /// Spawn `world` rank threads (one per GPU in `devices`), init NCCL, build an
 /// identical model per rank via `make_model`, and run `train_rank`. Returns each
-/// rank's loss trace (all identical). The launcher owns the thread-per-GPU model:
-/// rank 0 mints the `UniqueId`, every thread `cudaSetDevice`s its GPU, builds its
-/// `Var` graph locally (the graph is `!Send`), and joins at the end.
+/// rank's `DdpResult` (loss traces are identical; eval/best-val are on rank 0).
+/// The launcher owns the thread-per-GPU model: rank 0 mints the `UniqueId`, every
+/// thread `cudaSetDevice`s its GPU, builds its `Var` graph locally (the graph is
+/// `!Send`), and joins at the end.
 ///
-/// `make_model(device)` must be deterministic — same params on every rank — for
-/// the parameters to stay consistent.
-pub fn launch<F>(devices: &[u32], corpus: &Corpus, cfg: &DdpConfig, make_model: F) -> Vec<Vec<f32>>
+/// `valid` is the held-out corpus for rank 0's periodic eval (only used when
+/// `cfg.eval_every > 0`). `make_model(device)` must be deterministic — same params
+/// on every rank — for the parameters to stay consistent.
+pub fn launch<F>(
+    devices: &[u32],
+    corpus: &Corpus,
+    valid: Option<&Corpus>,
+    cfg: &DdpConfig,
+    make_model: F,
+) -> Vec<DdpResult>
 where
     F: Fn(Device) -> TinyTransformer + Send + Sync,
 {
@@ -144,7 +208,9 @@ where
                     let ctx = DdpContext::init(rank, world, id, dev);
                     let device = Device::Cuda(dev);
                     let model = make_model(device);
-                    train_rank(&ctx, &model, device, corpus, &cfg)
+                    // Only rank 0 holds the val corpus for eval.
+                    let v = if rank == 0 { valid } else { None };
+                    train_rank(&ctx, &model, device, corpus, v, &cfg)
                 })
             })
             .collect();