train: --bf16 flag (fp32-master AMP) + bf16 correctness test

- TinyTransformer::with_compute_dtype(BF16): embedding stays fp32
  master then casts to bf16; each linear casts its fp32 weight to bf16
  on the fly; logits cast back to fp32 for cross-entropy. Default F32
  reproduces the v0-v4 forward graph bit-for-bit.
- --bf16 flag on bin/train and bin/train_ddp (off by default).
- tests/bf16.rs: same fp32 master weights run fp32 vs bf16; assert
  loss/logits/grads within a loose bf16 tol, no NaN, and grads are
  fp32 (master untouched).

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

crates/xtrain-model/src/model.rs

@@ -5,7 +5,7 @@
 use crate::config::Config;
 use xtrain_autodiff::ops;
 use xtrain_autodiff::tape::Var;
-use xtrain_tensor::{Device, Tensor};
+use xtrain_tensor::{DType, Device, Tensor};
 
 /// One decoder block's learnable tensors.
 struct Block {
@@ -30,6 +30,13 @@ pub struct TinyTransformer {
     blocks: Vec<Block>,
     final_norm: Var, // [dim]
     lm_head: Var,    // [dim, vocab]
+    /// Compute dtype for the forward graph (Phase T12). `F32` (default) = the
+    /// original path, bit-identical to T10/T11. `BF16` = mixed precision: the
+    /// parameter leaves stay fp32 (master), but each linear's weight is cast to
+    /// bf16 on the fly and the activation stream flows bf16 (see
+    /// `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,
 }
 
 impl TinyTransformer {
@@ -71,6 +78,7 @@ impl TinyTransformer {
             blocks,
             final_norm,
             lm_head,
+            compute_dtype: DType::F32,
         }
     }
 
@@ -78,6 +86,35 @@ impl TinyTransformer {
         &self.cfg
     }
 
+    /// Set the forward compute dtype (Phase T12). `BF16` enables mixed precision
+    /// (fp32 master weights, bf16 linears + activations); `F32` (the default) is
+    /// the unchanged full-precision path. Builder-style so existing call sites
+    /// that don't opt in keep the fp32 numerics bit-for-bit.
+    pub fn with_compute_dtype(mut self, dtype: DType) -> Self {
+        assert!(
+            matches!(dtype, DType::F32 | DType::BF16),
+            "compute_dtype must be F32 or BF16"
+        );
+        self.compute_dtype = dtype;
+        self
+    }
+
+    pub fn compute_dtype(&self) -> DType {
+        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!(),
+        }
+    }
+
     /// 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()`.
@@ -127,21 +164,42 @@ impl TinyTransformer {
         );
         let seq = total / batch;
 
-        let mut h = ops::embedding(&self.embed, ids); // [batch*seq, dim]
+        // 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 {
             // --- Attention sub-block (pre-norm + residual) ---
-            let normed = ops::rms_norm(&h, &b.attn_norm, self.cfg.eps);
+            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, &b.ffn_norm, self.cfg.eps);
+            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);
         }
 
-        let h = ops::rms_norm(&h, &self.final_norm, self.cfg.eps);
-        ops::matmul(&h, &self.lm_head) // [batch*seq, vocab]
+        let h = ops::rms_norm(&h, &self.norm_gamma(&self.final_norm), self.cfg.eps);
+        // lm_head matmul in compute dtype; cast logits back to fp32 for CE.
+        let logits = self.linear(&h, &self.lm_head); // [batch*seq, vocab]
+        if self.compute_dtype == DType::BF16 {
+            ops::cast(&logits, DType::F32)
+        } else {
+            logits
+        }
+    }
+
+    /// 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!(),
+        }
     }
 
     /// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
@@ -186,7 +244,7 @@ impl TinyTransformer {
                 // 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, gamma, self.cfg.eps);
+                    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)
                 }
@@ -197,9 +255,9 @@ impl TinyTransformer {
             ops::reshape(&t, &[bh, seq, hd]) // [B*nh, S, hd]
         };
 
-        let q = to_bh(ops::matmul(x, &b.wq), Some(&b.q_norm));
-        let k = to_bh(ops::matmul(x, &b.wk), Some(&b.k_norm));
-        let v = to_bh(ops::matmul(x, &b.wv), None);
+        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);
 
         // 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).
@@ -210,15 +268,15 @@ impl TinyTransformer {
         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]
-        ops::matmul(&concat, &b.wo) // out projection
+        self.linear(&concat, &b.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 = ops::matmul(x, &b.w_gate); // [seq, ffn_hidden]
-        let up = ops::matmul(x, &b.w_up); // [seq, ffn_hidden]
+        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
-        ops::matmul(&act, &b.w_down) // [seq, dim]
+        self.linear(&act, &b.w_down) // [seq, dim]
     }
 }