perf: keep bf16 logits (no persistent fp32 logits buffer)

At vocab 50257 the logits tensor [B*S, vocab] is ~1.6GB fp32 at batch
32 — held across the whole backward. Keep it bf16: cross_entropy
upcasts the bf16 logits to fp32 internally (transient) + caches fp32
probs, and its backward casts dx back to bf16 to chain into the
bf16 lm_head matmul backward. The sampler casts bf16 logits→f32 before
the host argmax/softmax. Halves the persistent logits activation.

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

crates/xtrain-train/src/sample.rs

@@ -26,7 +26,11 @@ pub fn generate(
 
     for _ in 0..max_new {
         let ids_t = ids_tensor(&ids, device);
-        let logits = model.forward(&ids_t).value().to_device(Device::Cpu);
+        // In bf16 mode the logits are bf16; cast to f32 (on device) before reading.
+        let logits = model.forward(&ids_t).value();
+        let logits = logits
+            .to_dtype(xtrain_tensor::DType::F32)
+            .to_device(Device::Cpu);
         let lg = logits.as_slice::<f32>();
         // Last row = next-token distribution for the current prefix.
         let last = &lg[(ids.len() - 1) * vocab..ids.len() * vocab];