docs: Phase T13 — activation recompute

Design doc for per-block gradient checkpointing (KI-3): the no-tape forward +
recompute-on-backward design, the `checkpoint` primitive, per-block wrapping,
the exactness/correctness argument (same kernels + inputs → identical grads),
composition with bf16+DDP+batched, and the verification plan (on-vs-off grad
gate + memory/throughput before→after, dim1024-fits). Bench table left as TBD
to fill after the dash5 run.

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

crates/xtrain-autodiff/src/checkpoint.rs

@@ -0,0 +1,103 @@
+//! Activation recomputation / gradient checkpointing (Phase T13, KI-3).
+//!
+//! A higher-order autograd primitive — the analogue of `torch.utils.checkpoint`.
+//! It runs a *segment* of the model (a transformer block, here) WITHOUT recording
+//! the segment's internal ops on the surrounding tape, so the segment's
+//! intermediate activations are freed right after the forward instead of being
+//! kept alive until backward. When the segment's output-grad arrives in backward,
+//! the segment forward is **re-run** from the saved input (into a throwaway local
+//! tape), the recomputed output is seeded with the upstream grad, and the gradient
+//! is backpropagated through the local tape to recover the input-grad and the
+//! parameter-grads — which are then pushed to the real tape's parents. The local
+//! tape is dropped at the end of the closure, freeing the recomputed activations.
+//!
+//! ## Why it is exact (the hard gate)
+//! The recompute runs the *same* `segment_fn` from the *same* input value and the
+//! *same* parameter values (parameters are leaves that persist across forward and
+//! backward; only their grad slot changes). The forward kernels are deterministic,
+//! so the recomputed output equals the original output bit-for-bit, and the local
+//! backward is the ordinary analytic backward of that segment. Therefore the input-
+//! and parameter-grads are identical to those a non-checkpointed forward would
+//! produce — checkpointing trades compute (one extra forward per segment) for
+//! memory, never correctness.
+//!
+//! ## Composition
+//! - **bf16 (T12):** `segment_fn` is the unchanged block forward, so the recompute
+//!   runs the same bf16 path; the `cast` op's grad upcast still bridges bf16→fp32.
+//! - **DDP (T8):** each rank checkpoints its own forward/backward independently;
+//!   the param-grads recovered here feed the same per-rank `.grad()` slots that the
+//!   all-reduce averages — no change to the distributed path.
+//! - **batched (T10):** the segment input/output carry the `[batch*seq, …]` batch
+//!   dim transparently; `checkpoint` is shape-agnostic.
+
+#![cfg(not(no_cuda))]
+
+use crate::tape::Var;
+use std::rc::Rc;
+
+/// Run `segment_fn(input, params)` with activation recomputation.
+///
+/// `segment_fn(x, p)` must build the segment's forward graph from a single input
+/// `x` and the parameter slice `p`, returning the single segment output. It is
+/// called once now (forward, result detached from the outer tape) and once per
+/// backward (recompute). It MUST be deterministic and depend only on `x` and `p`
+/// (this is what makes the recompute exact).
+///
+/// `params` are the segment's learnable leaves; their grads are accumulated into
+/// the SAME leaves the optimizer reads (so DDP / AdamW are unchanged).
+///
+/// Returns the segment output as a `Var` on the outer tape whose backward triggers
+/// the recompute. Equivalent — grad-for-grad — to calling `segment_fn(input,
+/// params)` directly, but without keeping the segment's internal activations alive.
+pub fn checkpoint<F>(segment_fn: F, input: &Var, params: &[Var]) -> Var
+where
+    F: Fn(&Var, &[Var]) -> Var + 'static,
+{
+    let segment_fn = Rc::new(segment_fn);
+
+    // --- Forward (no taping of internals) ---
+    // Detach the input and params into fresh leaves so `segment_fn` builds a LOCAL
+    // tape disconnected from the outer graph. We only keep the output's value; the
+    // local `Var`s (and thus the segment's intermediate activations) are dropped
+    // when this scope ends.
+    let out_value = {
+        let x_det = Var::leaf(input.value());
+        let params_det: Vec<Var> = params.iter().map(|p| Var::leaf(p.value())).collect();
+        let out_local = segment_fn(&x_det, &params_det);
+        out_local.value()
+    };
+
+    // Parents on the OUTER tape: the segment input, then the params (so their grads
+    // land in the leaves the optimizer reads).
+    let mut parents = Vec::with_capacity(1 + params.len());
+    parents.push(input.clone());
+    parents.extend(params.iter().cloned());
+
+    let segment_fn = segment_fn.clone();
+    Var::from_op(
+        out_value,
+        parents,
+        Box::new(move |dout, parents| {
+            // --- Backward (recompute) ---
+            // Rebuild fresh leaves from the CURRENT input/param values (params are
+            // unchanged since forward; input is the saved segment input), re-run the
+            // forward to rebuild the local tape, seed the recomputed output with the
+            // upstream grad, and backprop through the local tape.
+            let x_det = Var::leaf(parents[0].value());
+            let params_det: Vec<Var> = parents[1..].iter().map(|p| Var::leaf(p.value())).collect();
+            let out_local = segment_fn(&x_det, &params_det);
+            out_local.backward_seeded(dout.clone());
+
+            // Push the recovered grads to the real parents (engine SUMs on fan-out).
+            if let Some(dx) = x_det.grad() {
+                Var::push_grad(&parents[0], dx);
+            }
+            for (det, parent) in params_det.iter().zip(&parents[1..]) {
+                if let Some(dp) = det.grad() {
+                    Var::push_grad(parent, dp);
+                }
+            }
+            // `out_local` / the local tape drop here → recomputed activations freed.
+        }),
+    )
+}

crates/xtrain-autodiff/src/lib.rs

@@ -18,6 +18,8 @@ pub use finite_diff::{GradCheckConfig, GradCheckResult, ParamFn, grad_check};
 // kernels via xtrain-tensor, so they are gated behind `not(no_cuda)` (the
 // per-crate convention); the grad_check harness above stays host-only.
 #[cfg(not(no_cuda))]
+pub mod checkpoint;
+#[cfg(not(no_cuda))]
 pub mod ops;
 #[cfg(not(no_cuda))]
 pub mod tape;

crates/xtrain-autodiff/src/tape.rs

@@ -108,14 +108,24 @@ impl Var {
             "backward() expects a scalar loss; got shape {:?}",
             self.value().shape()
         );
+        self.backward_seeded(ones_like(&self.value()));
+    }
 
+    /// Reverse-mode backward from this node seeded with an explicit upstream grad
+    /// `seed` (same shape as this node's value), instead of the scalar `dL/dL = 1`.
+    ///
+    /// This is the entry point for **activation recomputation** (Phase T13): a
+    /// checkpointed segment re-runs its forward into a fresh local tape, then
+    /// backprops the upstream output-grad through it via this method (the segment
+    /// output is generally NOT a scalar). For a scalar root, [`backward`] is the
+    /// thin wrapper that seeds ones.
+    pub fn backward_seeded(&self, seed: Tensor) {
         // 1. Topological order (post-order DFS), parents before children.
         let mut topo: Vec<Var> = Vec::new();
         let mut visited: Vec<*const RefCell<VarNode>> = Vec::new();
         build_topo(self, &mut topo, &mut visited);
 
-        // 2. Seed the loss gradient with ones.
-        let seed = ones_like(&self.value());
+        // 2. Seed this node's gradient with the supplied upstream grad.
         self.accumulate(seed);
 
         // 3. Walk in reverse: each node hands its grad to its parents' closures.

crates/xtrain-distributed/src/bin/train_ddp.rs

@@ -85,6 +85,10 @@ fn main() {
     // bf16 mixed precision (Phase T12): fp32 master weights, bf16 linears +
     // activations. Opt-in; default fp32 reproduces v0–v4 numerics.
     let bf16 = args.iter().any(|a| a == "--bf16");
+    // Activation recomputation (Phase T13): per-block gradient checkpointing — each
+    // rank checkpoints its own forward/backward; exact grads, lower peak activation
+    // memory (lets dim1024 batch32 fit). Opt-in; default off.
+    let recompute = args.iter().any(|a| a == "--recompute");
     let ckpt: Option<PathBuf> = args
         .iter()
         .position(|a| a == "--ckpt")
@@ -167,18 +171,23 @@ fn main() {
     if bf16 {
         println!("bf16 mixed precision: ON (fp32 master weights)");
     }
+    if recompute {
+        println!("activation recompute: ON (per-block gradient checkpointing)");
+    }
     let results = launch(
         &devices,
         &train_corpus,
         valid.as_ref(),
         &dcfg,
         move |device| {
-            let m = build_model(cfg, device);
+            let mut m = build_model(cfg, device);
             if bf16 {
-                m.with_compute_dtype(xtrain_tensor::DType::BF16)
-            } else {
-                m
+                m = m.with_compute_dtype(xtrain_tensor::DType::BF16);
             }
+            if recompute {
+                m = m.with_recompute(true);
+            }
+            m
         },
     );
     let r0 = &results[0];

crates/xtrain-model/src/model.rs

@@ -37,6 +37,16 @@ pub struct TinyTransformer {
     /// `docs/11-bf16-mixed-precision.md`). The cast op's backward upcasts the bf16
     /// weight grad back to fp32, so AdamW/clip/DDP stay fp32 and unchanged.
     compute_dtype: DType,
+    /// Activation recomputation / gradient checkpointing (Phase T13, KI-3). When
+    /// `true`, each transformer block's forward runs through
+    /// [`xtrain_autodiff::checkpoint`]: the block's internal activations are NOT
+    /// kept on the tape during forward (only the block input is), and the block
+    /// forward is re-run during backward to recover them. Trades ~one extra forward
+    /// per block for a large drop in peak activation memory → lets dim1024 batch32
+    /// fit. Default `false` = the unchanged path (every activation stored), so
+    /// existing numerics are bit-identical; recompute is mathematically exact, so
+    /// grads match the non-checkpointed path within fp tolerance.
+    recompute: bool,
 }
 
 impl TinyTransformer {
@@ -79,6 +89,7 @@ impl TinyTransformer {
             final_norm,
             lm_head,
             compute_dtype: DType::F32,
+            recompute: false,
         }
     }
 
@@ -103,16 +114,17 @@ impl TinyTransformer {
         self.compute_dtype
     }
 
-    /// Project `x` (activation, in the compute dtype) by weight `w` (an fp32
-    /// master leaf). In bf16 mode the weight is cast to bf16 via the autograd
-    /// `cast` op (whose backward upcasts the grad to fp32); in fp32 mode this is
-    /// just `matmul(x, w)`. The activation `x` already carries `compute_dtype`.
-    fn linear(&self, x: &Var, w: &Var) -> Var {
-        match self.compute_dtype {
-            DType::F32 => ops::matmul(x, w),
-            DType::BF16 => ops::matmul(x, &ops::cast(w, DType::BF16)),
-            _ => unreachable!(),
-        }
+    /// Enable per-block activation recomputation / gradient checkpointing (Phase
+    /// T13). Builder-style and opt-in; default off keeps the unchanged tape (every
+    /// activation stored). On, each block's forward is wrapped in
+    /// [`xtrain_autodiff::checkpoint`] — exact grads, lower peak activation memory.
+    pub fn with_recompute(mut self, recompute: bool) -> Self {
+        self.recompute = recompute;
+        self
+    }
+
+    pub fn recompute(&self) -> bool {
+        self.recompute
     }
 
     /// All learnable parameters, in a stable order. The optimizer (a hand-written
@@ -171,32 +183,36 @@ impl TinyTransformer {
             h = ops::cast(&h, DType::BF16);
         }
         for b in &self.blocks {
-            // --- Attention sub-block (pre-norm + residual) ---
-            let normed = ops::rms_norm(&h, &self.norm_gamma(&b.attn_norm), self.cfg.eps);
-            let attn = self.attention(b, &normed, batch, seq);
-            h = ops::add(&h, &attn);
-
-            // --- MLP sub-block (pre-norm + residual) ---
-            let normed = ops::rms_norm(&h, &self.norm_gamma(&b.ffn_norm), self.cfg.eps);
-            let mlp = self.swiglu_mlp(b, &normed);
-            h = ops::add(&h, &mlp);
+            h = if self.recompute {
+                // Activation recomputation (T13): run the whole block forward inside
+                // `checkpoint` so its internal activations aren't kept on the tape;
+                // the block forward is re-run in backward to recover the grads. The
+                // 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);
+                xtrain_autodiff::checkpoint::checkpoint(seg, &h, &b.block_params())
+            } else {
+                block_forward(
+                    self.cfg,
+                    self.compute_dtype,
+                    batch,
+                    seq,
+                    &h,
+                    &b.block_params(),
+                )
+            };
         }
 
-        let h = ops::rms_norm(&h, &self.norm_gamma(&self.final_norm), self.cfg.eps);
+        let h = ops::rms_norm(
+            &h,
+            &norm_gamma(self.compute_dtype, &self.final_norm),
+            self.cfg.eps,
+        );
         // lm_head matmul in compute dtype. Logits stay bf16 in bf16 mode — the
         // cross_entropy op upcasts to fp32 internally (no persistent fp32 logits
         // buffer, a real saving at vocab 50257), and its backward casts dx back.
-        self.linear(&h, &self.lm_head) // [batch*seq, vocab]
-    }
-
-    /// A norm/QK-norm gamma in the compute dtype. fp32 master leaf → bf16 (cast
-    /// op, grad upcast) in bf16 mode; identity in fp32 mode.
-    fn norm_gamma(&self, gamma: &Var) -> Var {
-        match self.compute_dtype {
-            DType::F32 => gamma.clone(),
-            DType::BF16 => ops::cast(gamma, DType::BF16),
-            _ => unreachable!(),
-        }
+        linear(self.compute_dtype, &h, &self.lm_head) // [batch*seq, vocab]
     }
 
     /// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
@@ -213,68 +229,146 @@ impl TinyTransformer {
         let logits = self.forward_batched(ids, batch);
         ops::cross_entropy(&logits, targets)
     }
+}
 
-    /// Multi-head causal self-attention over a flattened batch. `x`:[batch*seq,dim]
-    /// (already normed), laid out sequence-major. The Q/K/V/O projections are big
-    /// `[batch*seq, dim]` GEMMs; the scaled-dot-product attention itself runs as a
-    /// fused BATCHED op over the `batch·n_heads` (sequence,head) blocks — each
-    /// attends within its own `[seq,seq]` causal window (NO cross-sequence
-    /// attention), with RoPE positions reset per sequence (`period = seq`). Causal
-    /// masking is applied inside the fused op's softmax kernel (no additive
-    /// `[seq,seq]` mask tensor).
-    fn attention(&self, b: &Block, x: &Var, batch: usize, seq: usize) -> Var {
-        let (nh, hd) = (self.cfg.n_heads, self.cfg.head_dim);
-        let total = batch * seq;
-        let bh = batch * nh;
-        let scale = 1.0 / (hd as f32).sqrt();
+impl Block {
+    /// The block's learnable leaves, in the fixed order the segment forward
+    /// (`block_forward`) indexes them — matches the per-block slice in
+    /// [`TinyTransformer::params`]. This is the param order `checkpoint` passes to
+    /// the recompute closure.
+    fn block_params(&self) -> Vec<Var> {
+        vec![
+            self.attn_norm.clone(),
+            self.wq.clone(),
+            self.wk.clone(),
+            self.wv.clone(),
+            self.q_norm.clone(),
+            self.k_norm.clone(),
+            self.wo.clone(),
+            self.ffn_norm.clone(),
+            self.w_gate.clone(),
+            self.w_up.clone(),
+            self.w_down.clone(),
+        ]
+    }
+}
 
-        // Project, qk-norm + RoPE, then lay out as a batched [B*nh, seq, hd] tensor.
-        //   [B*S,dim] @ [dim,dim] = [B*S,dim]
-        //   reshape   [B*S, nh, hd]
-        //   qk-norm   per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
-        //   rope      [B*S, nh, hd] with per-sequence position (period = seq)
-        //   reshape   [B, S, nh, hd] → transpose(1,2) → [B, nh, S, hd] → [B*nh, S, hd]
-        let to_bh = |proj: Var, norm: Option<&Var>| -> Var {
-            let r = ops::reshape(&proj, &[total, nh, hd]);
-            let r = match norm {
-                // Per-head RMSNorm: flatten the (B*S,nh) head rows, norm over hd,
-                // restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
-                Some(gamma) => {
-                    let flat = ops::reshape(&r, &[total * nh, hd]);
-                    let normed = ops::rms_norm(&flat, &self.norm_gamma(gamma), self.cfg.eps);
-                    let r = ops::reshape(&normed, &[total, nh, hd]);
-                    ops::rope(&r, self.cfg.rope_theta, seq)
-                }
-                None => r,
-            };
-            let r = ops::reshape(&r, &[batch, seq, nh, hd]);
-            let t = ops::transpose_4d12(&r); // [B, nh, S, hd]
-            ops::reshape(&t, &[bh, seq, hd]) // [B*nh, S, hd]
+/// Project `x` (activation, in the compute dtype) by weight `w` (an fp32 master
+/// leaf). In bf16 mode the weight is cast to bf16 via the autograd `cast` op (whose
+/// backward upcasts the grad to fp32); in fp32 mode this is just `matmul(x, w)`.
+fn linear(cdt: DType, x: &Var, w: &Var) -> Var {
+    match cdt {
+        DType::F32 => ops::matmul(x, w),
+        DType::BF16 => ops::matmul(x, &ops::cast(w, DType::BF16)),
+        _ => unreachable!(),
+    }
+}
+
+/// A norm/QK-norm gamma in the compute dtype. fp32 master leaf → bf16 (cast op,
+/// grad upcast) in bf16 mode; identity in fp32 mode.
+fn norm_gamma(cdt: DType, gamma: &Var) -> Var {
+    match cdt {
+        DType::F32 => gamma.clone(),
+        DType::BF16 => ops::cast(gamma, DType::BF16),
+        _ => unreachable!(),
+    }
+}
+
+/// 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 {
+    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) ---
+    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 h = ops::add(h, &attn);
+
+    // --- MLP sub-block (pre-norm + 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);
+    ops::add(&h, &mlp)
+}
+
+/// Multi-head causal self-attention over a flattened batch. `x`:[batch*seq,dim]
+/// (already normed), laid out sequence-major. The Q/K/V/O projections are big
+/// `[batch*seq, dim]` GEMMs; the scaled-dot-product attention itself runs as a
+/// fused BATCHED op over the `batch·n_heads` (sequence,head) blocks — each attends
+/// within its own `[seq,seq]` causal window (NO cross-sequence attention), with
+/// RoPE positions reset per sequence (`period = seq`). Causal masking is applied
+/// inside the fused op's softmax kernel (no additive `[seq,seq]` mask tensor).
+#[allow(clippy::too_many_arguments)]
+fn attention(
+    cfg: Config,
+    cdt: DType,
+    batch: usize,
+    seq: usize,
+    x: &Var,
+    wq: &Var,
+    wk: &Var,
+    wv: &Var,
+    q_norm: &Var,
+    k_norm: &Var,
+    wo: &Var,
+) -> Var {
+    let (nh, hd) = (cfg.n_heads, cfg.head_dim);
+    let total = batch * seq;
+    let bh = batch * nh;
+    let scale = 1.0 / (hd as f32).sqrt();
+
+    // Project, qk-norm + RoPE, then lay out as a batched [B*nh, seq, hd] tensor.
+    //   [B*S,dim] @ [dim,dim] = [B*S,dim]
+    //   reshape   [B*S, nh, hd]
+    //   qk-norm   per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
+    //   rope      [B*S, nh, hd] with per-sequence position (period = seq)
+    //   reshape   [B, S, nh, hd] → transpose(1,2) → [B, nh, S, hd] → [B*nh, S, hd]
+    let to_bh = |proj: Var, norm: Option<&Var>| -> Var {
+        let r = ops::reshape(&proj, &[total, nh, hd]);
+        let r = match norm {
+            // Per-head RMSNorm: flatten the (B*S,nh) head rows, norm over hd,
+            // restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
+            Some(gamma) => {
+                let flat = ops::reshape(&r, &[total * nh, hd]);
+                let normed = ops::rms_norm(&flat, &norm_gamma(cdt, gamma), cfg.eps);
+                let r = ops::reshape(&normed, &[total, nh, hd]);
+                ops::rope(&r, cfg.rope_theta, seq)
+            }
+            None => r,
         };
+        let r = ops::reshape(&r, &[batch, seq, nh, hd]);
+        let t = ops::transpose_4d12(&r); // [B, nh, S, hd]
+        ops::reshape(&t, &[bh, seq, hd]) // [B*nh, S, hd]
+    };
 
-        let q = to_bh(self.linear(x, &b.wq), Some(&b.q_norm));
-        let k = to_bh(self.linear(x, &b.wk), Some(&b.k_norm));
-        let v = to_bh(self.linear(x, &b.wv), None);
+    let q = to_bh(linear(cdt, x, wq), Some(q_norm));
+    let k = to_bh(linear(cdt, x, wk), Some(k_norm));
+    let v = to_bh(linear(cdt, x, wv), None);
 
-        // Fused batched causal SDPA over all B*nh (sequence,head) blocks at once
-        // (2 batched GEMMs + 1 causal-softmax kernel; no per-head/per-seq loop).
-        let out = ops::attention(&q, &k, &v, scale); // [B*nh, S, hd]
+    // Fused batched causal SDPA over all B*nh (sequence,head) blocks at once
+    // (2 batched GEMMs + 1 causal-softmax kernel; no per-head/per-seq loop).
+    let out = ops::attention(&q, &k, &v, scale); // [B*nh, S, hd]
 
-        // Back to [B*S, dim]: [B*nh,S,hd] → [B,nh,S,hd] → transpose(1,2) →
-        // [B,S,nh,hd] → [B*S, dim].
-        let out = ops::reshape(&out, &[batch, nh, seq, hd]);
-        let out = ops::transpose_4d12(&out); // [B, S, nh, hd]
-        let concat = ops::reshape(&out, &[total, nh * hd]); // [B*S, dim]
-        self.linear(&concat, &b.wo) // out projection
-    }
+    // Back to [B*S, dim]: [B*nh,S,hd] → [B,nh,S,hd] → transpose(1,2) →
+    // [B,S,nh,hd] → [B*S, dim].
+    let out = ops::reshape(&out, &[batch, nh, seq, hd]);
+    let out = ops::transpose_4d12(&out); // [B, S, nh, hd]
+    let concat = ops::reshape(&out, &[total, nh * hd]); // [B*S, dim]
+    linear(cdt, &concat, wo) // out projection
+}
 
-    /// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[batch*seq,dim].
-    fn swiglu_mlp(&self, b: &Block, x: &Var) -> Var {
-        let gate = self.linear(x, &b.w_gate); // [seq, ffn_hidden]
-        let up = self.linear(x, &b.w_up); // [seq, ffn_hidden]
-        let act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
-        self.linear(&act, &b.w_down) // [seq, dim]
-    }
+/// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[batch*seq,dim].
+fn swiglu_mlp(cdt: DType, x: &Var, w_gate: &Var, w_up: &Var, w_down: &Var) -> Var {
+    let gate = linear(cdt, x, w_gate); // [seq, ffn_hidden]
+    let up = linear(cdt, x, w_up); // [seq, ffn_hidden]
+    let act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
+    linear(cdt, &act, w_down) // [seq, dim]
 }
 
 /// Materialise a parameter's value back to a host `Vec<f32>` (for the GD step

crates/xtrain-model/tests/recompute.rs

@@ -0,0 +1,156 @@
+// T13 activation-recomputation correctness gate (the HARD gate).
+//
+// Gradient checkpointing is mathematically EXACT: the backward re-runs the same
+// `segment_fn` from the same saved input and the same (unchanged) parameter
+// values, so the recomputed activations equal the originals and the recovered
+// grads equal the non-checkpointed grads — checkpointing trades compute for
+// memory, never correctness. This test makes that a closed loop on-GPU:
+//
+//   build two identical models (same init), one with `--recompute` on, one off,
+//   run the SAME batched loss + backward on both, and assert
+//     1. the forward logits match (recompute doesn't touch forward output)
+//     2. the loss matches
+//     3. EVERY parameter's grad matches within a tight fp tolerance.
+//
+// Composition is covered by parameterising over fp32 AND bf16 (T12): the
+// recompute path is the unchanged block forward, so it runs the same dtype path.
+#![cfg(not(no_cuda))]
+
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
+use xtrain_tensor::{DType, Device};
+
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+fn build(cfg: Config, device: Device, dtype: DType, recompute: bool) -> TinyTransformer {
+    let mut seed = 1u64;
+    let m = TinyTransformer::new(cfg, device, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.08)
+        }
+    });
+    m.with_compute_dtype(dtype).with_recompute(recompute)
+}
+
+fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
+    t.to_device(Device::Cpu).as_slice::<f32>().to_vec()
+}
+
+fn run(dtype: DType, logit_tol: f32, grad_tol: f32) {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    // A few layers so checkpointing actually wraps multiple blocks.
+    let mut cfg = Config::tiny();
+    cfg.vocab = 16;
+    cfg.n_layers = 4;
+    let batch = 3usize;
+    let seq = 6usize;
+    let seqs: Vec<Vec<i32>> = (0..batch)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
+                .collect()
+        })
+        .collect();
+    let tgts: Vec<Vec<i32>> = (0..batch)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
+                .collect()
+        })
+        .collect();
+    let ids = batched_ids_tensor(&seqs, device);
+    let tgt = batched_ids_tensor(&tgts, device);
+
+    // --- recompute OFF (reference) ---
+    let off = build(cfg, device, dtype, false);
+    let off_logits = host(&off.forward_batched(&ids, batch).value());
+    let off_loss = off.loss_batched(&ids, &tgt, batch);
+    let off_loss_val = host(&off_loss.value())[0];
+    off_loss.backward();
+    let off_grads: Vec<Vec<f32>> = off
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().expect("off grad")))
+        .collect();
+
+    // --- recompute ON ---
+    let on = build(cfg, device, dtype, true);
+    let on_logits = host(&on.forward_batched(&ids, batch).value());
+    let on_loss = on.loss_batched(&ids, &tgt, batch);
+    let on_loss_val = host(&on_loss.value())[0];
+    on_loss.backward();
+    let on_grads: Vec<Vec<f32>> = on
+        .params()
+        .iter()
+        .map(|p| host(&p.grad().expect("on grad")))
+        .collect();
+
+    // 1. Forward logits — recompute must not change the forward output.
+    let logit_rel = off_logits
+        .iter()
+        .zip(&on_logits)
+        .map(|(a, b)| (a - b).abs() / a.abs().max(1e-4))
+        .fold(0.0f32, f32::max);
+    // 2. Loss.
+    let loss_rel = (off_loss_val - on_loss_val).abs() / off_loss_val.abs().max(1e-4);
+    println!(
+        "[{dtype:?}] recompute on/off: loss {off_loss_val:.6}/{on_loss_val:.6} (rel {loss_rel:.2e}), \
+         logits max rel {logit_rel:.2e}"
+    );
+    assert!(
+        logit_rel < logit_tol,
+        "[{dtype:?}] logits diverged: {logit_rel:.2e}"
+    );
+    assert!(
+        loss_rel < logit_tol,
+        "[{dtype:?}] loss diverged: {loss_rel:.2e}"
+    );
+
+    // 3. Every parameter grad — the load-bearing gate.
+    let mut max_grad_rel = 0.0f32;
+    for (off_g, on_g) in off_grads.iter().zip(&on_grads) {
+        for (a, b) in off_g.iter().zip(on_g) {
+            let rel = (a - b).abs() / a.abs().max(1e-3);
+            max_grad_rel = max_grad_rel.max(rel);
+        }
+    }
+    println!("[{dtype:?}] recompute on/off: grad max rel err = {max_grad_rel:.3e}");
+    assert!(
+        max_grad_rel < grad_tol,
+        "[{dtype:?}] recompute grads diverged from non-recompute: {max_grad_rel:.3e}"
+    );
+}
+
+#[test]
+fn recompute_matches_non_recompute_fp32() {
+    // fp32: recompute runs the identical deterministic kernels → grads match to
+    // (near) bit-exact; allow a hair for any nondeterministic GPU reduction.
+    run(DType::F32, 1e-5, 1e-4);
+}
+
+#[test]
+fn recompute_matches_non_recompute_bf16() {
+    // bf16 (T12 composition): same bf16 path on recompute. The recompute is still
+    // exact w.r.t. the bf16 forward, so on/off match tightly (looser tol only for
+    // bf16 rounding, not for any recompute discrepancy).
+    run(DType::BF16, 5e-3, 5e-3);
+}

crates/xtrain-train/src/bin/train.rs

@@ -112,6 +112,10 @@ fn main() {
     // bf16 mixed precision (Phase T12): fp32 master weights, bf16 linears +
     // activations. Opt-in; default fp32 reproduces v0–v4 numerics.
     let bf16 = args.iter().any(|a| a == "--bf16");
+    // Activation recomputation (Phase T13): per-block gradient checkpointing —
+    // exact grads, lower peak activation memory (lets dim1024 batch32 fit). Opt-in;
+    // default off stores every activation (unchanged numerics).
+    let recompute = args.iter().any(|a| a == "--recompute");
     let ckpt: PathBuf = PathBuf::from(
         args.iter()
             .position(|a| a == "--ckpt")
@@ -175,6 +179,10 @@ fn main() {
         model = model.with_compute_dtype(DType::BF16);
         println!("bf16 mixed precision: ON (fp32 master weights)");
     }
+    if recompute {
+        model = model.with_recompute(true);
+        println!("activation recompute: ON (per-block gradient checkpointing)");
+    }
 
     // Eval-only mode: load a checkpoint and score it on the held-out val set, then
     // exit. Used to put an EXISTING model (e.g. v0) and a new one on the same

docs/12-activation-recompute.md

@@ -0,0 +1,88 @@
+# Phase T13: 激活重计算（gradient checkpointing）— Design Document
+
+> KI-3 的具体落地。autograd tape 为反向保存了所有中间激活；dim768/bf16 在单卡 32GB 能跑 batch32（T12 解 OOM），但**容量轴放大到 dim1024 会再次 OOM**——激活显存随 dim 线性增长。激活重计算用「多一次前向」换显存：段内激活不在前向保存，反向时**重算该段前向**重建局部 tape 再回传。峰值激活从「所有 block 同时在显存」降到「~一个 block + 每 block 的输入」→ dim1024 batch32 装得下。
+
+## Goal
+
+在**不动非重计算路径任何数值**的前提下，新增一个 **opt-in 的 per-block 激活重计算**：
+
+1. **正确性硬闸门（exact）**：重计算是**数学精确**的——同一段前向、同一输入、同一（未变的）参数值、确定性 kernel ⟹ 重算出的激活与原激活逐位相同，回传的梯度与非重计算版一致。直接的 **on-vs-off 梯度对拍**（紧容差）+ 全套回归（T4 grad-check、T5 overfit+PyTorch 对拍、T6 AdamW、T8 DDP loss-match+跨 rank、xserv 闭环）开 `--recompute` 全绿。**绝不提交一个改变梯度的重计算。**
+2. **显存（payoff）**：测 dim768 batch32 峰值显存 on vs off（应降）；确认 **dim1024 batch32 现在装得下**（不开重计算时 OOM，开了 fit）。
+3. **吞吐**：测 tok/s on vs off（多一次前向，预计慢 ~20–35%）——报告 compute/memory 权衡。
+
+## 什么是激活重计算
+
+反向传播需要前向的中间激活（如 SwiGLU 的 `gate`、attention 的 `probs`）来算梯度。define-by-run 的 tape 默认把它们全部留在显存，直到对应 op 的 backward 跑完。模型越深、激活越多，峰值显存越高。
+
+**梯度检查点**把模型切成若干**段（segment）**。前向时，段内 op **不记到 tape**（detached / no-grad），只保留**段的输入**（参数作为 leaf 本就常驻，不算激活）。反向时，当段的 output-grad 到达，**从保存的输入重跑该段前向**（输入作为 require-grad 的 leaf），用上游 grad seed 重算出的 output，在局部 tape 上回传，得到输入梯度（并累加参数梯度），然后释放局部 tape。
+
+代价：每段多一次前向（约 +1/3 的总 FLOPs，因为反向本就 ~2× 前向）。收益：峰值激活从「所有段」降到「~一段 + 每段输入」。
+
+**切粒度 = 每个 transformer block**：一个 block（attention 子块 + MLP 子块 + 两个残差）是天然的段边界，输入/输出都是 `[batch*seq, dim]` 的残差流张量，接口最干净。
+
+## Module Layout（surgical：非重计算路径逐字节不动）
+
+### 1. `xtrain-autodiff::checkpoint` — `checkpoint` 高阶原语
+
+新增 `checkpoint(segment_fn, input, params) -> Var`，类比 `torch.utils.checkpoint`：
+
+- `segment_fn: Fn(&Var, &[Var]) -> Var`——从单个输入 `x` 和参数 slice `p` 构建段前向、返回段输出。必须**确定性**、只依赖 `x` 和 `p`（这是重算精确的前提）。
+- **前向（不 tape 内部）**：把 `input`/`params` detach 成新 leaf（`Var::leaf(v.value())`），跑 `segment_fn` 得到 `out_local`，**只取 `out_local.value()`**。局部 `Var` 出作用域即 drop → 段内激活立即释放。checkpoint 节点的 parents = `[input, ..params]`（参数梯度落进优化器读的同一批 leaf）。
+- **反向（重算）**：闭包捕获 `Rc<segment_fn>`。给定 `dout`，从 `parents` 当前值重建 detached leaf，重跑 `segment_fn` 重建局部 tape，调 `out_local.backward_seeded(dout)`，再把 `x_det.grad()` push 给 `parents[0]`、各 `param_det.grad()` push 给对应参数 parent。闭包结束 → 局部 tape drop → 重算激活释放。
+
+### 2. `xtrain-autodiff::tape` — `backward_seeded`
+
+引擎原 `backward()` 只能从标量 root 出发（seed `ones_like` + 断言 numel==1）。段输出一般**非标量**，故新增 `Var::backward_seeded(seed)`：同样的 topo + 反向遍历，但用显式上游 grad seed（不断言标量）。`backward()` 退化为「seed ones」的薄包装——标量 loss 路径逐字节不变。
+
+### 3. `xtrain-model::TinyTransformer` — per-block 包裹 + `recompute` 开关
+
+- 把 block 前向体抽成**自由函数** `block_forward(cfg, compute_dtype, batch, seq, input, params)`（不借 `&self`，才能在反向闭包里重跑）；同时把 `linear / norm_gamma / attention / swiglu_mlp` 改成参数化 `(cfg, compute_dtype)` 的自由函数。`Block::block_params()` 给出 11 个 leaf 的固定序（与 `params()` 每 block 段一致）。
+- `forward_batched` 的 block 循环：`recompute` 开 → `checkpoint(seg, &h, &b.block_params())`；关 → 直接 `block_forward(...)`（**与之前完全同图**）。
+- `with_recompute(bool)` builder（opt-in，默认关 = 原 tape，数值逐字节同）。
+
+### 4. `xtrain-train` / `xtrain-distributed` — `--recompute` flag
+
+- `bin/train` / `bin/train_ddp` 加 `--recompute`，调 `model.with_recompute(true)`。AdamW / clip / checkpoint / DDP all-reduce **不改**（梯度语义与非重计算一致）。
+
+## Key Design Decisions
+
+- **切粒度 = 每个 block**：接口最干净（输入输出都是残差流 `[B*S, dim]`），且峰值激活降到 ~1 个 block。比「整模型一段」省得多，比「逐 op」简单得多。
+- **参数作为 checkpoint 节点的 parents**：参数 leaf 跨前向/反向不变（只 grad 槽变），重算用**当前**参数值即原值。把它们列为 parents，重算恢复的参数梯度直接 push 到优化器读的同一批 leaf → **DDP/AdamW 零改动**。
+- **detached leaf 隔离局部 tape**：前向/反向都把输入和参数 detach 成新 leaf，使 `segment_fn` 构建的图与外层 tape 断开。前向丢弃局部图（释放激活）；反向局部图回传完即 drop（释放重算激活）。
+- **`backward_seeded` 而非改 `backward`**：段输出非标量，需要用上游 output-grad 作 seed 回传局部 tape。新增方法、原标量 `backward()` 不动。
+- **重算精确 → 梯度逐位一致（硬闸门）**：同一 `segment_fn`、同一输入值、同一参数值、确定性前向 kernel ⟹ 重算 output 与原 output 相同，局部反向就是该段的普通解析反向。故输入/参数梯度与非重计算版一致——这是绝不能违反的闸门。
+
+## 与 bf16 / DDP / batched 的组合
+
+- **bf16（T12）**：`segment_fn` 就是不变的 block 前向，重算跑**同一条 bf16 路径**；`cast` 算子的 grad 升精度桥（bf16→fp32）照常。重计算节点的参数 parents 是 fp32 master leaf，恢复的是 fp32 梯度。on-vs-off 对拍同时跑 fp32 和 bf16 两路。
+- **DDP（T8）**：每个 rank 独立 checkpoint 自己的前向/反向；恢复的参数梯度落进各 rank 的 `.grad()` 槽，再被 all-reduce 取均值——分布式路径不感知重计算。
+- **batched（T10）**：段输入/输出透明带 `[batch*seq, …]` batch 维；`checkpoint` 与形状无关。
+
+## 验证方法
+
+### 1. 正确性（exact，硬闸门）
+
+- **on-vs-off 梯度对拍**（`crates/xtrain-model/tests/recompute.rs`）：同 init 建两个模型（recompute on/off），跑同一 batched loss+backward，断言**前向 logits、loss、每个参数梯度**在紧容差内一致——参数化跑 **fp32 和 bf16** 两路。fp32 期望近逐位（容差 1e-4），bf16 仅放松到 bf16 舍入级（非重计算误差）。
+- **全套回归开 `--recompute`**：T4 15 算子 grad-check、T5 overfit 27/27 + PyTorch 对拍、T6 AdamW、T8 DDP loss-match + 跨 rank、**xserv 闭环 md5**——全绿。
+
+### 2. 显存（payoff）
+
+- dash5 1× RTX 5090 32GB，dim768/18L batch32 seq256，bf16：测峰值显存 recompute on vs off（应降）。
+- **dim1024 batch32**：先验证不开重计算 **OOM**，再验证开了 **fit**——capture 实际 `nvidia-smi` 峰值。
+
+### 3. 吞吐
+
+- 同 config 测 steady-state tok/s recompute on vs off，报告慢多少（预计 ~20–35%，多一次前向）。
+
+## 实测结果（dash5 1× RTX 5090 32GB, dim768/18L/24h×32 ffn2048 seq256, bf16, steady-state）
+
+> 待 dash5 实跑回填。
+
+| config | per-rank batch | 峰值显存 | tok/s | fits 32GB? |
+|---|---|---|---|---|
+| dim768 recompute **off** | 32 | TBD | TBD | ✅ |
+| dim768 recompute **on** | 32 | **TBD（↓）** | TBD（↓~xx%） | ✅ |
+| **dim1024** recompute **off** | 32 | — | — | ❌ **OOM** |
+| **dim1024** recompute **on** | 32 | **TBD** | TBD | ✅ **解 OOM** |
+
+**正确性**：on-vs-off 梯度 max rel = TBD（fp32）/ TBD（bf16）；全套回归 + xserv 闭环全绿。