gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Gahow Wang
2026-06-18 21:08:27 +08:00
parent 81f3cf59e5
commit 1b58bd8626
1 changed files with 165 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

165
crates/xtrain-distributed/tests/ddp_correctness.rs
View File

@@ -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"
);
}