Gahow Wang
1b58bd8626
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>
554 lines
21 KiB
Rust
554 lines
21 KiB
Rust
//! 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 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>>) {
|
||
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(¶ms, dcfg.max_grad_norm, 1.0);
|
||
opt.step(lr, ¶ms);
|
||
for p in ¶ms {
|
||
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, 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"
|
||
);
|
||
}
|