Gahow Wang
5353b38402
forward_batched(ids[B*S], batch)/loss_batched: run B equal-length sequences as ONE forward over flattened [B*S] ids, so every linear is one big [B*S,dim] GEMM. Attention reshapes to [B*nh,S,hd], runs the fused batched causal SDPA (per-seq mask + RoPE period=S, no cross-sequence attention), writes back [B*S,dim]. The old per-(batch,head) loop + host-round-tripping split/merge_heads + the additive causal_mask leaf are gone. forward(ids[seq]) is now forward_batched(ids,1), so the sampler / inference path (batch=1) is unchanged. +batched_ids_tensor helper. New batched.rs test: batched forward == looped single-sequence (logits identical 0.0, grads 6.4e-4, loss identical). PyTorch parity now exercises B>1 (B=2,S=4): loss 5e-8, logits 6.9e-6, all 25 param grads within rtol — verifying per-seq RoPE position + per-seq causal masking. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
249 lines
10 KiB
Rust
249 lines
10 KiB
Rust
//! 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]
|
||
q_norm: Var, // [head_dim] — per-head QK-norm (Qwen3-style)
|
||
k_norm: Var, // [head_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]
|
||
}
|
||
|
||
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]),
|
||
q_norm: mk(&[cfg.head_dim]),
|
||
k_norm: mk(&[cfg.head_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,
|
||
}
|
||
}
|
||
|
||
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.q_norm.clone(),
|
||
b.k_norm.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]`. This
|
||
/// is the batch-1 special case of [`forward_batched`](Self::forward_batched)
|
||
/// (used by the autoregressive sampler / inference path).
|
||
pub fn forward(&self, ids: &Tensor) -> Var {
|
||
self.forward_batched(ids, 1)
|
||
}
|
||
|
||
/// Batched forward over `batch` sequences of equal length `seq`, flattened to
|
||
/// `[batch*seq]` I32 ids in sequence-major order (sequence 0's `seq` tokens,
|
||
/// then sequence 1's, …). Returns logits `[batch*seq, vocab]` in the SAME flat
|
||
/// layout. The whole graph runs on the flattened tokens so every linear
|
||
/// projection is ONE big `[batch*seq, dim] × [dim, out]` GEMM (the
|
||
/// GPU-filling win); only attention is sequence-aware (per-sequence causal
|
||
/// mask + RoPE position, NO cross-sequence attention).
|
||
pub fn forward_batched(&self, ids: &Tensor, batch: usize) -> Var {
|
||
let total = ids.shape()[0];
|
||
assert_eq!(
|
||
total % batch,
|
||
0,
|
||
"ids len {total} not divisible by batch {batch}"
|
||
);
|
||
let seq = total / batch;
|
||
|
||
let mut h = ops::embedding(&self.embed, ids); // [batch*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, batch, 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) // [batch*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)
|
||
}
|
||
|
||
/// Batched cross-entropy mean loss: `forward_batched(ids, batch)` against
|
||
/// flat `targets` (`[batch*seq]` I32, same sequence-major layout). The CE mean
|
||
/// is over all `batch*seq` rows — identical to averaging the per-sequence
|
||
/// losses, so the loss value matches the looped single-sequence path.
|
||
pub fn loss_batched(&self, ids: &Tensor, targets: &Tensor, batch: usize) -> Var {
|
||
let logits = self.forward_batched(ids, batch);
|
||
ops::cross_entropy(&logits, targets)
|
||
}
|
||
|
||
/// Multi-head causal self-attention over a flattened batch. `x`:[batch*seq,dim]
|
||
/// (already normed), laid out sequence-major. The Q/K/V/O projections are big
|
||
/// `[batch*seq, dim]` GEMMs; the scaled-dot-product attention itself runs as a
|
||
/// fused BATCHED op over the `batch·n_heads` (sequence,head) blocks — each
|
||
/// attends within its own `[seq,seq]` causal window (NO cross-sequence
|
||
/// attention), with RoPE positions reset per sequence (`period = seq`). Causal
|
||
/// masking is applied inside the fused op's softmax kernel (no additive
|
||
/// `[seq,seq]` mask tensor).
|
||
fn attention(&self, b: &Block, x: &Var, batch: usize, seq: usize) -> Var {
|
||
let (nh, hd) = (self.cfg.n_heads, self.cfg.head_dim);
|
||
let total = batch * seq;
|
||
let bh = batch * nh;
|
||
let scale = 1.0 / (hd as f32).sqrt();
|
||
|
||
// Project, qk-norm + RoPE, then lay out as a batched [B*nh, seq, hd] tensor.
|
||
// [B*S,dim] @ [dim,dim] = [B*S,dim]
|
||
// reshape [B*S, nh, hd]
|
||
// qk-norm per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
|
||
// rope [B*S, nh, hd] with per-sequence position (period = seq)
|
||
// reshape [B, S, nh, hd] → transpose(1,2) → [B, nh, S, hd] → [B*nh, S, hd]
|
||
let to_bh = |proj: Var, norm: Option<&Var>| -> Var {
|
||
let r = ops::reshape(&proj, &[total, nh, hd]);
|
||
let r = match norm {
|
||
// Per-head RMSNorm: flatten the (B*S,nh) head rows, norm over hd,
|
||
// restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
|
||
Some(gamma) => {
|
||
let flat = ops::reshape(&r, &[total * nh, hd]);
|
||
let normed = ops::rms_norm(&flat, gamma, self.cfg.eps);
|
||
let r = ops::reshape(&normed, &[total, nh, hd]);
|
||
ops::rope(&r, self.cfg.rope_theta, seq)
|
||
}
|
||
None => r,
|
||
};
|
||
let r = ops::reshape(&r, &[batch, seq, nh, hd]);
|
||
let t = ops::transpose_4d12(&r); // [B, nh, S, hd]
|
||
ops::reshape(&t, &[bh, seq, hd]) // [B*nh, S, hd]
|
||
};
|
||
|
||
let q = to_bh(ops::matmul(x, &b.wq), Some(&b.q_norm));
|
||
let k = to_bh(ops::matmul(x, &b.wk), Some(&b.k_norm));
|
||
let v = to_bh(ops::matmul(x, &b.wv), None);
|
||
|
||
// Fused batched causal SDPA over all B*nh (sequence,head) blocks at once
|
||
// (2 batched GEMMs + 1 causal-softmax kernel; no per-head/per-seq loop).
|
||
let out = ops::attention(&q, &k, &v, scale); // [B*nh, S, hd]
|
||
|
||
// Back to [B*S, dim]: [B*nh,S,hd] → [B,nh,S,hd] → transpose(1,2) →
|
||
// [B,S,nh,hd] → [B*S, dim].
|
||
let out = ops::reshape(&out, &[batch, nh, seq, hd]);
|
||
let out = ops::transpose_4d12(&out); // [B, S, nh, hd]
|
||
let concat = ops::reshape(&out, &[total, nh * hd]); // [B*S, dim]
|
||
ops::matmul(&concat, &b.wo) // out projection
|
||
}
|
||
|
||
/// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[batch*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]
|
||
}
|
||
}
|
||
|
||
/// 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)
|
||
}
|
||
|
||
/// Flatten `batch` equal-length sequences into one `[batch*seq]` I32 tensor in
|
||
/// sequence-major order (the layout `forward_batched` expects). Each row of
|
||
/// `seqs` is one sequence; all must have the same length.
|
||
pub fn batched_ids_tensor(seqs: &[Vec<i32>], device: Device) -> Tensor {
|
||
assert!(!seqs.is_empty(), "empty batch");
|
||
let seq = seqs[0].len();
|
||
let mut flat = Vec::with_capacity(seqs.len() * seq);
|
||
for s in seqs {
|
||
assert_eq!(s.len(), seq, "ragged batch: sequences must be equal length");
|
||
flat.extend_from_slice(s);
|
||
}
|
||
Tensor::from_slice(&flat, &[flat.len()]).to_device(device)
|
||
}
|