model: tiny RoPE+RMSNorm+SwiGLU transformer + overfit test

New crate xtrain-model: a from-scratch decoder built entirely from the
autodiff op set.
- Config (tiny: dim=32, 2 layers, 2 heads, head_dim=16, ffn=64).
- TinyTransformer: embedding -> N x {pre-RMSNorm -> multi-head causal
  attention (RoPE, additive causal mask, per-head SDPA) -> residual;
  pre-RMSNorm -> SwiGLU MLP -> residual} -> final RMSNorm -> LM head.
  x@W weight convention (engine GEMM is plain A@B); dim=n_heads*head_dim.
- params()/zero_grad-able leaves for the optimizer; param_to_host export.
- overfit test: char-level bring-up (embedded text -> vocab -> shifted
  targets), minimal hand-written GD (p -= lr*grad) memorises one fixed
  batch -> loss ~0 + greedy argmax matches targets. End-to-end fwd+bwd
  correctness signal. Gated #![cfg(not(no_cuda))].

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

crates/xtrain-model/src/model.rs

@@ -0,0 +1,205 @@
+//! The tiny transformer forward graph + parameter container (Phase T5).
+
+#![cfg(not(no_cuda))]
+
+use crate::config::Config;
+use xtrain_autodiff::ops;
+use xtrain_autodiff::tape::Var;
+use xtrain_tensor::{Device, Tensor};
+
+/// One decoder block's learnable tensors.
+struct Block {
+    attn_norm: Var, // [dim]
+    wq: Var,        // [dim, dim]
+    wk: Var,        // [dim, dim]
+    wv: Var,        // [dim, dim]
+    wo: Var,        // [dim, dim]
+    ffn_norm: Var,  // [dim]
+    w_gate: Var,    // [dim, ffn_hidden]
+    w_up: Var,      // [dim, ffn_hidden]
+    w_down: Var,    // [ffn_hidden, dim]
+}
+
+/// A tiny RoPE+RMSNorm+SwiGLU decoder. Holds every parameter as a leaf [`Var`];
+/// `forward` builds an autograd graph over them.
+pub struct TinyTransformer {
+    cfg: Config,
+    embed: Var, // [vocab, dim]
+    blocks: Vec<Block>,
+    final_norm: Var, // [dim]
+    lm_head: Var,    // [dim, vocab]
+    device: Device,
+}
+
+impl TinyTransformer {
+    /// Build a model with parameters initialised from `init(shape) -> host data`.
+    /// The caller controls initialisation (deterministic for tests / PyTorch
+    /// parity). `init` receives the logical shape and returns row-major data.
+    pub fn new(cfg: Config, device: Device, mut init: impl FnMut(&[usize]) -> Vec<f32>) -> Self {
+        let leaf = |data: Vec<f32>, shape: &[usize]| -> Var {
+            Var::leaf(Tensor::from_slice(&data, shape).to_device(device))
+        };
+        let mut mk = |shape: &[usize]| -> Var {
+            let data = init(shape);
+            assert_eq!(data.len(), shape.iter().product::<usize>(), "init size");
+            leaf(data, shape)
+        };
+
+        let embed = mk(&[cfg.vocab, cfg.dim]);
+        let blocks = (0..cfg.n_layers)
+            .map(|_| Block {
+                attn_norm: mk(&[cfg.dim]),
+                wq: mk(&[cfg.dim, cfg.dim]),
+                wk: mk(&[cfg.dim, cfg.dim]),
+                wv: mk(&[cfg.dim, cfg.dim]),
+                wo: mk(&[cfg.dim, cfg.dim]),
+                ffn_norm: mk(&[cfg.dim]),
+                w_gate: mk(&[cfg.dim, cfg.ffn_hidden]),
+                w_up: mk(&[cfg.dim, cfg.ffn_hidden]),
+                w_down: mk(&[cfg.ffn_hidden, cfg.dim]),
+            })
+            .collect();
+        let final_norm = mk(&[cfg.dim]);
+        let lm_head = mk(&[cfg.dim, cfg.vocab]);
+
+        Self {
+            cfg,
+            embed,
+            blocks,
+            final_norm,
+            lm_head,
+            device,
+        }
+    }
+
+    pub fn config(&self) -> &Config {
+        &self.cfg
+    }
+
+    /// 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()`.
+    pub fn params(&self) -> Vec<Var> {
+        let mut ps = vec![self.embed.clone()];
+        for b in &self.blocks {
+            ps.extend([
+                b.attn_norm.clone(),
+                b.wq.clone(),
+                b.wk.clone(),
+                b.wv.clone(),
+                b.wo.clone(),
+                b.ffn_norm.clone(),
+                b.w_gate.clone(),
+                b.w_up.clone(),
+                b.w_down.clone(),
+            ]);
+        }
+        ps.push(self.final_norm.clone());
+        ps.push(self.lm_head.clone());
+        ps
+    }
+
+    /// Forward over a single sequence of token `ids` (`[seq]` I32 on this
+    /// model's device). Returns the logits [`Var`] of shape `[seq, vocab]`.
+    pub fn forward(&self, ids: &Tensor) -> Var {
+        let seq = ids.shape()[0];
+        let mask = self.causal_mask(seq);
+
+        let mut h = ops::embedding(&self.embed, ids); // [seq, dim]
+        for b in &self.blocks {
+            // --- Attention sub-block (pre-norm + residual) ---
+            let normed = ops::rms_norm(&h, &b.attn_norm, self.cfg.eps);
+            let attn = self.attention(b, &normed, &mask, 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 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) // [seq, vocab]
+    }
+
+    /// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
+    pub fn loss(&self, ids: &Tensor, targets: &Tensor) -> Var {
+        let logits = self.forward(ids);
+        ops::cross_entropy(&logits, targets)
+    }
+
+    /// Multi-head causal self-attention. `x`:[seq,dim] (already normed).
+    fn attention(&self, b: &Block, x: &Var, mask: &Var, seq: usize) -> Var {
+        let (nh, hd) = (self.cfg.n_heads, self.cfg.head_dim);
+        let scale = 1.0 / (hd as f32).sqrt();
+
+        // Project, then lay out as per-head [seq, head_dim] tensors.
+        //   [seq,dim] @ [dim,dim] = [seq,dim]
+        //   reshape   [seq, nh, hd]
+        //   rope      (kernel expects exactly [tokens, heads, head_dim])
+        //   transpose [nh, seq, hd]  → split into nh × [seq, hd]
+        let to_heads = |proj: Var, rope: bool| -> Vec<Var> {
+            let r = ops::reshape(&proj, &[seq, nh, hd]);
+            let r = if rope {
+                ops::rope(&r, self.cfg.rope_theta)
+            } else {
+                r
+            };
+            let t = ops::transpose_3d01(&r); // [nh, seq, hd]
+            ops::split_heads(&t)
+        };
+
+        let q = to_heads(ops::matmul(x, &b.wq), true);
+        let k = to_heads(ops::matmul(x, &b.wk), true);
+        let v = to_heads(ops::matmul(x, &b.wv), false);
+
+        // Per-head scaled-dot-product attention with causal mask.
+        let heads_out: Vec<Var> = (0..nh)
+            .map(|i| {
+                let kt = ops::transpose_2d(&k[i]); // [hd, seq]
+                let scores = ops::scale(&ops::matmul(&q[i], &kt), scale); // [seq,seq]
+                let scores = ops::add(&scores, mask); // causal
+                let probs = ops::softmax(&scores);
+                ops::matmul(&probs, &v[i]) // [seq, hd]
+            })
+            .collect();
+
+        // Stack heads back: nh × [seq,hd] → [nh,seq,hd] → [seq,nh,hd] → [seq,dim].
+        let merged = ops::merge_heads(&heads_out); // [nh, seq, hd]
+        let t = ops::transpose_3d01(&merged); // [seq, nh, hd]
+        let concat = ops::reshape(&t, &[seq, nh * hd]); // [seq, dim]
+        ops::matmul(&concat, &b.wo) // out projection
+    }
+
+    /// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[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 act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
+        ops::matmul(&act, &b.w_down) // [seq, dim]
+    }
+
+    /// Additive causal mask `[seq,seq]`: 0 on/below the diagonal, −1e9 above it
+    /// (so softmax zeros out future positions). A constant leaf (no grad needed,
+    /// but harmless if it accumulates one — it has no consumers downstream of x).
+    fn causal_mask(&self, seq: usize) -> Var {
+        let mut m = vec![0.0f32; seq * seq];
+        for i in 0..seq {
+            for j in (i + 1)..seq {
+                m[i * seq + j] = -1.0e9;
+            }
+        }
+        Var::leaf(Tensor::from_slice(&m, &[seq, seq]).to_device(self.device))
+    }
+}
+
+/// Materialise a parameter's value back to a host `Vec<f32>` (for the GD step
+/// and PyTorch parity export).
+pub fn param_to_host(v: &Var) -> Vec<f32> {
+    v.value().to_device(Device::Cpu).as_slice::<f32>().to_vec()
+}
+
+/// Build an I32 id tensor on `device` from token ids.
+pub fn ids_tensor(ids: &[i32], device: Device) -> Tensor {
+    Tensor::from_slice(ids, &[ids.len()]).to_device(device)
+}