test: T21 — DDP-dropout regression (live under DDP + p=0 bit-identical)

Adds ddp_dropout_is_live_and_p0_bit_identical, run via the real launcher
path (train_rank, world=2). It would have caught the original bug:

- GATE A: a p=0 DDP run is BIT-IDENTICAL (loss trace + final params) to the
  no-dropout path — the regression guard for --dropout 0 (default).
- GATE B: a p=0.2 DDP run's loss trace DIFFERS measurably (>1e-3) from p=0.
  On the pre-T21 code the model stays in eval mode, so p=0.2 would be an
  identity and the trace would be bit-identical to p=0 — this gate fails.
- GATE C: model.is_training() == true after the run (direct proof that
  train_rank called model.train() and it survived the final-step eval).
- p>0 run is finite (no NaN/Inf).

eval_every < steps so a periodic eval fires mid-run (flipping to eval mode),
exercising the per-step model.train() restore discipline the pilot called out.
Run with --test-threads=1 like the other DDP tests (shared-GPU deadlock).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-distributed/tests/ddp_correctness.rs

@@ -38,6 +38,46 @@ fn test_config(vocab: usize) -> Config {
     cfg
 }
 
+/// Run `cfg`/`dcfg` as a 2-rank DDP job (the same launcher path as production) and
+/// return rank 0's (loss trace, final params on host, final `is_training()` flag).
+/// `cfg` carries the dropout prob; `dcfg` carries the loop knobs. Caller asserts.
+fn run_ddp2(
+    cfg: Config,
+    corpus: &Corpus,
+    valid: Option<&Corpus>,
+    dcfg: &DdpConfig,
+) -> (Vec<f32>, Vec<Vec<f32>>, bool) {
+    let world = 2usize;
+    let devices = [0u32, 1u32];
+    let id = get_unique_id();
+    let results: Vec<(Vec<f32>, Vec<Vec<f32>>, bool)> = std::thread::scope(|s| {
+        let handles: Vec<_> = devices
+            .iter()
+            .enumerate()
+            .map(|(rank, &dev)| {
+                let dcfg = dcfg.clone();
+                let corpus = &corpus;
+                s.spawn(move || {
+                    let ctx = DdpContext::init(rank, world, id, dev);
+                    let device = Device::Cuda(dev);
+                    let model = build_model(cfg, device);
+                    // Only rank 0 holds the val corpus (mirrors launch()).
+                    let v = if rank == 0 { valid } else { None };
+                    let res = train_rank(&ctx, &model, device, corpus, v, &dcfg);
+                    let host = model
+                        .params()
+                        .iter()
+                        .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
+                        .collect::<Vec<_>>();
+                    (res.losses, host, model.is_training())
+                })
+            })
+            .collect();
+        handles.into_iter().map(|h| h.join().unwrap()).collect()
+    });
+    results.into_iter().next().unwrap()
+}
+
 // Single-GPU baseline: the SAME loop as the DDP rank but world=1, so the global
 // batch is processed on one device. Returns (loss trace, final params on host).
 fn run_single_gpu(cfg: Config, corpus: &Corpus, dcfg: &DdpConfig) -> (Vec<f32>, Vec<Vec<f32>>) {
@@ -386,3 +426,128 @@ fn ddp_throughput_scaling() {
         );
     }
 }
+
+/// T21 regression: prove dropout is actually LIVE under DDP, and that p=0 is
+/// bit-identical to the no-dropout path. Guards the V9-PILOT launcher-wiring gap —
+/// `train_ddp` had no `--dropout` flag and `train_rank` never called `model.train()`,
+/// so under DDP every forward ran in the default eval mode and dropout was a silent
+/// identity regardless of config. Op/single-GPU tests never exercised dropout-under-
+/// DDP, so it slipped through; this test runs the real launcher path (`train_rank`).
+///
+/// With dropout fixed across the 4 sub-runs, all three checks below would FAIL on the
+/// pre-T21 code: (a) the p>0 trace would be bit-identical to p=0 (model stuck in eval
+/// mode → identity), and (c) `is_training()` would be false after the run.
+#[test]
+fn ddp_dropout_is_live_and_p0_bit_identical() {
+    let world = 2usize;
+    if device::device_count().unwrap_or(0) < world as i32 {
+        eprintln!("skip: need >= {world} GPUs");
+        return;
+    }
+
+    let vocab = 64usize;
+    let corpus = synth_corpus(vocab, 4096);
+    let steps = 20usize;
+    // eval_every < steps so a periodic eval fires MID-run (flipping the model to
+    // eval mode via eval_loss → model.eval()). The per-step model.train() must
+    // restore training mode so dropout stays live across the eval boundary — this is
+    // exactly the train/eval discipline the pilot called out. A held-out slice gives
+    // rank 0 something to eval on.
+    let valid = synth_corpus(vocab, 512);
+    let base_dcfg = DdpConfig {
+        seq_len: 32,
+        batch_size: 8, // global; 4 per rank with world=2
+        accum_steps: 1,
+        steps,
+        schedule: LrSchedule {
+            max_lr: 3e-3,
+            min_lr: 3e-4,
+            warmup: 3,
+            total: steps,
+        },
+        weight_decay: 0.1,
+        max_grad_norm: 1.0,
+        log_every: 1_000_000, // silence per-step logging
+        seed: 7,
+        eval_every: 7, // fires at steps 6, 13, 19 — flips to eval mode mid-run
+        eval_batches: 4,
+        ckpt_path: None,
+    };
+
+    // (1) p=0 config — the no-dropout baseline. cfg.dropout defaults to 0.0.
+    let cfg_p0 = test_config(vocab);
+    assert_eq!(cfg_p0.dropout, 0.0, "baseline cfg must have dropout 0");
+    let (loss_p0, params_p0, _) = run_ddp2(cfg_p0, &corpus, Some(&valid), &base_dcfg);
+
+    // (2) Same config, dropout disabled by p=0 but explicitly set — must be the
+    //     SAME run (sanity: setting dropout=0 doesn't perturb anything).
+    let mut cfg_p0b = test_config(vocab);
+    cfg_p0b.dropout = 0.0;
+    let (loss_p0b, params_p0b, _) = run_ddp2(cfg_p0b, &corpus, Some(&valid), &base_dcfg);
+
+    // (3) Same config + data + seed, but dropout p=0.2 ON.
+    let mut cfg_p = test_config(vocab);
+    cfg_p.dropout = 0.2;
+    let (loss_p, _params_p, train_flag_p) = run_ddp2(cfg_p, &corpus, Some(&valid), &base_dcfg);
+
+    // GATE A — p=0 is bit-identical to the no-dropout path (regression guard).
+    // ops::dropout(p=0) is a clone no-op regardless of training mode, so these two
+    // runs must agree to the last bit on BOTH the loss trace and the final params.
+    let mut max_loss_diff = 0.0f32;
+    for (a, b) in loss_p0.iter().zip(&loss_p0b) {
+        max_loss_diff = max_loss_diff.max((a - b).abs());
+    }
+    let mut max_param_diff = 0.0f32;
+    for (a, b) in params_p0.iter().zip(&params_p0b) {
+        for (x, y) in a.iter().zip(b) {
+            max_param_diff = max_param_diff.max((x - y).abs());
+        }
+    }
+    println!(
+        "T21 GATE A (p=0 bit-identical): max |loss diff| = {max_loss_diff:.3e}, \
+         max |param diff| = {max_param_diff:.3e}"
+    );
+    assert_eq!(
+        max_loss_diff, 0.0,
+        "p=0 DDP loss trace not bit-identical to no-dropout path"
+    );
+    assert_eq!(
+        max_param_diff, 0.0,
+        "p=0 DDP final params not bit-identical to no-dropout path"
+    );
+
+    // GATE B — dropout is LIVE with p>0 under DDP. If model.train() were not wired
+    // (the pre-T21 bug), the model would stay in eval mode and the p=0.2 forward
+    // would be IDENTITY → loss trace bit-identical to p=0. A real, sizeable
+    // difference proves dropout masks are actually applied during the training
+    // forward (and survive the mid-run eval flips, since model.train() is re-asserted
+    // each step). Inverted scaling + masking perturbs every step, so the gap is large.
+    let mut max_live_diff = 0.0f32;
+    for (a, b) in loss_p0.iter().zip(&loss_p) {
+        max_live_diff = max_live_diff.max((a - b).abs());
+    }
+    println!(
+        "T21 GATE B (dropout live): p0[last]={:.6} p0.2[last]={:.6} max |loss diff| = {max_live_diff:.3e}",
+        loss_p0.last().unwrap(),
+        loss_p.last().unwrap()
+    );
+    assert!(
+        max_live_diff > 1e-3,
+        "p=0.2 DDP loss trace matches p=0 — dropout is NOT live under DDP \
+         (model.train() not wired): max |loss diff| {max_live_diff:.3e}"
+    );
+
+    // GATE C — train_rank leaves the model in TRAINING mode (direct proof that
+    // model.train() was called and survives the final-step eval). On the pre-T21
+    // code this would be false (model never left the default eval mode).
+    assert!(
+        train_flag_p,
+        "model not in training mode after DDP run — model.train() not wired in train_rank"
+    );
+
+    // No NaN/Inf in the p>0 run (dropout converges normally under DDP).
+    assert!(
+        loss_p.iter().all(|l| l.is_finite()),
+        "p=0.2 DDP loss has non-finite values"
+    );
+}