train: loop + checkpoint save/load + sampler + train binary

Training loop (train_loop.rs): sample batch_size sequences, forward loss +
backward (tape SUMs grads), clip_grad_norm with ×1/batch averaging, AdamW step
with scheduled lr, zero_grad; logs loss/lr/gnorm/tok-s and checkpoints
periodically; returns the loss trace.

Checkpoint (checkpoint.rs): flat little-endian dump of params() in order
(magic/version/count + per-param ndim/dims/f32 data); load_into validates and
overwrites a matching model's params via set_value (exact f32 round-trip).

Sampler (sample.rs): autoregressive greedy / temperature generation — re-runs
forward on the growing prefix (model is single-sequence, RoPE pos=row).

bin/train.rs: end-to-end entry — load tokenizer+corpus, train a tiny 4-layer
model for a bounded budget, checkpoint, print samples. no_cuda stub keeps it
buildable on a GPU-less host.

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

crates/xtrain-train/src/train_loop.rs

@@ -0,0 +1,107 @@
+//! 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 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.
+
+#![cfg(not(no_cuda))]
+
+use std::path::PathBuf;
+use std::time::Instant;
+
+use xtrain_model::{TinyTransformer, ids_tensor};
+use xtrain_optim::AdamW;
+use xtrain_tensor::Device;
+
+use crate::checkpoint;
+use crate::clip::clip_grad_norm;
+use crate::data::Corpus;
+use crate::schedule::LrSchedule;
+
+/// Knobs for a training run.
+pub struct TrainConfig {
+    pub seq_len: usize,
+    pub batch_size: usize,
+    pub steps: usize,
+    pub schedule: LrSchedule,
+    pub weight_decay: f32,
+    pub max_grad_norm: f32,
+    pub log_every: usize,
+    /// Optional checkpoint path written every `ckpt_every` steps (and at the end).
+    pub ckpt_path: Option<PathBuf>,
+    pub ckpt_every: usize,
+    /// Seed for reproducible sequence sampling.
+    pub seed: u64,
+}
+
+/// Train `model` on `corpus` for `cfg.steps` AdamW steps. Returns the per-step
+/// loss trace (one mean loss per step, read from the first sequence of the
+/// batch — cheap and representative). Logs progress and checkpoints as configured.
+pub fn train(
+    model: &TinyTransformer,
+    device: Device,
+    corpus: &Corpus,
+    cfg: &TrainConfig,
+) -> Vec<f32> {
+    let params = model.params();
+    let mut opt = AdamW::new(cfg.schedule.max_lr, cfg.weight_decay);
+    let mut rng = cfg.seed;
+    let mut losses = Vec::with_capacity(cfg.steps);
+    let inv_batch = 1.0 / cfg.batch_size as f32;
+    let start = Instant::now();
+    let mut tokens_seen: u64 = 0;
+
+    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;
+        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;
+        }
+        step_loss *= inv_batch;
+        losses.push(step_loss);
+
+        // Average the summed grads (×1/batch) and clip to the global norm.
+        let gnorm = clip_grad_norm(&params, cfg.max_grad_norm, inv_batch);
+        opt.step(lr, &params);
+        for p in &params {
+            p.zero_grad();
+        }
+
+        if step % cfg.log_every == 0 || step == cfg.steps - 1 {
+            let elapsed = start.elapsed().as_secs_f32();
+            let tps = tokens_seen as f32 / elapsed.max(1e-6);
+            println!(
+                "step {step:5}/{}: loss {step_loss:.4} lr {lr:.2e} gnorm {gnorm:.3} \
+                 ({tps:.0} tok/s)",
+                cfg.steps
+            );
+        }
+
+        if let Some(path) = &cfg.ckpt_path {
+            if cfg.ckpt_every > 0 && (step + 1) % cfg.ckpt_every == 0 {
+                checkpoint::save(path, &params).expect("checkpoint save");
+            }
+        }
+    }
+
+    if let Some(path) = &cfg.ckpt_path {
+        checkpoint::save(path, &params).expect("final checkpoint save");
+        println!("saved checkpoint → {}", path.display());
+    }
+    losses
+}
+
+fn read_scalar(v: &xtrain_autodiff::tape::Var) -> f32 {
+    v.value().to_device(Device::Cpu).as_slice::<f32>()[0]
+}