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-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];