distributed: T21 — wire dropout into the DDP path (--dropout + model.train())

V9-PILOT caught a launcher-level integration gap: T18 wired dropout into
the single-GPU bin/train, but the DDP path never did. train_ddp had no
--dropout flag and never set cfg.dropout, and ddp.rs::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.

Fix, mirroring the single-GPU train/eval discipline:
- train_ddp.rs: add a --dropout <p> flag (default 0 = off, matching the
  prior behavior) and set cfg.dropout from it; log it when on.
- ddp.rs::train_rank: call model.train() at the start of each step (before
  the micro-batch loop). eval_loss() flips the model to eval mode and does
  not restore it, so re-asserting train() each step keeps dropout live
  across eval boundaries.

--dropout 0 (default) is bit-identical to the prior DDP path: cfg.dropout
stays 0 and ops::dropout(p=0) is a clone no-op regardless of training mode.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-18 21:08:17 +08:00
parent db70abe450
commit 81f3cf59e5
2 changed files with 20 additions and 1 deletions

View File

@@ -124,6 +124,14 @@ pub fn train_rank(
// all-reduce fires ONLY after the last micro-step (intermediate micro-steps
// are local-only, no NCCL).
let mut local_sum = 0.0f32; // Σ over micro of (local_mean · b_local)
// Training mode → dropout active (T18; no-op when cfg.dropout == 0). Set
// each step so it is restored after a periodic eval flips the model to eval
// mode (eval_loss calls model.eval() and does not restore). Mirrors the
// single-GPU loop's train/eval discipline — without this, DDP forwards run
// in the default eval (identity) mode and --dropout is silently ignored
// (the T21 launcher-wiring gap the V9-PILOT caught). Each micro-step's
// forward bumps the per-step seed → fresh masks.
model.train();
for _ in 0..accum {
let mut inputs = Vec::with_capacity(batch_local);
let mut targets_v = Vec::with_capacity(batch_local);