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>

Cargo.lock

@@ -103,6 +103,15 @@ dependencies = [
  "cc",
 ]
 
+[[package]]
+name = "xtrain-model"
+version = "0.1.0"
+dependencies = [
+ "xtrain-autodiff",
+ "xtrain-cuda",
+ "xtrain-tensor",
+]
+
 [[package]]
 name = "xtrain-tensor"
 version = "0.1.0"

Cargo.toml

@@ -4,6 +4,7 @@ members = [
     "crates/xtrain-cuda",
     "crates/xtrain-tensor",
     "crates/xtrain-autodiff",
+    "crates/xtrain-model",
 ]
 
 [workspace.package]

crates/xtrain-model/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "xtrain-model"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+xtrain-tensor = { path = "../xtrain-tensor" }
+xtrain-autodiff = { path = "../xtrain-autodiff" }
+
+[dev-dependencies]
+# Acceptance tests drive the GPU (device selection) directly.
+xtrain-cuda = { path = "../xtrain-cuda" }

crates/xtrain-model/build.rs

@@ -0,0 +1,26 @@
+use std::env;
+use std::path::Path;
+use std::process::Command;
+
+// Same per-crate convention as the other crates: this crate's tiny-transformer
+// forward/backward calls GPU ops (via xtrain-autodiff / xtrain-tensor), so it
+// gates GPU code + tests behind `not(no_cuda)`. cfg does not propagate across
+// crates, so each crate re-detects nvcc. No CUDA is compiled here.
+fn main() {
+    println!("cargo:rustc-check-cfg=cfg(no_cuda)");
+
+    let cuda_path = env::var("CUDA_HOME")
+        .or_else(|_| env::var("CUDA_PATH"))
+        .unwrap_or_else(|_| "/usr/local/cuda".to_string());
+
+    if !nvcc_available(&cuda_path) {
+        println!("cargo:rustc-cfg=no_cuda");
+    }
+}
+
+fn nvcc_available(cuda_path: &str) -> bool {
+    if Command::new("nvcc").arg("--version").output().is_ok() {
+        return true;
+    }
+    Path::new(&format!("{cuda_path}/bin/nvcc")).exists()
+}

crates/xtrain-model/src/config.rs

@@ -0,0 +1,54 @@
+//! Tiny-transformer hyperparameters. Host-only (no GPU), always compiled.
+
+/// Architecture config for [`crate::TinyTransformer`]. Keep it tiny — T5 is a
+/// correctness bring-up, not a real training run.
+#[derive(Debug, Clone, Copy)]
+pub struct Config {
+    /// Vocabulary size (char-level in the bring-up).
+    pub vocab: usize,
+    /// Model / residual width. Must equal `n_heads * head_dim`.
+    pub dim: usize,
+    /// Number of decoder blocks.
+    pub n_layers: usize,
+    /// Number of attention heads.
+    pub n_heads: usize,
+    /// Per-head dimension (`dim / n_heads`).
+    pub head_dim: usize,
+    /// SwiGLU hidden width (gate/up project to this, down projects back).
+    pub ffn_hidden: usize,
+    /// RMSNorm epsilon.
+    pub eps: f32,
+    /// RoPE base frequency (theta).
+    pub rope_theta: f32,
+}
+
+impl Config {
+    /// A minimal config used by the bring-up / overfit test.
+    pub fn tiny() -> Self {
+        let n_heads = 2;
+        let head_dim = 16;
+        Config {
+            vocab: 0, // set by the caller from the char vocab
+            dim: n_heads * head_dim,
+            n_layers: 2,
+            n_heads,
+            head_dim,
+            ffn_hidden: 64,
+            eps: 1e-5,
+            rope_theta: 10000.0,
+        }
+    }
+
+    /// Total learnable parameter count (for logging / sanity).
+    pub fn num_params(&self) -> usize {
+        let per_layer = 2 * self.dim                       // 2 rmsnorm gammas
+            + 3 * self.dim * self.dim                       // q/k/v proj
+            + self.dim * self.dim                           // out proj
+            + 2 * self.dim * self.ffn_hidden                // gate/up proj
+            + self.ffn_hidden * self.dim; // down proj
+        self.vocab * self.dim          // embedding
+            + self.n_layers * per_layer
+            + self.dim                 // final norm
+            + self.dim * self.vocab // lm head
+    }
+}

crates/xtrain-model/src/lib.rs

@@ -0,0 +1,26 @@
+//! Tiny modern-architecture transformer (Phase T5).
+//!
+//! A from-scratch decoder built entirely from the [`xtrain_autodiff`] op set:
+//! token embedding → `n_layers` × {pre-RMSNorm → multi-head causal attention
+//! (RoPE) → residual; pre-RMSNorm → SwiGLU MLP → residual} → final RMSNorm →
+//! LM-head matmul. The forward builds an autograd graph; calling `.backward()`
+//! on the cross-entropy loss fills every parameter's `.grad()`.
+//!
+//! Conventions (matching the engine, not HuggingFace):
+//! - Linear weights are `[in, out]` and applied as `x @ W` (no transpose), since
+//!   the engine's GEMM is plain `A @ B`.
+//! - `dim == n_heads * head_dim` (no separate attention projection size).
+//! - RoPE position = token row index (the kernel's built-in convention).
+//! - Causal masking is an additive `[seq,seq]` constant (−1e9 above the diagonal)
+//!   added to the attention scores before softmax.
+//!
+//! Everything GPU-facing is gated behind `not(no_cuda)`; on a GPU-less host the
+//! crate still `cargo check`s (only [`Config`] is visible there).
+
+mod config;
+pub use config::Config;
+
+#[cfg(not(no_cuda))]
+mod model;
+#[cfg(not(no_cuda))]
+pub use model::{TinyTransformer, ids_tensor, param_to_host};

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)
+}

crates/xtrain-model/tests/overfit.rs

@@ -0,0 +1,133 @@
+// End-to-end acceptance for the Phase T5 tiny transformer: overfit one fixed
+// char-level batch with a hand-written gradient-descent step and assert the loss
+// collapses toward 0. This is THE signal that the whole fwd+bwd graph (embedding,
+// RMSNorm, RoPE, multi-head attention, SwiGLU, LM head, cross-entropy) is wired
+// correctly — a single buggy backward would stall the loss.
+//
+// The optimizer here is deliberately minimal (`p ← p − lr·grad`); AdamW / LR
+// schedule / real data are T6. Gated behind `not(no_cuda)` (runs on dash5).
+#![cfg(not(no_cuda))]
+
+use xtrain_autodiff::tape::Var;
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, ids_tensor};
+use xtrain_tensor::Device;
+
+// Deterministic LCG fill in [-scale, scale).
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+fn require_gpu() {
+    assert!(
+        device::device_count().expect("device count") > 0,
+        "no CUDA device"
+    );
+    device::set_device(0).unwrap();
+}
+
+// One GD step over every parameter: p ← p − lr·grad, then zero the grad.
+fn gd_step(params: &[Var], lr: f32) {
+    for p in params {
+        if let Some(g) = p.grad() {
+            let updated = p.value().add(&g.scale(-lr));
+            p.set_value(updated);
+        }
+        p.zero_grad();
+    }
+}
+
+#[test]
+fn overfit_tiny_batch() {
+    require_gpu();
+    let device = Device::Cuda(0);
+
+    // --- Char-level bring-up: tiny embedded text → vocab → (input, target). ---
+    let text = "hello tiny transformer world";
+    let mut vocab_chars: Vec<char> = text.chars().collect();
+    vocab_chars.sort_unstable();
+    vocab_chars.dedup();
+    let vocab = vocab_chars.len();
+    let stoi = |c: char| vocab_chars.iter().position(|&x| x == c).unwrap() as i32;
+
+    let tokens: Vec<i32> = text.chars().map(stoi).collect();
+    // Next-token prediction: input = tokens[..n-1], target = tokens[1..].
+    let input: Vec<i32> = tokens[..tokens.len() - 1].to_vec();
+    let target: Vec<i32> = tokens[1..].to_vec();
+    let ids = ids_tensor(&input, device);
+    let targets = ids_tensor(&target, device);
+
+    // --- Tiny model with small-scale deterministic init. ---
+    let mut cfg = Config::tiny();
+    cfg.vocab = vocab;
+    let mut seed = 1u64;
+    let model = TinyTransformer::new(cfg, device, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        // RMSNorm gammas ([dim]) init to ~1; everything else small random.
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.08)
+        }
+    });
+    let params = model.params();
+    println!(
+        "overfit: vocab={vocab} seq={} params={}",
+        input.len(),
+        cfg.num_params()
+    );
+
+    let read_loss = |l: &Var| -> f32 { l.value().to_device(Device::Cpu).as_slice::<f32>()[0] };
+
+    let lr = 0.3f32;
+    let steps = 200;
+    let start = read_loss(&model.loss(&ids, &targets));
+    let mut last = start;
+    for step in 0..steps {
+        let loss = model.loss(&ids, &targets);
+        last = read_loss(&loss);
+        if step % 20 == 0 || step == steps - 1 {
+            println!("step {step:3}: loss = {last:.6}");
+        }
+        loss.backward();
+        gd_step(&params, lr);
+    }
+
+    println!("overfit: start loss = {start:.6} → final loss = {last:.6} ({steps} steps)");
+    // A correct fwd+bwd memorises this tiny fixed batch: loss → ~0.
+    assert!(
+        last < 0.05,
+        "overfit failed to drive loss to ~0: start {start:.4} final {last:.4}"
+    );
+    assert!(last < start, "loss did not decrease");
+
+    // Sanity: greedy argmax should reproduce the target sequence after overfit.
+    let logits = model.forward(&ids).value().to_device(Device::Cpu);
+    let lg = logits.as_slice::<f32>();
+    let mut correct = 0;
+    for (r, &t) in target.iter().enumerate() {
+        let row = &lg[r * vocab..(r + 1) * vocab];
+        let argmax = row
+            .iter()
+            .enumerate()
+            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
+            .unwrap()
+            .0 as i32;
+        if argmax == t {
+            correct += 1;
+        }
+    }
+    println!("overfit: greedy match {correct}/{}", target.len());
+    assert_eq!(correct, target.len() as i32, "did not memorise the batch");
+}