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-train/src/bin/train.rs

@@ -31,6 +31,8 @@ use xtrain_cuda::device;
 #[cfg(not(no_cuda))]
 use xtrain_model::{Config, TinyTransformer};
 #[cfg(not(no_cuda))]
+use xtrain_tensor::DType;
+#[cfg(not(no_cuda))]
 use xtrain_tensor::Device;
 #[cfg(not(no_cuda))]
 use xtrain_train::data::Corpus;
@@ -107,6 +109,9 @@ fn main() {
     let val_tokens: usize = flag(&args, "--val-tokens", 0);
     let eval_every: usize = flag(&args, "--eval-every", 0);
     let eval_batches: usize = flag(&args, "--eval-batches", 64);
+    // 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");
     let ckpt: PathBuf = PathBuf::from(
         args.iter()
             .position(|a| a == "--ckpt")
@@ -155,7 +160,7 @@ fn main() {
     );
 
     let mut seed = 1u64;
-    let model = TinyTransformer::new(cfg, device, |shape| {
+    let mut model = TinyTransformer::new(cfg, device, |shape| {
         seed = seed.wrapping_add(1);
         let n: usize = shape.iter().product();
         if shape.len() == 1 {
@@ -166,6 +171,10 @@ fn main() {
             fill(n, seed, 0.04)
         }
     });
+    if bf16 {
+        model = model.with_compute_dtype(DType::BF16);
+        println!("bf16 mixed precision: ON (fp32 master weights)");
+    }
 
     // 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