xtrain/crates/xtrain-distributed/tests/ddp_correctness.rs

//! DDP acceptance (Phase T8). Gated to a GPU host; skips when fewer than 2 GPUs.
//!
//! 1. **Correctness**: K steps single-GPU (world=1, global batch B) vs 2-rank DDP
//!    (B/2 of the SAME data in the same order each) → loss trajectories match
//!    within tight fp tolerance (it's just gradient averaging), and the two
//!    ranks' parameters are identical after the run.
//! 2. **Throughput**: 1 / 2 / 4 GPU global tok/s on the SAME per-GPU workload →
//!    near-linear scaling. Prints the table (run with `--nocapture`).

#![cfg(not(no_cuda))]

use std::time::Instant;

use xtrain_cuda::device;
use xtrain_distributed::{DdpConfig, DdpContext, build_model, get_unique_id, launch, train_rank};
use xtrain_model::{Config, batched_ids_tensor};
use xtrain_optim::GpuAdamW;
use xtrain_tensor::Device;
use xtrain_train::clip::clip_grad_norm_gpu;
use xtrain_train::data::Corpus;
use xtrain_train::schedule::LrSchedule;

// A self-contained synthetic corpus so the test needs no tokenizer/data files.
fn synth_corpus(vocab: usize, n_tokens: usize) -> Corpus {
    let tokens: Vec<i32> = (0..n_tokens)
        .map(|i| (i * 7 + 3) as i32 % vocab as i32)
        .collect();
    Corpus {
        tokens,
        vocab_size: vocab,
    }
}

fn test_config(vocab: usize) -> Config {
    let mut cfg = Config::tiny();
    cfg.vocab = vocab;
    cfg.n_layers = 2;
    cfg
}

/// Run `cfg`/`dcfg` as a DDP job over `devices` (the same launcher path as
/// production — `DdpContext::init` + `train_rank` per rank) 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.
///
/// `world == 1` is the deterministic path: `all_reduce_average_grads` short-circuits
/// (no NCCL collective), so the run is bit-reproducible — used for the bit-identity
/// gate. `world >= 2` exercises the real cross-rank NCCL all-reduce, which is not
/// bit-reproducible run-to-run on this PCIe box (KI-5), so those gates use the same
/// ULP/relative tolerances as the rest of this file.
fn run_ddp(
    devices: &[u32],
    cfg: Config,
    corpus: &Corpus,
    valid: Option<&Corpus>,
    dcfg: &DdpConfig,
) -> (Vec<f32>, Vec<Vec<f32>>, bool) {
    let world = devices.len();
    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>>) {
    device::set_device(0).unwrap();
    let device = Device::Cuda(0);
    let model = build_model(cfg, device);
    let params = model.params();
    let mut opt = GpuAdamW::new(dcfg.weight_decay);
    let mut rng = dcfg.seed;
    let mut losses = Vec::new();

    for step in 0..dcfg.steps {
        let lr = dcfg.schedule.lr(step);
        // Sample the whole global batch and run it as ONE batched forward/backward
        // (matches the T10 DDP path: backward yields the global-batch mean grad).
        let mut inputs = Vec::with_capacity(dcfg.batch_size);
        let mut targets_v = Vec::with_capacity(dcfg.batch_size);
        for _ in 0..dcfg.batch_size {
            let (input, target) = corpus.sample(dcfg.seq_len, &mut rng);
            inputs.push(input);
            targets_v.push(target);
        }
        let ids = batched_ids_tensor(&inputs, device);
        let targets = batched_ids_tensor(&targets_v, device);
        let loss = model.loss_batched(&ids, &targets, dcfg.batch_size);
        losses.push(loss.value().to_device(Device::Cpu).as_slice::<f32>()[0]);
        loss.backward();
        clip_grad_norm_gpu(&params, dcfg.max_grad_norm, 1.0);
        opt.step(lr, &params);
        for p in &params {
            p.zero_grad();
        }
    }

    let host = params
        .iter()
        .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
        .collect();
    (losses, host)
}

#[test]
fn ddp_matches_single_gpu_and_params_consistent() {
    let world = 2usize;
    if device::device_count().unwrap_or(0) < world as i32 {
        eprintln!("skip: need >= {world} GPUs");
        return;
    }

    let vocab = 64usize;
    let cfg = test_config(vocab);
    let corpus = synth_corpus(vocab, 4096);
    let steps = 20usize;
    let 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 in the test
        seed: 7,
        eval_every: 0,
        eval_batches: 0,
        ckpt_path: None,
    };

    // Single-GPU baseline (world=1) over the global batch.
    let (single_losses, single_params) = run_single_gpu(cfg, &corpus, &dcfg);

    // 2-rank DDP over the SAME corpus/config; returns per-rank (losses, params).
    let devices = [0u32, 1u32];
    let id = get_unique_id();
    let results: Vec<(Vec<f32>, Vec<Vec<f32>>)> = 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);
                    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<_>>();
                    (res.losses, host)
                })
            })
            .collect();
        handles.into_iter().map(|h| h.join().unwrap()).collect()
    });

    let (ddp_losses, ddp_p0) = &results[0];
    let (_, ddp_p1) = &results[1];

    // (a) DDP loss trajectory matches single-GPU within tight tolerance.
    let mut max_rel = 0.0f32;
    for (s, d) in single_losses.iter().zip(ddp_losses) {
        let rel = (s - d).abs() / s.abs().max(1e-6);
        max_rel = max_rel.max(rel);
    }
    println!(
        "DDP vs single-GPU loss: single[last]={:.6} ddp[last]={:.6} max_rel={max_rel:.2e}",
        single_losses.last().unwrap(),
        ddp_losses.last().unwrap()
    );
    assert!(
        max_rel < 1e-3,
        "DDP loss trajectory diverged from single-GPU: max_rel {max_rel:.3e}"
    );

    // (b) Cross-rank parameter identity (same init + same averaged grad + same
    //     optimizer state ⇒ identical params).
    let mut max_pdiff = 0.0f32;
    for (a, b) in ddp_p0.iter().zip(ddp_p1) {
        for (x, y) in a.iter().zip(b) {
            max_pdiff = max_pdiff.max((x - y).abs());
        }
    }
    println!("cross-rank max |param diff| = {max_pdiff:.3e}");
    // On this PCIe-only box, NCCL's all-reduce is not bit-reproducible run-to-run
    // across ranks (algorithm/chunk choice is unstable), so cross-rank params can
    // differ by a few ULP (observed ≤1.2e-7) even with identical init + averaged
    // grads. The load-bearing gate is the loss-trajectory match (a, ~5.7e-7); a
    // tight tolerance here, not bit-identity, is the honest invariant (KI-5).
    assert!(
        max_pdiff < 1e-6,
        "ranks' params drifted apart: {max_pdiff:.3e}"
    );

    // (c) DDP final params match single-GPU final params within fp tolerance.
    // Looser than (a)/(b): DDP and single-GPU differ only in the gradient SUMMATION
    // ORDER (single-GPU sums B sequences in tape order; DDP sums per-rank shards
    // then NCCL-sums across ranks). fp addition isn't associative, so that tiny
    // per-step rounding compounds over the AdamW steps — a few e-3 relative on
    // individual params is expected and benign. The loss-trajectory match (a, ~1e-7)
    // and tight cross-rank agreement (b, <1e-6) are the load-bearing checks.
    let mut max_sdiff = 0.0f32;
    for (a, b) in ddp_p0.iter().zip(&single_params) {
        for (x, y) in a.iter().zip(b) {
            max_sdiff = max_sdiff.max((x - y).abs() / y.abs().max(1e-6));
        }
    }
    println!("DDP vs single-GPU max rel |param diff| = {max_sdiff:.3e}");
    assert!(max_sdiff < 1e-2, "DDP params diverged from single-GPU");
}

#[test]
fn ddp_with_accum_matches_single_gpu_big_batch() {
    // T16: DDP + gradient accumulation must match a single-GPU big-batch baseline
    // of the SAME effective batch. world=2, accum=2, per-rank micro-batch 2 →
    // effective global batch = world·accum·b_local = 2·2·2 = 8. Compared against a
    // single-GPU run with batch 8, accum 1 (the big-batch baseline). The all-reduce
    // fires only at the accumulation boundary (once per optimizer step, not per
    // micro-step) — enforced by the train_rank implementation; the load-bearing
    // gate here is that loss + final params still match the big-batch baseline.
    let world = 2usize;
    if device::device_count().unwrap_or(0) < world as i32 {
        eprintln!("skip: need >= {world} GPUs");
        return;
    }

    let vocab = 64usize;
    let cfg = test_config(vocab);
    let corpus = synth_corpus(vocab, 4096);
    let steps = 20usize;
    let effective_batch = 8usize; // world(2) · accum(2) · b_local(2)
    let sched = LrSchedule {
        max_lr: 3e-3,
        min_lr: 3e-4,
        warmup: 3,
        total: steps,
    };

    // Single-GPU big-batch baseline: world=1, accum=1, batch = effective_batch.
    let baseline_cfg = DdpConfig {
        seq_len: 32,
        batch_size: effective_batch,
        accum_steps: 1,
        steps,
        schedule: sched,
        weight_decay: 0.1,
        max_grad_norm: 1.0,
        log_every: 1_000_000,
        seed: 7,
        eval_every: 0,
        eval_batches: 0,
        ckpt_path: None,
    };
    let (single_losses, single_params) = run_single_gpu(cfg, &corpus, &baseline_cfg);

    // DDP + accumulation: world=2, accum=2 → per-rank micro-batch = batch/world = 2.
    let ddp_cfg = DdpConfig {
        batch_size: effective_batch / 2, // per-step global batch; ×accum = effective
        accum_steps: 2,
        ..baseline_cfg
    };
    let devices = [0u32, 1u32];
    let id = get_unique_id();
    let results: Vec<(Vec<f32>, Vec<Vec<f32>>)> = std::thread::scope(|s| {
        let handles: Vec<_> = devices
            .iter()
            .enumerate()
            .map(|(rank, &dev)| {
                let ddp_cfg = ddp_cfg.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);
                    let res = train_rank(&ctx, &model, device, corpus, None, &ddp_cfg);
                    let host = model
                        .params()
                        .iter()
                        .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
                        .collect::<Vec<_>>();
                    (res.losses, host)
                })
            })
            .collect();
        handles.into_iter().map(|h| h.join().unwrap()).collect()
    });

    let (ddp_losses, ddp_p0) = &results[0];
    let (_, ddp_p1) = &results[1];

    // (a) Loss trajectory matches the single-GPU big-batch baseline.
    let mut max_rel = 0.0f32;
    for (s, d) in single_losses.iter().zip(ddp_losses) {
        max_rel = max_rel.max((s - d).abs() / s.abs().max(1e-6));
    }
    println!(
        "DDP+accum(w2·a2·b2) vs single-GPU big-batch(8): single[last]={:.6} ddp[last]={:.6} max_rel={max_rel:.2e}",
        single_losses.last().unwrap(),
        ddp_losses.last().unwrap()
    );
    assert!(
        max_rel < 1e-3,
        "DDP+accum loss diverged from big-batch baseline: {max_rel:.3e}"
    );

    // (b) Cross-rank parameter agreement (same KI-5 ULP tolerance as the base test).
    let mut max_pdiff = 0.0f32;
    for (a, b) in ddp_p0.iter().zip(ddp_p1) {
        for (x, y) in a.iter().zip(b) {
            max_pdiff = max_pdiff.max((x - y).abs());
        }
    }
    println!("DDP+accum cross-rank max |param diff| = {max_pdiff:.3e}");
    assert!(
        max_pdiff < 1e-6,
        "ranks' params drifted apart: {max_pdiff:.3e}"
    );

    // (c) Final params match single-GPU big-batch within fp tolerance.
    let mut max_sdiff = 0.0f32;
    for (a, b) in ddp_p0.iter().zip(&single_params) {
        for (x, y) in a.iter().zip(b) {
            max_sdiff = max_sdiff.max((x - y).abs() / y.abs().max(1e-6));
        }
    }
    println!("DDP+accum vs single-GPU big-batch max rel |param diff| = {max_sdiff:.3e}");
    assert!(
        max_sdiff < 1e-2,
        "DDP+accum params diverged from big-batch baseline"
    );
}

#[test]
fn ddp_throughput_scaling() {
    let max_gpus = device::device_count().unwrap_or(0) as usize;
    if max_gpus < 1 {
        eprintln!("skip: no GPU");
        return;
    }
    // Same PER-GPU workload at each world size (batch scales with world), so the
    // per-rank cost is fixed and global tok/s should scale ~linearly. Use enough
    // steps that the one-time NCCL init + model-build overhead (which is larger at
    // world=4 and absent at world=1) amortizes — otherwise the wall-clock ratio
    // understates steady-state scaling.
    let per_gpu_batch = 8usize;
    let vocab = 256usize;
    let cfg = test_config(vocab);
    let corpus = synth_corpus(vocab, 8192);
    let steps = 150usize;
    let seq_len = 64usize;

    let worlds: Vec<usize> = [1, 2, 4, 8]
        .into_iter()
        .filter(|&w| w <= max_gpus)
        .collect();
    println!("\n=== DDP throughput scaling (per-GPU batch {per_gpu_batch}, seq {seq_len}) ===");
    println!(
        "{:>6} | {:>14} | {:>8}",
        "GPUs", "tok/s (global)", "speedup"
    );

    let mut base = 0.0f64;
    for &world in &worlds {
        let devices: Vec<u32> = (0..world as u32).collect();
        let dcfg = DdpConfig {
            seq_len,
            batch_size: per_gpu_batch * world,
            accum_steps: 1,
            steps,
            schedule: LrSchedule {
                max_lr: 1e-3,
                min_lr: 1e-3,
                warmup: 1,
                total: steps,
            },
            weight_decay: 0.0,
            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, None, &dcfg, move |device| {
            build_model(cfg, device)
        });
        let secs = t.elapsed().as_secs_f64();
        let tps = total_tokens / secs;
        if world == 1 {
            base = tps;
        }
        println!(
            "{:>6} | {:>14.0} | {:>7.2}x",
            world,
            tps,
            tps / base.max(1e-9)
        );
    }
}

/// T21 regression: prove dropout is actually LIVE under DDP (with `p>0`), 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 (`DdpContext::init` + `train_rank`).
///
/// On the pre-T21 code, both load-bearing gates FAIL: GATE B (p>0 trace would be
/// 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.
#[test]
fn ddp_dropout_is_live_and_p0_bit_identical() {
    if device::device_count().unwrap_or(0) < 2 {
        eprintln!("skip: need >= 2 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,
    };

    // --- 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).
    let d1 = [0u32];
    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 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}"
    );
    assert_eq!(
        max_loss_diff_1, 0.0,
        "world=1 p=0 loss trace 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"
    );

    // --- world=2 runs: real cross-rank NCCL all-reduce (the production path). ---
    let d2 = [0u32, 1u32];
    let mut cfg_p0_w2 = test_config(vocab);
    cfg_p0_w2.dropout = 0.0;
    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);

    // 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
    // confirms enabling dropout=0 doesn't perturb the DDP path beyond that noise floor.
    let (loss_nd_2, _, _) = run_ddp(&d2, test_config(vocab), &corpus, Some(&valid), &base_dcfg);
    let max_loss_diff_2 = loss_nd_2
        .iter()
        .zip(&loss_p0_2)
        .map(|(a, b)| (a - b).abs())
        .fold(0.0f32, f32::max);
    println!("T21 GATE A2 (world=2 p=0 vs no-dropout, KI-5 noise): max |loss diff| = {max_loss_diff_2:.3e}");
    assert!(
        max_loss_diff_2 < 1e-6,
        "world=2 p=0 diverged from no-dropout beyond NCCL noise: {max_loss_diff_2:.3e}"
    );

    // 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 (diff at the ~1e-7 NCCL noise
    // floor). A difference orders of magnitude above that proves dropout masks are
    // actually applied during the training forward — and that they survive the mid-run
    // eval flips (model.train() is re-asserted each step). Inverted scaling + masking
    // perturbs every step, so the gap is large (>1e-3 ≫ KI-5 noise ~2.4e-7).
    let max_live_diff = loss_p0_2
        .iter()
        .zip(&loss_p_2)
        .map(|(a, b)| (a - b).abs())
        .fold(0.0f32, f32::max);
    println!(
        "T21 GATE B (dropout live, world=2): p0[last]={:.6} p0.2[last]={:.6} max |loss diff| = {max_live_diff:.3e}",
        loss_p0_2.last().unwrap(),
        loss_p_2.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_2.iter().all(|l| l.is_finite()),
        "p=0.2 DDP loss has non-finite values"
    );
}