diff --git a/crates/xtrain-distributed/tests/ddp_correctness.rs b/crates/xtrain-distributed/tests/ddp_correctness.rs index ebda9b0..89a1562 100644 --- a/crates/xtrain-distributed/tests/ddp_correctness.rs +++ b/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, Vec>, bool) { + let world = 2usize; + let devices = [0u32, 1u32]; + let id = get_unique_id(); + let results: Vec<(Vec, Vec>, 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::().to_vec()) + .collect::>(); + (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, Vec>) { @@ -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(¶ms_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" + ); +}