train: real batched step (drop loop+SUM)

Feed a real batch of B sequences as ONE batched forward/backward, replacing the
"loop B times + let the tape SUM grads + clip ×1/B" hack. CE mean over B*S rows
is already the batch-mean loss, so backward yields the batch-mean gradient
directly → clip pre-scale = 1.0.

DDP stays equivalent: each rank runs one batched forward over its b_local =
B_global/world sequences (local-mean grad Σ_local/b_local); all_reduce_average
(sum across ranks /world) = Σ_global/B_global = global batch-mean → clip
pre-scale 1.0. The ddp_correctness single-GPU baseline batches the same way.
DDP loss matches single-GPU 5.7e-7, cross-rank params bit-identical (0.0).

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

crates/xtrain-train/src/train_loop.rs

@@ -1,18 +1,20 @@
-//! The training loop: sample sequences → forward `loss` → backward → grad clip
-//! (with batch averaging) → AdamW step → zero grads; with an LR schedule,
-//! periodic loss logging, and periodic checkpointing.
+//! The training loop: sample a batch of sequences → ONE batched forward `loss` →
+//! backward → grad clip → AdamW step → zero grads; with an LR schedule, periodic
+//! loss logging, and periodic checkpointing.
 //!
-//! The T5 model is single-sequence, so a "batch" of `batch_size` sequences is
-//! handled by running forward+backward on each and letting the tape SUM their
-//! grads (its fan-out rule); the clip pass then multiplies by `1/batch_size` to
-//! recover the batch-mean gradient before clipping + the optimizer step.
+//! Since T10 the model is batched (`loss_batched`): `batch_size` sequences are
+//! flattened to `[batch*seq]` and run as a SINGLE forward/backward, so the linear
+//! projections become big `[batch*seq, dim]` GEMMs that fill the GPU. The
+//! cross-entropy mean is over all `batch*seq` rows — already the batch-mean loss,
+//! so backward yields the batch-mean gradient directly (clip pre-scale = 1.0; no
+//! more "loop B times + SUM + ×1/batch" hack).
 
 #![cfg(not(no_cuda))]
 
 use std::path::PathBuf;
 use std::time::Instant;
 
-use xtrain_model::{TinyTransformer, ids_tensor};
+use xtrain_model::{TinyTransformer, batched_ids_tensor, ids_tensor};
 use xtrain_optim::GpuAdamW;
 use xtrain_tensor::Device;
 
@@ -67,7 +69,6 @@ pub fn train(
     let mut losses = Vec::with_capacity(cfg.steps);
     let mut evals = Vec::new();
     let mut best_val: Option<f32> = None;
-    let inv_batch = 1.0 / cfg.batch_size as f32;
     let start = Instant::now();
     let mut tokens_seen: u64 = 0;
     // Best-val checkpointing only kicks in when we actually evaluate.
@@ -76,22 +77,26 @@ pub fn train(
     for step in 0..cfg.steps {
         let lr = cfg.schedule.lr(step);
 
-        // Accumulate grads over `batch_size` sequences (tape SUMs them).
-        let mut step_loss = 0.0f32;
+        // Sample `batch_size` sequences and run them as ONE batched forward/
+        // backward. The CE mean over all batch*seq rows is the batch-mean loss, so
+        // backward already yields the batch-mean gradient (clip pre-scale = 1.0).
+        let mut inputs = Vec::with_capacity(cfg.batch_size);
+        let mut targets_v = Vec::with_capacity(cfg.batch_size);
         for _ in 0..cfg.batch_size {
             let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
-            let ids = ids_tensor(&input, device);
-            let targets = ids_tensor(&target, device);
-            let loss = model.loss(&ids, &targets);
-            step_loss += read_scalar(&loss);
-            loss.backward();
-            tokens_seen += cfg.seq_len as u64;
+            inputs.push(input);
+            targets_v.push(target);
         }
-        step_loss *= inv_batch;
+        let ids = batched_ids_tensor(&inputs, device);
+        let targets = batched_ids_tensor(&targets_v, device);
+        let loss = model.loss_batched(&ids, &targets, cfg.batch_size);
+        let step_loss = read_scalar(&loss);
+        loss.backward();
+        tokens_seen += (cfg.batch_size * cfg.seq_len) as u64;
         losses.push(step_loss);
 
-        // Average the summed grads (×1/batch) and clip to the global norm.
-        let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, inv_batch);
+        // Backward already produced the batch-mean gradient — just clip it.
+        let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, 1.0);
         opt.step(lr, &params);
         for p in &params {
             p.zero_grad();