diff --git a/Cargo.lock b/Cargo.lock index 0c5d7d5..48db26d 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 63d1380..456b6d7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ members = [ "crates/xtrain-cuda", "crates/xtrain-tensor", "crates/xtrain-autodiff", + "crates/xtrain-model", ] [workspace.package] diff --git a/crates/xtrain-model/Cargo.toml b/crates/xtrain-model/Cargo.toml new file mode 100644 index 0000000..98be4da --- /dev/null +++ b/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" } diff --git a/crates/xtrain-model/build.rs b/crates/xtrain-model/build.rs new file mode 100644 index 0000000..0c1d53a --- /dev/null +++ b/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() +} diff --git a/crates/xtrain-model/src/config.rs b/crates/xtrain-model/src/config.rs new file mode 100644 index 0000000..8d52109 --- /dev/null +++ b/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 + } +} diff --git a/crates/xtrain-model/src/lib.rs b/crates/xtrain-model/src/lib.rs new file mode 100644 index 0000000..61d5e5d --- /dev/null +++ b/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}; diff --git a/crates/xtrain-model/src/model.rs b/crates/xtrain-model/src/model.rs new file mode 100644 index 0000000..807d2c3 --- /dev/null +++ b/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, + 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) -> Self { + let leaf = |data: Vec, 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::(), "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 { + 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 { + 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 = (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` (for the GD step +/// and PyTorch parity export). +pub fn param_to_host(v: &Var) -> Vec { + v.value().to_device(Device::Cpu).as_slice::().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) +} diff --git a/crates/xtrain-model/tests/overfit.rs b/crates/xtrain-model/tests/overfit.rs new file mode 100644 index 0000000..5533bbe --- /dev/null +++ b/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 { + 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 = 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 = text.chars().map(stoi).collect(); + // Next-token prediction: input = tokens[..n-1], target = tokens[1..]. + let input: Vec = tokens[..tokens.len() - 1].to_vec(); + let target: Vec = 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::()[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(¶ms, 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::(); + 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"); +}