model: per-block activation recompute (--recompute)

Wrap each transformer block's forward in the checkpoint primitive when
recompute is enabled (Phase T13 / KI-3). To make the block forward a pure
segment fn (no `&self` borrow, so it can re-run in the backward closure),
extract the block body + its helpers (linear / norm_gamma / attention /
swiglu_mlp) into free functions parameterised by (cfg, compute_dtype) and add
`Block::block_params()` (the 11 leaves in the params() per-block order). The
non-recompute path calls `block_forward` directly — identical graph to before.

- `TinyTransformer::with_recompute(bool)` builder (opt-in; default off keeps the
  unchanged tape / bit-identical numerics).
- `--recompute` flag wired into bin/train and bin/train_ddp (DDP: each rank
  checkpoints independently).

Correctness gate: tests/recompute.rs builds two identical models (recompute
on/off), runs the same batched loss+backward, and asserts the forward logits,
the loss, and EVERY parameter grad match within tight fp tol — parameterised
over fp32 and bf16 (T12 composition).

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

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