model: batched forward [B,S]
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>
This commit is contained in:
@@ -24,4 +24,4 @@ pub use config::Config;
|
||||
#[cfg(not(no_cuda))]
|
||||
mod model;
|
||||
#[cfg(not(no_cuda))]
|
||||
pub use model::{TinyTransformer, ids_tensor, param_to_host};
|
||||
pub use model::{TinyTransformer, batched_ids_tensor, ids_tensor, param_to_host};
|
||||
|
||||
@@ -30,7 +30,6 @@ pub struct TinyTransformer {
|
||||
blocks: Vec<Block>,
|
||||
final_norm: Var, // [dim]
|
||||
lm_head: Var, // [dim, vocab]
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl TinyTransformer {
|
||||
@@ -72,7 +71,6 @@ impl TinyTransformer {
|
||||
blocks,
|
||||
final_norm,
|
||||
lm_head,
|
||||
device,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -106,16 +104,34 @@ impl TinyTransformer {
|
||||
}
|
||||
|
||||
/// Forward over a single sequence of token `ids` (`[seq]` I32 on this
|
||||
/// model's device). Returns the logits [`Var`] of shape `[seq, vocab]`.
|
||||
/// 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 {
|
||||
let seq = ids.shape()[0];
|
||||
let mask = self.causal_mask(seq);
|
||||
self.forward_batched(ids, 1)
|
||||
}
|
||||
|
||||
let mut h = ops::embedding(&self.embed, ids); // [seq, dim]
|
||||
/// 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, &mask, seq);
|
||||
let attn = self.attention(b, &normed, batch, seq);
|
||||
h = ops::add(&h, &attn);
|
||||
|
||||
// --- MLP sub-block (pre-norm + residual) ---
|
||||
@@ -125,7 +141,7 @@ impl TinyTransformer {
|
||||
}
|
||||
|
||||
let h = ops::rms_norm(&h, &self.final_norm, self.cfg.eps);
|
||||
ops::matmul(&h, &self.lm_head) // [seq, vocab]
|
||||
ops::matmul(&h, &self.lm_head) // [batch*seq, vocab]
|
||||
}
|
||||
|
||||
/// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
|
||||
@@ -134,76 +150,76 @@ impl TinyTransformer {
|
||||
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 {
|
||||
/// 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, then lay out as per-head [seq, head_dim] tensors.
|
||||
// [seq,dim] @ [dim,dim] = [seq,dim]
|
||||
// reshape [seq, nh, hd]
|
||||
// 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 (kernel expects exactly [tokens, heads, head_dim])
|
||||
// transpose [nh, seq, hd] → split into nh × [seq, hd]
|
||||
let to_heads = |proj: Var, norm: Option<&Var>| -> Vec<Var> {
|
||||
let r = ops::reshape(&proj, &[seq, nh, hd]);
|
||||
// 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 (seq,nh) head rows, norm over hd,
|
||||
// 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, &[seq * nh, hd]);
|
||||
let flat = ops::reshape(&r, &[total * nh, hd]);
|
||||
let normed = ops::rms_norm(&flat, gamma, self.cfg.eps);
|
||||
let r = ops::reshape(&normed, &[seq, nh, hd]);
|
||||
ops::rope(&r, self.cfg.rope_theta)
|
||||
let r = ops::reshape(&normed, &[total, nh, hd]);
|
||||
ops::rope(&r, self.cfg.rope_theta, seq)
|
||||
}
|
||||
None => r,
|
||||
};
|
||||
let t = ops::transpose_3d01(&r); // [nh, seq, hd]
|
||||
ops::split_heads(&t)
|
||||
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_heads(ops::matmul(x, &b.wq), Some(&b.q_norm));
|
||||
let k = to_heads(ops::matmul(x, &b.wk), Some(&b.k_norm));
|
||||
let v = to_heads(ops::matmul(x, &b.wv), None);
|
||||
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);
|
||||
|
||||
// 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();
|
||||
// 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]
|
||||
|
||||
// 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]
|
||||
// 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`:[seq,dim].
|
||||
/// 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]
|
||||
}
|
||||
|
||||
/// 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
|
||||
@@ -216,3 +232,17 @@ pub fn param_to_host(v: &Var) -> Vec<f32> {
|
||||
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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user