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

crates/xtrain-distributed/tests/ddp_correctness.rs

@@ -446,12 +446,15 @@ fn ddp_throughput_scaling() {
 /// bit-identical to p=0 — model stuck in eval mode → dropout is identity) and GATE C
 /// (`is_training()` would be false after the run).
 ///
-/// Bit-identity (GATE A) is asserted at `world=1`, where `all_reduce_average_grads`
-/// short-circuits (no NCCL) so the run is deterministic. The cross-rank NCCL
-/// all-reduce (`world>=2`) is not bit-reproducible run-to-run on this PCIe box (KI-5,
-/// observed ≤~2.4e-7), so the `world=2` p=0-vs-no-dropout check (GATE A2) uses the
-/// same KI-5 ULP tolerance as the rest of this file, while GATE B's live-dropout
-/// signal (>1e-3) sits orders of magnitude above that noise floor.
+/// p=0 regression (GATE A) is checked at `world=1`, ONE step, where the NCCL
+/// all-reduce short-circuits: the p=0 FORWARD is byte-identical to no-dropout so the
+/// loss is BIT-IDENTICAL (== 0.0), and the post-step params match within the engine's
+/// atomicAdd backward-reduction ULP floor (< 1e-7, dropout-independent — the
+/// fresh-train md5 caveat). The cross-rank NCCL all-reduce (`world>=2`) is not
+/// bit-reproducible run-to-run on this PCIe box (KI-5, observed ≤~2.4e-7), so the
+/// `world=2` p=0-vs-no-dropout check (GATE A2) uses the same KI-5 ULP tolerance as the
+/// rest of this file. GATE B's live-dropout signal (>1e-3) sits ~4 orders of magnitude
+/// above every noise floor here, so it carries the load.
 #[test]
 fn ddp_dropout_is_live_and_p0_bit_identical() {
     if device::device_count().unwrap_or(0) < 2 {
@@ -488,39 +491,48 @@ fn ddp_dropout_is_live_and_p0_bit_identical() {
         ckpt_path: None,
     };
 
-    // --- GATE A: bit-identity at world=1 (deterministic — no NCCL collective). ---
-    // The regression guard for `--dropout 0`: a p=0 run must be bit-for-bit the same
-    // as the no-dropout path, since ops::dropout(p=0) is a clone no-op regardless of
-    // training mode. At world=1, all_reduce_average_grads short-circuits, so the run
-    // is fully deterministic and bit-identity is the honest invariant (no NCCL noise).
+    // --- GATE A: p=0 == no-dropout at world=1, ONE step (the deterministic scope). ---
+    // The regression guard for `--dropout 0`. ops::dropout(p=0) returns x.clone() (a
+    // graph no-op) regardless of training mode, so the p=0 FORWARD graph is byte-for-
+    // byte the no-dropout forward → loss[0] must be BIT-IDENTICAL (the load-bearing
+    // claim, asserted == 0.0). At world=1 the NCCL all-reduce short-circuits, and one
+    // step has no optimizer-state compounding; the only residual non-determinism is
+    // the engine's atomicAdd backward-reduction ORDER (the documented fresh-train md5
+    // caveat — dropout-INDEPENDENT, present with or without the dropout op), which
+    // moves the post-step params by a single grad ULP. So params are checked against
+    // that tight reduction floor (< 1e-7), the same nature as the cross-rank KI-5
+    // tolerance used elsewhere in this file — not a dropout signal. GATE B (live) has
+    // a >1e-3 signal, ~4 orders of magnitude above this floor, so it carries the load.
     let d1 = [0u32];
+    let dcfg_1step = DdpConfig {
+        steps: 1,
+        eval_every: 0,
+        ..base_dcfg.clone()
+    };
     let cfg_nodrop = test_config(vocab); // cfg.dropout defaults to 0.0
     assert_eq!(cfg_nodrop.dropout, 0.0, "baseline cfg must have dropout 0");
     let mut cfg_p0 = test_config(vocab);
     cfg_p0.dropout = 0.0; // explicitly set p=0 — must not perturb anything
-    let (loss_nd1, params_nd1, _) = run_ddp(&d1, cfg_nodrop, &corpus, Some(&valid), &base_dcfg);
-    let (loss_p01, params_p01, _) = run_ddp(&d1, cfg_p0, &corpus, Some(&valid), &base_dcfg);
-    let max_loss_diff_1 = loss_nd1
-        .iter()
-        .zip(&loss_p01)
-        .map(|(a, b)| (a - b).abs())
-        .fold(0.0f32, f32::max);
+    let (loss_nd1, params_nd1, _) = run_ddp(&d1, cfg_nodrop, &corpus, None, &dcfg_1step);
+    let (loss_p01, params_p01, _) = run_ddp(&d1, cfg_p0, &corpus, None, &dcfg_1step);
+    let max_loss_diff_1 = (loss_nd1[0] - loss_p01[0]).abs();
     let max_param_diff_1 = params_nd1
         .iter()
         .zip(&params_p01)
         .flat_map(|(a, b)| a.iter().zip(b).map(|(x, y)| (x - y).abs()))
         .fold(0.0f32, f32::max);
     println!(
-        "T21 GATE A (world=1 p=0 bit-identical): max |loss diff| = {max_loss_diff_1:.3e}, \
-         max |param diff| = {max_param_diff_1:.3e}"
+        "T21 GATE A (world=1, 1 step, p=0 vs no-dropout): |loss diff| = {max_loss_diff_1:.3e} \
+         (bit-identical forward), max |param diff| = {max_param_diff_1:.3e} (atomicAdd floor)"
     );
     assert_eq!(
         max_loss_diff_1, 0.0,
-        "world=1 p=0 loss trace not bit-identical to no-dropout path"
+        "world=1 p=0 forward loss not bit-identical to no-dropout path"
     );
-    assert_eq!(
-        max_param_diff_1, 0.0,
-        "world=1 p=0 final params not bit-identical to no-dropout path"
+    assert!(
+        max_param_diff_1 < 1e-7,
+        "world=1 p=0 post-step params diverged from no-dropout beyond the atomicAdd \
+         reduction floor: {max_param_diff_1:.3e}"
     );
 
     // --- world=2 runs: real cross-rank NCCL all-reduce (the production path). ---
@@ -530,8 +542,7 @@ fn ddp_dropout_is_live_and_p0_bit_identical() {
     let mut cfg_p_w2 = test_config(vocab);
     cfg_p_w2.dropout = 0.2;
     let (loss_p0_2, _params_p0_2, _) = run_ddp(&d2, cfg_p0_w2, &corpus, Some(&valid), &base_dcfg);
-    let (loss_p_2, _params_p_2, train_flag_p) =
-        run_ddp(&d2, cfg_p_w2, &corpus, Some(&valid), &base_dcfg);
+    let (loss_p_2, _params_p_2, _) = run_ddp(&d2, cfg_p_w2, &corpus, Some(&valid), &base_dcfg);
 
     // GATE A2 — under DDP (world=2), p=0 matches a separate no-dropout baseline within
     // NCCL's run-to-run ULP noise (KI-5; the all-reduce is not bit-reproducible). This
@@ -571,17 +582,30 @@ fn ddp_dropout_is_live_and_p0_bit_identical() {
          (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_2.iter().all(|l| l.is_finite()),
         "p=0.2 DDP loss has non-finite values"
     );
+
+    // GATE C — train_rank actually sets TRAINING mode (direct, complementary proof of
+    // model.train() being wired). Use a dedicated short run with eval_every=0 so no
+    // eval fires: a model that finishes a training step in training mode proves
+    // train_rank called model.train(). (With eval enabled, eval_loss → model.eval()
+    // runs LAST on the final step and legitimately leaves the model in eval mode —
+    // same as the single-GPU loop — so is_training() after an eval-enabled run reflects
+    // the final eval, not the training-mode wiring. GATE B already proves dropout
+    // survives the mid-run eval flips via the per-step model.train() restore.) On the
+    // pre-T21 code is_training() stays false (model never left the default eval mode).
+    let dcfg_noeval = DdpConfig {
+        steps: 2,
+        eval_every: 0,
+        ..base_dcfg.clone()
+    };
+    let (_, _, train_flag) = run_ddp(&d1, cfg_p_w2, &corpus, None, &dcfg_noeval);
+    assert!(
+        train_flag,
+        "model not in training mode after a no-eval DDP run — model.train() not wired in train_rank"
+    );
+    println!("T21 GATE C (train_rank sets training mode): is_training() == true ✅");
 }