dropout: wire into model (residual sites) + train/eval switch + flag (T18)

Config.dropout (default 0). TinyTransformer gets a Cell<bool> training switch
(train()/eval()/with_training, default eval = safe) + a Cell<u64> step_seed bumped
once per training forward. forward_batched derives a per-layer block_seed (pure fn
of step_seed×layer) and block_forward derives two per-site seeds, inserting
ops::dropout at the attn and ffn sub-block outputs (before each residual). The
seed is a pure function of (step_seed, layer, site) so the checkpoint (T13)
recompute re-derives the same masks → grads stay exact. p=0 or eval → no dropout
node → graph bit-identical to pre-T18.

train_loop: model.train() per step (restored after eval flips to eval); eval_loss
runs model.eval(). bin/train: --dropout flag → cfg.dropout. Export/sampling run in
eval (default), so exported weights are dropout-free (xserv closed loop unaffected).

Model-level tests (dropout.rs): p=0 bit-identical to no-dropout (logits/loss/grads);
eval(p>0) == p=0 identity; train differs from eval + finite; recompute-with-dropout
grads match non-recompute (fp32 + bf16).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-model/src/config.rs

@@ -20,6 +20,11 @@ pub struct Config {
     pub eps: f32,
     /// RoPE base frequency (theta).
     pub rope_theta: f32,
+    /// Dropout probability `p` (Phase T18). Applied at the attention/MLP sub-block
+    /// outputs (before each residual add) at TRAINING time, with inverted scaling
+    /// `1/(1-p)`; disabled (identity) at eval. Default `0.0` = no dropout, and the
+    /// forward graph is then bit-identical to the pre-T18 path.
+    pub dropout: f32,
 }
 
 impl Config {
@@ -36,6 +41,7 @@ impl Config {
             ffn_hidden: 64,
             eps: 1e-5,
             rope_theta: 10000.0,
+            dropout: 0.0,
         }
     }
 
@@ -60,6 +66,7 @@ impl Config {
             ffn_hidden,
             eps: 1e-5,
             rope_theta: 10000.0,
+            dropout: 0.0,
         }
     }
 

crates/xtrain-model/src/model.rs

@@ -2,6 +2,8 @@
 
 #![cfg(not(no_cuda))]
 
+use std::cell::Cell;
+
 use crate::config::Config;
 use xtrain_autodiff::ops;
 use xtrain_autodiff::tape::Var;
@@ -47,6 +49,19 @@ pub struct TinyTransformer {
     /// existing numerics are bit-identical; recompute is mathematically exact, so
     /// grads match the non-checkpointed path within fp tolerance.
     recompute: bool,
+    /// Training mode for dropout (Phase T18). `true` → the attn/MLP sub-block
+    /// outputs pass through `ops::dropout` (with `cfg.dropout` and a per-step,
+    /// per-site seed); `false` (default) → dropout is identity (eval/sampling/
+    /// export). `Cell` so `train()`/`eval()` flip it through `&self` (the forward
+    /// takes `&self`). When `cfg.dropout == 0` this flag is irrelevant — the graph
+    /// is bit-identical to the no-dropout path either way.
+    training: Cell<bool>,
+    /// Per-step dropout RNG seed (Phase T18). Bumped once at the start of each
+    /// TRAINING forward so every step draws fresh masks; combined with the layer
+    /// index + a per-site constant to give each dropout site its own seed. The RNG
+    /// is counter-based, so re-running a checkpointed block's forward in backward
+    /// (T13) reproduces the same seed → the same mask (recompute stays exact).
+    step_seed: Cell<u64>,
 }
 
 impl TinyTransformer {
@@ -90,6 +105,8 @@ impl TinyTransformer {
             lm_head,
             compute_dtype: DType::F32,
             recompute: false,
+            training: Cell::new(false),
+            step_seed: Cell::new(0),
         }
     }
 
@@ -127,6 +144,30 @@ impl TinyTransformer {
         self.recompute
     }
 
+    /// Switch to training mode (Phase T18): dropout (if `cfg.dropout > 0`) is
+    /// active in subsequent forwards. The training loop calls this before stepping.
+    pub fn train(&self) {
+        self.training.set(true);
+    }
+
+    /// Switch to eval mode (Phase T18): dropout is identity. Held-out eval,
+    /// autoregressive sampling, and weight export all run in this mode (default).
+    pub fn eval(&self) {
+        self.training.set(false);
+    }
+
+    pub fn is_training(&self) -> bool {
+        self.training.get()
+    }
+
+    /// Builder-style train/eval toggle (Phase T18) — handy for tests that want a
+    /// model fixed in one mode. Equivalent to [`train`](Self::train) /
+    /// [`eval`](Self::eval) but chains off `new(..)`.
+    pub fn with_training(self, training: bool) -> Self {
+        self.training.set(training);
+        self
+    }
+
     /// All learnable parameters, in a stable order. The optimizer (a hand-written
     /// GD step in T5, AdamW in T6) iterates this; each holds its `.grad()` after
     /// `backward()`.
@@ -176,13 +217,34 @@ impl TinyTransformer {
         );
         let seq = total / batch;
 
+        // Dropout (T18) is active only in training mode with p>0; otherwise it is
+        // identity (`ops::dropout` no-ops at p==0). Bump the per-step seed ONCE per
+        // training forward so each step draws fresh masks (counter-based RNG, so a
+        // checkpointed block's recompute reproduces the same seed → same mask).
+        let dropout_p = if self.training.get() {
+            self.cfg.dropout
+        } else {
+            0.0
+        };
+        if dropout_p > 0.0 {
+            self.step_seed.set(self.step_seed.get().wrapping_add(1));
+        }
+        let base_seed = self.step_seed.get();
+
         // Embedding gathers from the fp32 master table; in bf16 mode cast the
         // activation stream to bf16 here (norms are cast to bf16 gammas too).
         let mut h = ops::embedding(&self.embed, ids); // [batch*seq, dim], fp32
         if self.compute_dtype == DType::BF16 {
             h = ops::cast(&h, DType::BF16);
         }
-        for b in &self.blocks {
+        for (li, b) in self.blocks.iter().enumerate() {
+            // Per-layer dropout seed: a deterministic function of (base_seed,
+            // layer index) — NOT a mutable counter — so the checkpoint recompute
+            // (which re-derives it from the captured base_seed/li) gets the same
+            // masks. The block derives its two per-site seeds from this.
+            let block_seed = base_seed
+                .wrapping_mul(0x100000001B3)
+                .wrapping_add(li as u64);
             h = if self.recompute {
                 // Activation recomputation (T13): run the whole block forward inside
                 // `checkpoint` so its internal activations aren't kept on the tape;
@@ -190,7 +252,9 @@ impl TinyTransformer {
                 // segment fn captures only `Copy` config (no borrow of `self`) and
                 // receives the block's params via the slice, in `block_params` order.
                 let (cfg, cdt) = (self.cfg, self.compute_dtype);
-                let seg = move |x: &Var, p: &[Var]| block_forward(cfg, cdt, batch, seq, x, p);
+                let seg = move |x: &Var, p: &[Var]| {
+                    block_forward(cfg, cdt, batch, seq, dropout_p, block_seed, x, p)
+                };
                 xtrain_autodiff::checkpoint::checkpoint(seg, &h, &b.block_params())
             } else {
                 block_forward(
@@ -198,6 +262,8 @@ impl TinyTransformer {
                     self.compute_dtype,
                     batch,
                     seq,
+                    dropout_p,
+                    block_seed,
                     &h,
                     &b.block_params(),
                 )
@@ -275,25 +341,46 @@ fn norm_gamma(cdt: DType, gamma: &Var) -> Var {
 }
 
 /// One transformer block's forward: pre-norm + multi-head causal attention +
-/// residual, then pre-norm + SwiGLU MLP + residual. Pure in `(cfg, cdt, batch,
-/// seq, input, params)` (no `&self`) so it can be the segment fn of
-/// [`xtrain_autodiff::checkpoint`] for activation recomputation (T13). `params` is
-/// the block's leaves in [`Block::block_params`] order.
-fn block_forward(cfg: Config, cdt: DType, batch: usize, seq: usize, h: &Var, p: &[Var]) -> Var {
+/// (T18) dropout + residual, then pre-norm + SwiGLU MLP + dropout + residual.
+/// Pure in `(cfg, cdt, batch, seq, dropout_p, block_seed, input, params)` (no
+/// `&self`, all `Copy`) so it can be the segment fn of
+/// [`xtrain_autodiff::checkpoint`] for activation recomputation (T13) — the
+/// recompute re-derives the same per-site seeds, so the dropout masks are
+/// reproduced bit-for-bit. `dropout_p == 0` makes `ops::dropout` a no-op (the
+/// graph is then identical to the pre-T18 path). `params` is the block's leaves in
+/// [`Block::block_params`] order.
+#[allow(clippy::too_many_arguments)]
+fn block_forward(
+    cfg: Config,
+    cdt: DType,
+    batch: usize,
+    seq: usize,
+    dropout_p: f32,
+    block_seed: u64,
+    h: &Var,
+    p: &[Var],
+) -> Var {
     let (attn_norm, wq, wk, wv) = (&p[0], &p[1], &p[2], &p[3]);
     let (q_norm, k_norm, wo) = (&p[4], &p[5], &p[6]);
     let (ffn_norm, w_gate, w_up, w_down) = (&p[7], &p[8], &p[9], &p[10]);
 
-    // --- Attention sub-block (pre-norm + residual) ---
+    // Per-site dropout seeds (XOR a site constant into the block seed) so the two
+    // residual-path dropouts draw independent masks within the same step/layer.
+    let attn_seed = block_seed ^ 0x0A7700;
+    let ffn_seed = block_seed ^ 0x0FF700;
+
+    // --- Attention sub-block (pre-norm + dropout + residual) ---
     let normed = ops::rms_norm(h, &norm_gamma(cdt, attn_norm), cfg.eps);
     let attn = attention(
         cfg, cdt, batch, seq, &normed, wq, wk, wv, q_norm, k_norm, wo,
     );
+    let attn = ops::dropout(&attn, dropout_p, attn_seed);
     let h = ops::add(h, &attn);
 
-    // --- MLP sub-block (pre-norm + residual) ---
+    // --- MLP sub-block (pre-norm + dropout + residual) ---
     let normed = ops::rms_norm(&h, &norm_gamma(cdt, ffn_norm), cfg.eps);
     let mlp = swiglu_mlp(cdt, &normed, w_gate, w_up, w_down);
+    let mlp = ops::dropout(&mlp, dropout_p, ffn_seed);
     ops::add(&h, &mlp)
 }