diff --git a/crates/xtrain-autodiff/src/ops.rs b/crates/xtrain-autodiff/src/ops.rs index 7e944df..fe8bb6c 100644 --- a/crates/xtrain-autodiff/src/ops.rs +++ b/crates/xtrain-autodiff/src/ops.rs @@ -120,15 +120,17 @@ pub fn swiglu(gate: &Var, up: &Var) -> Var { mul(&silu(gate), up) } -/// RoPE (rotate_half) over `x:[tokens,heads,head_dim]`. Orthogonal map, so the -/// backward is the inverse rotation of `dy` — no cached forward values needed. -pub fn rope(x: &Var, theta: f32) -> Var { - let out = x.value().rope(theta); +/// RoPE (rotate_half) over `x:[tokens,heads,head_dim]` with per-sequence position +/// `row % period` (`period` = sequence length; `period == tokens` for a single +/// sequence). Orthogonal map, so the backward is the inverse rotation of `dy` — no +/// cached forward values needed. +pub fn rope(x: &Var, theta: f32, period: usize) -> Var { + let out = x.value().rope(theta, period); Var::from_op( out, vec![x.clone()], Box::new(move |dy, parents| { - Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta)); + Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta, period)); }), ) } diff --git a/crates/xtrain-autodiff/tests/autograd.rs b/crates/xtrain-autodiff/tests/autograd.rs index d9e44e9..cc27cf9 100644 --- a/crates/xtrain-autodiff/tests/autograd.rs +++ b/crates/xtrain-autodiff/tests/autograd.rs @@ -327,12 +327,12 @@ fn rope_bwd() { let w = fill(n, 82); let x = Var::leaf(cuda(&x_h, &[tokens, heads, head_dim])); - let out = ops::rope(&x, theta); + let out = ops::rope(&x, theta, tokens); scalar_loss(&out, &w).backward(); let dx = x.grad().unwrap().to_device(Device::Cpu); let wf = w.clone(); - let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).rope(theta), &wf); + let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).rope(theta, tokens), &wf); report( "rope dX", &grad_check( @@ -345,6 +345,38 @@ fn rope_bwd() { ); } +// ---- rope batched (per-sequence position = row % period) ---- +// tokens = B*S laid end to end; period = S. Sequences 2 and 3 re-use positions +// 0..S, so the kernel's `tok % period` must reset RoPE per sequence. +#[test] +fn rope_batched_bwd() { + require_gpu(); + let (b, s, heads, head_dim) = (3, 4, 2, 8); + let tokens = b * s; + let n = tokens * heads * head_dim; + let theta = 10000.0; + let x_h = fill(n, 83); + let w = fill(n, 84); + + let x = Var::leaf(cuda(&x_h, &[tokens, heads, head_dim])); + let out = ops::rope(&x, theta, s); + scalar_loss(&out, &w).backward(); + + let dx = x.grad().unwrap().to_device(Device::Cpu); + let wf = w.clone(); + let lx = move |v: &[f32], sh: &[usize]| weighted_sum(&cuda(v, sh).rope(theta, s), &wf); + report( + "rope batched dX", + &grad_check( + &x_h, + &[tokens, heads, head_dim], + &lx, + dx.as_slice::(), + cfg_linear(), + ), + ); +} + // ---- softmax ---- #[test] fn softmax_bwd() { diff --git a/crates/xtrain-cuda/src/ffi.rs b/crates/xtrain-cuda/src/ffi.rs index e18e027..5ccc9f0 100644 --- a/crates/xtrain-cuda/src/ffi.rs +++ b/crates/xtrain-cuda/src/ffi.rs @@ -125,7 +125,9 @@ unsafe extern "C" { pub fn launch_silu_f32(x: *const f32, y: *mut f32, n: i32, s: CudaStream); pub fn launch_silu_dx_f32(x: *const f32, dy: *const f32, dx: *mut f32, n: i32, s: CudaStream); - // RoPE (rotate_half), x:[tokens,heads,head_dim], position = token index. + // RoPE (rotate_half), x:[tokens,heads,head_dim], position = (token index % + // period). `period` = sequence length, so a flattened batch of sequences gets + // per-sequence positions; period == tokens reproduces the single-sequence case. pub fn launch_rope_f32( x: *const f32, y: *mut f32, @@ -133,6 +135,7 @@ unsafe extern "C" { heads: i32, head_dim: i32, theta: f32, + period: i32, s: CudaStream, ); pub fn launch_rope_dx_f32( @@ -142,6 +145,7 @@ unsafe extern "C" { heads: i32, head_dim: i32, theta: f32, + period: i32, s: CudaStream, ); diff --git a/crates/xtrain-distributed/src/ddp.rs b/crates/xtrain-distributed/src/ddp.rs index dd217d0..33043be 100644 --- a/crates/xtrain-distributed/src/ddp.rs +++ b/crates/xtrain-distributed/src/ddp.rs @@ -17,7 +17,7 @@ use std::thread; use std::time::Instant; use xtrain_autodiff::tape::Var; -use xtrain_model::{Config, TinyTransformer, ids_tensor}; +use xtrain_model::{Config, TinyTransformer, batched_ids_tensor}; use xtrain_optim::GpuAdamW; use xtrain_tensor::Device; use xtrain_train::checkpoint; @@ -91,10 +91,11 @@ pub fn train_rank( let mut losses = Vec::with_capacity(cfg.steps); let mut evals = Vec::new(); let mut best_val: Option = None; - // Each rank reaches the global batch mean as (Σ_global / world) · (1/b_local), - // where b_local = batch_size / world (see DdpContext::all_reduce_average_grads). + // Each rank runs ONE batched forward over its b_local = batch_size/world + // sequences → backward grad = local mean (Σ_local / b_local). all_reduce_average + // (sum across ranks, /world) then gives Σ_global/(world·b_local) = Σ_global/ + // B_global — already the global-batch mean — so the clip pre-scale is 1.0. let batch_local = cfg.batch_size / ctx.world; - let inv_batch_local = 1.0 / batch_local as f32; let start = Instant::now(); let mut tokens_seen: u64 = 0; // Rank 0 owns the held-out eval + best-val checkpoint (params are identical @@ -105,31 +106,36 @@ pub fn train_rank( let lr = cfg.schedule.lr(step); // Draw the whole global batch from the shared RNG (same on every rank); - // run forward+backward only on this rank's shard. The tape SUMs the - // shard's grads; the union of shards == the single-GPU batch. - let mut local_loss_sum = 0.0f32; + // collect only this rank's shard (global index % world == rank) and run it + // as ONE batched forward/backward. The union of shards == the single-GPU + // batch; each rank's backward yields its local mean (Σ_local / b_local). + let mut inputs = Vec::with_capacity(batch_local); + let mut targets_v = Vec::with_capacity(batch_local); for i in 0..cfg.batch_size { let (input, target) = corpus.sample(cfg.seq_len, &mut rng); - if i % ctx.world != ctx.rank { - continue; // not this rank's sequence + if i % ctx.world == ctx.rank { + inputs.push(input); + targets_v.push(target); } - let ids = ids_tensor(&input, device); - let targets = ids_tensor(&target, device); - let loss = model.loss(&ids, &targets); - local_loss_sum += read_scalar(&loss); - loss.backward(); - tokens_seen += cfg.seq_len as u64; } + let ids = batched_ids_tensor(&inputs, device); + let targets = batched_ids_tensor(&targets_v, device); + let loss = model.loss_batched(&ids, &targets, batch_local); + let local_mean = read_scalar(&loss); // Σ_local / b_local + loss.backward(); + tokens_seen += (batch_local * cfg.seq_len) as u64; - // AllReduce(sum) + /world the grads → every rank holds Σ_global/world. + // AllReduce(sum) + /world the grads → every rank holds Σ_global/B_global + // (local means summed over ranks, /world = global mean). See note above. ctx.all_reduce_average_grads(¶ms); - // The reported loss is the global mean: average local sums across ranks. - let step_loss = all_reduce_loss(ctx, local_loss_sum) / cfg.batch_size as f32; + // Reported loss = global mean: sum the per-rank local sums (= mean·b_local) + // across ranks, /B_global. With equal b_local this is mean over ranks. + let step_loss = + all_reduce_loss(ctx, local_mean * batch_local as f32) / cfg.batch_size as f32; losses.push(step_loss); - // clip pre_scale = 1/b_local finishes the average to Σ_global/B_global, - // identical to the single-GPU clip(pre_scale = 1/B_global). - let gnorm = clip_grad_norm_gpu(¶ms, cfg.max_grad_norm, inv_batch_local); + // Grads are already the global-batch mean — just clip (pre-scale 1.0). + let gnorm = clip_grad_norm_gpu(¶ms, cfg.max_grad_norm, 1.0); opt.step(lr, ¶ms); for p in ¶ms { p.zero_grad(); diff --git a/crates/xtrain-distributed/tests/ddp_correctness.rs b/crates/xtrain-distributed/tests/ddp_correctness.rs index 371ad3c..faa8f5e 100644 --- a/crates/xtrain-distributed/tests/ddp_correctness.rs +++ b/crates/xtrain-distributed/tests/ddp_correctness.rs @@ -13,7 +13,7 @@ use std::time::Instant; use xtrain_cuda::device; use xtrain_distributed::{DdpConfig, DdpContext, build_model, get_unique_id, launch, train_rank}; -use xtrain_model::{Config, ids_tensor}; +use xtrain_model::{Config, batched_ids_tensor}; use xtrain_optim::GpuAdamW; use xtrain_tensor::Device; use xtrain_train::clip::clip_grad_norm_gpu; @@ -47,22 +47,25 @@ fn run_single_gpu(cfg: Config, corpus: &Corpus, dcfg: &DdpConfig) -> (Vec, let params = model.params(); let mut opt = GpuAdamW::new(dcfg.weight_decay); let mut rng = dcfg.seed; - let inv_batch = 1.0 / dcfg.batch_size as f32; let mut losses = Vec::new(); for step in 0..dcfg.steps { let lr = dcfg.schedule.lr(step); - let mut loss_sum = 0.0f32; + // Sample the whole global batch and run it as ONE batched forward/backward + // (matches the T10 DDP path: backward yields the global-batch mean grad). + let mut inputs = Vec::with_capacity(dcfg.batch_size); + let mut targets_v = Vec::with_capacity(dcfg.batch_size); for _ in 0..dcfg.batch_size { let (input, target) = corpus.sample(dcfg.seq_len, &mut rng); - let ids = ids_tensor(&input, device); - let targets = ids_tensor(&target, device); - let loss = model.loss(&ids, &targets); - loss_sum += loss.value().to_device(Device::Cpu).as_slice::()[0]; - loss.backward(); + inputs.push(input); + targets_v.push(target); } - losses.push(loss_sum * inv_batch); - clip_grad_norm_gpu(¶ms, dcfg.max_grad_norm, inv_batch); + let ids = batched_ids_tensor(&inputs, device); + let targets = batched_ids_tensor(&targets_v, device); + let loss = model.loss_batched(&ids, &targets, dcfg.batch_size); + losses.push(loss.value().to_device(Device::Cpu).as_slice::()[0]); + loss.backward(); + clip_grad_norm_gpu(¶ms, dcfg.max_grad_norm, 1.0); opt.step(lr, ¶ms); for p in ¶ms { p.zero_grad(); diff --git a/crates/xtrain-model/src/lib.rs b/crates/xtrain-model/src/lib.rs index 43ba6a1..6b4c651 100644 --- a/crates/xtrain-model/src/lib.rs +++ b/crates/xtrain-model/src/lib.rs @@ -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}; diff --git a/crates/xtrain-model/src/model.rs b/crates/xtrain-model/src/model.rs index 25cd7b3..ad8e33c 100644 --- a/crates/xtrain-model/src/model.rs +++ b/crates/xtrain-model/src/model.rs @@ -106,16 +106,35 @@ 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]; + 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 mask = self.causal_mask(seq); - let mut h = ops::embedding(&self.embed, ids); // [seq, dim] + 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, &mask, batch, seq); h = ops::add(&h, &attn); // --- MLP sub-block (pre-norm + residual) --- @@ -125,7 +144,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,57 +153,93 @@ 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; only the scaled-dot-product attention is + /// sequence-aware — each (head, sequence) attends within its own `[seq,seq]` + /// causal window, with NO cross-sequence attention, and RoPE positions reset + /// per sequence (`period = seq`). + fn attention(&self, b: &Block, x: &Var, mask: &Var, batch: usize, seq: usize) -> Var { let (nh, hd) = (self.cfg.n_heads, self.cfg.head_dim); + let total = batch * seq; 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, then lay out as per-head [batch*seq, head_dim] tensors. + // [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] + // rope [B*S, nh, hd] with per-sequence position (period = seq) + // transpose [nh, B*S, hd] → split into nh × [B*S, hd] let to_heads = |proj: Var, norm: Option<&Var>| -> Vec { - let r = ops::reshape(&proj, &[seq, nh, hd]); + 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 t = ops::transpose_3d01(&r); // [nh, B*S, hd] + ops::split_heads(&t) // nh × [B*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); - // Per-head scaled-dot-product attention with causal mask. + // Per-head SDPA. Within a head, attention must NOT cross sequences, so we + // split the [B*S, hd] head into its B per-sequence [seq, hd] blocks + // (reshape to [B,seq,hd] + split over the batch axis), run masked SDPA on + // each, then reassemble. The big GEMMs above stay batched; this inner + // attention is the only sequence-aware (looped) part. Backward falls out of + // the matmul/softmax/scale/split/merge nodes — no new gradient to verify. + let attend_seq = |qh: &Var, kh: &Var, vh: &Var| -> Var { + let kt = ops::transpose_2d(kh); // [hd, seq] + let scores = ops::scale(&ops::matmul(qh, &kt), scale); // [seq,seq] + let scores = ops::add(&scores, mask); // causal + let probs = ops::softmax(&scores); + ops::matmul(&probs, vh) // [seq, hd] + }; + 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] + if batch == 1 { + return attend_seq(&q[i], &k[i], &v[i]); // [seq, hd] + } + // [B*S, hd] → [B, seq, hd] → B × [seq, hd] + let qb = ops::split_heads(&ops::reshape(&q[i], &[batch, seq, hd])); + let kb = ops::split_heads(&ops::reshape(&k[i], &[batch, seq, hd])); + let vb = ops::split_heads(&ops::reshape(&v[i], &[batch, seq, hd])); + let per_seq: Vec = (0..batch) + .map(|s| attend_seq(&qb[s], &kb[s], &vb[s])) + .collect(); + // B × [seq, hd] → [B, seq, hd] → [B*S, hd] + let merged = ops::merge_heads(&per_seq); + ops::reshape(&merged, &[total, 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] + // Stack heads back: nh × [B*S,hd] → [nh,B*S,hd] → [B*S,nh,hd] → [B*S,dim]. + let merged = ops::merge_heads(&heads_out); // [nh, B*S, hd] + let t = ops::transpose_3d01(&merged); // [B*S, nh, hd] + let concat = ops::reshape(&t, &[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] @@ -216,3 +271,17 @@ pub fn param_to_host(v: &Var) -> Vec { 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], 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) +} diff --git a/crates/xtrain-model/tests/batched.rs b/crates/xtrain-model/tests/batched.rs new file mode 100644 index 0000000..093c77f --- /dev/null +++ b/crates/xtrain-model/tests/batched.rs @@ -0,0 +1,142 @@ +// T10 batched-forward equivalence: a batched forward over B sequences must equal +// the old single-sequence path (run each sequence on its own, concatenate the +// logits) — both for the forward logits AND every parameter's gradient. +// +// This is THE on-GPU correctness gate for batching (no PyTorch needed): if the +// per-sequence RoPE position, per-sequence causal masking, or any flattened op +// were wrong, the batched logits/grads would drift from the looped reference. +// +// Forward equivalence: batched logits[b*S+i] == single-seq-b logits[i]. +// Gradient equivalence: the batched loss is the mean over all B*S rows, i.e. +// (1/B)·Σ_b mean_i(loss_b); summing the B single-sequence losses and scaling by +// 1/B gives the SAME scalar, so their summed grads (tape fan-out) ×1/B match the +// batched grads. We check that. +#![cfg(not(no_cuda))] + +use xtrain_cuda::device; +use xtrain_model::{Config, TinyTransformer, batched_ids_tensor, ids_tensor}; +use xtrain_tensor::Device; + +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 build(cfg: Config, device: Device) -> TinyTransformer { + let mut seed = 1u64; + TinyTransformer::new(cfg, device, |shape| { + seed = seed.wrapping_add(1); + let n: usize = shape.iter().product(); + if shape.len() == 1 { + fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect() + } else { + fill(n, seed, 0.08) + } + }) +} + +fn host(t: &xtrain_tensor::Tensor) -> Vec { + t.to_device(Device::Cpu).as_slice::().to_vec() +} + +#[test] +fn batched_matches_looped_single_sequence() { + assert!(device::device_count().unwrap() > 0, "no CUDA device"); + device::set_device(0).unwrap(); + let device = Device::Cuda(0); + + let mut cfg = Config::tiny(); + cfg.vocab = 16; + let batch = 3usize; + let seq = 5usize; + // B distinct sequences (sequence-major), within vocab. + let seqs: Vec> = (0..batch) + .map(|b| { + (0..seq) + .map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32) + .collect() + }) + .collect(); + let tgts: Vec> = (0..batch) + .map(|b| { + (0..seq) + .map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32) + .collect() + }) + .collect(); + + // --- Batched forward: ONE pass over [B*S]. --- + let bmodel = build(cfg, device); + let bids = batched_ids_tensor(&seqs, device); + let blogits = host(&bmodel.forward_batched(&bids, batch).value()); + + // --- Looped reference: each sequence on its own, concatenate logits. --- + let smodel = build(cfg, device); + let mut slogits = Vec::with_capacity(batch * seq * cfg.vocab); + for s in &seqs { + let ids = ids_tensor(s, device); + slogits.extend(host(&smodel.forward(&ids).value())); + } + + // Forward equivalence (fp GEMM rounding only differs in summation order). + let max_rel = blogits + .iter() + .zip(&slogits) + .map(|(b, s)| (b - s).abs() / s.abs().max(1e-4)) + .fold(0.0f32, f32::max); + println!("batched vs looped: logits max rel err = {max_rel:.3e}"); + assert!(max_rel < 1e-3, "batched logits diverged: {max_rel:.3e}"); + + // --- Gradient equivalence. --- + // Batched: loss = mean over B*S rows; one backward. + let bparams = bmodel.params(); + let btgt = batched_ids_tensor(&tgts, device); + let bloss = bmodel.loss_batched(&bids, &btgt, batch); + let bloss_val = host(&bloss.value())[0]; + bloss.backward(); + + // Looped: Σ_b loss_b (each a per-sequence mean), then grad ×(1/B) == batched. + let sparams = smodel.params(); + let mut sloss_sum = 0.0f32; + for (s, t) in seqs.iter().zip(&tgts) { + let ids = ids_tensor(s, device); + let tg = ids_tensor(t, device); + let l = smodel.loss(&ids, &tg); + sloss_sum += host(&l.value())[0]; + l.backward(); + } + println!( + "batched loss = {bloss_val:.6} looped mean = {:.6}", + sloss_sum / batch as f32 + ); + assert!( + (bloss_val - sloss_sum / batch as f32).abs() < 1e-4, + "batched loss != looped mean" + ); + + let mut max_grad_rel = 0.0f32; + for (bp, sp) in bparams.iter().zip(&sparams) { + let bg = host(&bp.grad().expect("batched grad")); + let sg = host(&sp.grad().expect("looped grad")); + for (g_b, g_s) in bg.iter().zip(&sg) { + // looped grad is the SUM over B sequences; ×(1/B) recovers the mean. + let g_s = g_s / batch as f32; + let rel = (g_b - g_s).abs() / g_s.abs().max(1e-4); + max_grad_rel = max_grad_rel.max(rel); + } + } + println!("batched vs looped: grad max rel err = {max_grad_rel:.3e}"); + assert!( + max_grad_rel < 5e-3, + "batched grads diverged: {max_grad_rel:.3e}" + ); +} diff --git a/crates/xtrain-model/tests/parity.py b/crates/xtrain-model/tests/parity.py index 249388d..3f4e310 100644 --- a/crates/xtrain-model/tests/parity.py +++ b/crates/xtrain-model/tests/parity.py @@ -55,10 +55,13 @@ NH = int(cfg["n_heads"]) HD = int(cfg["head_dim"]) EPS = float(cfg["eps"]) THETA = float(cfg["rope_theta"]) +# Batched: B sequences of length SEQ, flattened sequence-major to [B*SEQ] ids. +B = int(cfg.get("batch", "1")) +SEQ = int(cfg["seq"]) ids = read_ids("ids.txt") targets = read_ids("targets.txt") -SEQ = len(ids) +assert len(ids) == B * SEQ, f"ids {len(ids)} != B*SEQ {B*SEQ}" # Load params as leaf tensors requiring grad (float64 for a clean reference). P = {} @@ -76,15 +79,16 @@ def rms_norm(x, gamma): return x * torch.rsqrt(ms + EPS) * gamma -def rope(x): # x: [seq, nh, hd], position = token index, matching the kernel +def rope(x): # x: [B*SEQ, nh, hd], position = (row % SEQ) — resets per sequence half = HD // 2 out = torch.empty_like(x) i = torch.arange(half, dtype=torch.float64) freq = THETA ** (-(2.0 * i) / HD) # [half] - pos = torch.arange(SEQ, dtype=torch.float64).reshape(SEQ, 1) # [seq,1] - ang = pos * freq # [seq, half] - c = torch.cos(ang).reshape(SEQ, 1, half) - s = torch.sin(ang).reshape(SEQ, 1, half) + # Position within each sequence: rows 0..SEQ for seq 0, 0..SEQ for seq 1, ... + pos = (torch.arange(B * SEQ, dtype=torch.float64) % SEQ).reshape(B * SEQ, 1) + ang = pos * freq # [B*SEQ, half] + c = torch.cos(ang).reshape(B * SEQ, 1, half) + s = torch.sin(ang).reshape(B * SEQ, 1, half) x0 = x[..., :half] x1 = x[..., half:] out[..., :half] = x0 * c - x1 * s @@ -102,26 +106,30 @@ for l in range(NL): "ffn_norm", "w_gate", "w_up", "w_down"]}) idx = torch.tensor(ids, dtype=torch.long) +# Per-sequence causal mask (broadcast over the batch); NO cross-sequence attention. mask = torch.triu(torch.full((SEQ, SEQ), -1.0e9, dtype=torch.float64), diagonal=1) -h = emb[idx] # [seq, dim] +h = emb[idx] # [B*SEQ, dim] (everything stays flattened, matching the Rust path) for L in layers: # Attention x = rms_norm(h, L["attn_norm"]) - q = (x @ L["wq"]).reshape(SEQ, NH, HD) - k = (x @ L["wk"]).reshape(SEQ, NH, HD) - v = (x @ L["wv"]).reshape(SEQ, NH, HD) + q = (x @ L["wq"]).reshape(B * SEQ, NH, HD) + k = (x @ L["wk"]).reshape(B * SEQ, NH, HD) + v = (x @ L["wv"]).reshape(B * SEQ, NH, HD) # Per-head QK-norm (Qwen3-style), before RoPE. q = rms_norm(q, L["q_norm"]) k = rms_norm(k, L["k_norm"]) - q = rope(q).transpose(0, 1) # [nh, seq, hd] - k = rope(k).transpose(0, 1) - v = v.transpose(0, 1) + q = rope(q) # [B*SEQ, nh, hd] + k = rope(k) + # Reshape to [B, NH, SEQ, HD] so attention runs within each sequence. + q = q.reshape(B, SEQ, NH, HD).transpose(1, 2) # [B, nh, seq, hd] + k = k.reshape(B, SEQ, NH, HD).transpose(1, 2) + v = v.reshape(B, SEQ, NH, HD).transpose(1, 2) scale = 1.0 / math.sqrt(HD) - scores = (q @ k.transpose(-1, -2)) * scale + mask # [nh, seq, seq] + scores = (q @ k.transpose(-1, -2)) * scale + mask # [B, nh, seq, seq] probs = torch.softmax(scores, dim=-1) - out = probs @ v # [nh, seq, hd] - out = out.transpose(0, 1).reshape(SEQ, DIM) # [seq, dim] + out = probs @ v # [B, nh, seq, hd] + out = out.transpose(1, 2).reshape(B * SEQ, DIM) # [B*SEQ, dim] attn = out @ L["wo"] h = h + attn # MLP @@ -133,7 +141,7 @@ for L in layers: h = h + mlp h = rms_norm(h, final_norm) -logits = h @ lm_head # [seq, vocab] +logits = h @ lm_head # [B*SEQ, vocab] loss = torch.nn.functional.cross_entropy( logits, torch.tensor(targets, dtype=torch.long), reduction="mean") diff --git a/crates/xtrain-model/tests/parity_dump.rs b/crates/xtrain-model/tests/parity_dump.rs index 3181148..a64ef33 100644 --- a/crates/xtrain-model/tests/parity_dump.rs +++ b/crates/xtrain-model/tests/parity_dump.rs @@ -53,12 +53,17 @@ fn dump_for_parity() { ); fs::create_dir_all(&dir).unwrap(); - // Fixed config + ids (independent of any text, for reproducibility). + // Fixed config + ids (independent of any text, for reproducibility). B>1 so + // the batched forward is exercised: 2 sequences of length 4, flattened + // sequence-major to [B*S]=8 ids. Per-sequence RoPE position (resets at the + // sequence boundary) + per-sequence causal masking (no cross-sequence + // attention) are both checked against PyTorch. let mut cfg = Config::tiny(); cfg.vocab = 12; - let ids: Vec = vec![3, 1, 4, 1, 5, 9, 2, 6]; + let batch = 2usize; + let seq = 4usize; + let ids: Vec = vec![3, 1, 4, 1, 5, 9, 2, 6]; // [B*S], sequence-major let targets: Vec = vec![1, 4, 1, 5, 9, 2, 6, 0]; - let seq = ids.len(); // Same deterministic init as the overfit test. let mut seed = 1u64; @@ -83,6 +88,7 @@ fn dump_for_parity() { writeln!(f, "ffn_hidden {}", cfg.ffn_hidden).unwrap(); writeln!(f, "eps {:e}", cfg.eps).unwrap(); writeln!(f, "rope_theta {:e}", cfg.rope_theta).unwrap(); + writeln!(f, "batch {batch}").unwrap(); writeln!(f, "seq {seq}").unwrap(); } { @@ -105,10 +111,11 @@ fn dump_for_parity() { write_vec(&dir, &format!("w_{name}.txt"), ¶m_to_host(p), &shape); } - // Forward logits + loss, then backward → per-param grads. + // Batched forward logits + loss (B sequences as one forward), then backward + // → per-param grads. let ids_t = ids_tensor(&ids, device); let targets_t = ids_tensor(&targets, device); - let logits = model.forward(&ids_t); + let logits = model.forward_batched(&ids_t, batch); write_vec( &dir, "logits.txt", @@ -116,7 +123,7 @@ fn dump_for_parity() { logits.value().shape(), ); - let loss = model.loss(&ids_t, &targets_t); + let loss = model.loss_batched(&ids_t, &targets_t, batch); let loss_val = param_to_host(&loss)[0]; { let mut f = fs::File::create(dir.join("loss.txt")).unwrap(); diff --git a/crates/xtrain-tensor/src/tensor.rs b/crates/xtrain-tensor/src/tensor.rs index 5ad55f5..60d2e31 100644 --- a/crates/xtrain-tensor/src/tensor.rs +++ b/crates/xtrain-tensor/src/tensor.rs @@ -454,13 +454,20 @@ impl Tensor { dx } - /// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; the position - /// of each token is its row index. Returns the rotated tensor. + /// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; each token's + /// position is `row % period`. `period` = sequence length, so a flattened + /// batch `[B*S,heads,head_dim]` gets per-sequence positions (pass `period=S`); + /// pass `period=tokens` for a single sequence (position = row). Returns the + /// rotated tensor. #[cfg(not(no_cuda))] - pub fn rope(&self, theta: f32) -> Self { + pub fn rope(&self, theta: f32, period: usize) -> Self { assert_eq!(self.ndim(), 3, "rope requires [tokens,heads,head_dim]"); let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]); assert_eq!(head_dim % 2, 0, "head_dim must be even"); + assert!( + period > 0 && tokens % period == 0, + "tokens must be a multiple of period" + ); let out = Tensor::zeros(&self.shape, DType::F32, self.device()); unsafe { xtrain_cuda::ffi::launch_rope_f32( @@ -470,6 +477,7 @@ impl Tensor { heads as i32, head_dim as i32, theta, + period as i32, std::ptr::null_mut(), ); } @@ -477,9 +485,9 @@ impl Tensor { } /// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an - /// orthogonal map, so it needs no cached forward values, only `theta`. + /// orthogonal map, so it needs no cached forward values, only `theta`/`period`. #[cfg(not(no_cuda))] - pub fn rope_backward(dy: &Tensor, theta: f32) -> Self { + pub fn rope_backward(dy: &Tensor, theta: f32, period: usize) -> Self { let (tokens, heads, head_dim) = (dy.shape[0], dy.shape[1], dy.shape[2]); let dx = Tensor::zeros(&dy.shape, DType::F32, dy.device()); unsafe { @@ -490,6 +498,7 @@ impl Tensor { heads as i32, head_dim as i32, theta, + period as i32, std::ptr::null_mut(), ); } diff --git a/crates/xtrain-train/src/train_loop.rs b/crates/xtrain-train/src/train_loop.rs index 0185e5a..4d93b7f 100644 --- a/crates/xtrain-train/src/train_loop.rs +++ b/crates/xtrain-train/src/train_loop.rs @@ -1,18 +1,20 @@ -//! The training loop: sample sequences → forward `loss` → backward → grad clip -//! (with batch averaging) → AdamW step → zero grads; with an LR schedule, -//! periodic loss logging, and periodic checkpointing. +//! The training loop: sample a batch of sequences → ONE batched forward `loss` → +//! backward → grad clip → AdamW step → zero grads; with an LR schedule, periodic +//! loss logging, and periodic checkpointing. //! -//! The T5 model is single-sequence, so a "batch" of `batch_size` sequences is -//! handled by running forward+backward on each and letting the tape SUM their -//! grads (its fan-out rule); the clip pass then multiplies by `1/batch_size` to -//! recover the batch-mean gradient before clipping + the optimizer step. +//! Since T10 the model is batched (`loss_batched`): `batch_size` sequences are +//! flattened to `[batch*seq]` and run as a SINGLE forward/backward, so the linear +//! projections become big `[batch*seq, dim]` GEMMs that fill the GPU. The +//! cross-entropy mean is over all `batch*seq` rows — already the batch-mean loss, +//! so backward yields the batch-mean gradient directly (clip pre-scale = 1.0; no +//! more "loop B times + SUM + ×1/batch" hack). #![cfg(not(no_cuda))] use std::path::PathBuf; use std::time::Instant; -use xtrain_model::{TinyTransformer, ids_tensor}; +use xtrain_model::{TinyTransformer, batched_ids_tensor, ids_tensor}; use xtrain_optim::GpuAdamW; use xtrain_tensor::Device; @@ -67,7 +69,6 @@ pub fn train( let mut losses = Vec::with_capacity(cfg.steps); let mut evals = Vec::new(); let mut best_val: Option = None; - let inv_batch = 1.0 / cfg.batch_size as f32; let start = Instant::now(); let mut tokens_seen: u64 = 0; // Best-val checkpointing only kicks in when we actually evaluate. @@ -76,22 +77,26 @@ pub fn train( for step in 0..cfg.steps { let lr = cfg.schedule.lr(step); - // Accumulate grads over `batch_size` sequences (tape SUMs them). - let mut step_loss = 0.0f32; + // Sample `batch_size` sequences and run them as ONE batched forward/ + // backward. The CE mean over all batch*seq rows is the batch-mean loss, so + // backward already yields the batch-mean gradient (clip pre-scale = 1.0). + let mut inputs = Vec::with_capacity(cfg.batch_size); + let mut targets_v = Vec::with_capacity(cfg.batch_size); for _ in 0..cfg.batch_size { let (input, target) = corpus.sample(cfg.seq_len, &mut rng); - let ids = ids_tensor(&input, device); - let targets = ids_tensor(&target, device); - let loss = model.loss(&ids, &targets); - step_loss += read_scalar(&loss); - loss.backward(); - tokens_seen += cfg.seq_len as u64; + inputs.push(input); + targets_v.push(target); } - step_loss *= inv_batch; + let ids = batched_ids_tensor(&inputs, device); + let targets = batched_ids_tensor(&targets_v, device); + let loss = model.loss_batched(&ids, &targets, cfg.batch_size); + let step_loss = read_scalar(&loss); + loss.backward(); + tokens_seen += (cfg.batch_size * cfg.seq_len) as u64; losses.push(step_loss); - // Average the summed grads (×1/batch) and clip to the global norm. - let gnorm = clip_grad_norm_gpu(¶ms, cfg.max_grad_norm, inv_batch); + // Backward already produced the batch-mean gradient — just clip it. + let gnorm = clip_grad_norm_gpu(¶ms, cfg.max_grad_norm, 1.0); opt.step(lr, ¶ms); for p in ¶ms { p.zero_grad(); diff --git a/csrc/ops/nn.cu b/csrc/ops/nn.cu index e8dd32a..1cbac91 100644 --- a/csrc/ops/nn.cu +++ b/csrc/ops/nn.cu @@ -215,14 +215,20 @@ void launch_silu_dx_f32(const float* x, const float* dy, float* dx, int n, void* // dx[i+h] = dy[i+h]*cos - dy[i]*sin // ===================================================================== -__global__ void rope_k(const float* x, float* y, int heads, int head_dim, float theta) { +// `period` is the sequence length: a flattened batch lays B sequences end to end +// along the `tokens` axis, so each token's RoPE position is its index WITHIN its +// own sequence, `tok % period`. With period == tokens (single sequence) this is +// the original position = row. +__global__ void rope_k(const float* x, float* y, int heads, int head_dim, + float theta, int period) { int tok = blockIdx.x; int head = blockIdx.y; int half = head_dim / 2; int i = threadIdx.x; if (i >= half) return; + int pos = tok % period; float freq = powf(theta, -(float)(2 * i) / (float)head_dim); - float angle = (float)tok * freq; + float angle = (float)pos * freq; float c = cosf(angle), sn = sinf(angle); int base = (tok * heads + head) * head_dim; float x0 = x[base + i], x1 = x[base + i + half]; @@ -230,20 +236,22 @@ __global__ void rope_k(const float* x, float* y, int heads, int head_dim, float y[base + i + half] = x1 * c + x0 * sn; } void launch_rope_f32(const float* x, float* y, int tokens, int heads, - int head_dim, float theta, void* s) { + int head_dim, float theta, int period, void* s) { dim3 grid(tokens, heads); int blk = head_dim / 2; - rope_k<<>>(x, y, heads, head_dim, theta); + rope_k<<>>(x, y, heads, head_dim, theta, period); } -__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, float theta) { +__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, + float theta, int period) { int tok = blockIdx.x; int head = blockIdx.y; int half = head_dim / 2; int i = threadIdx.x; if (i >= half) return; + int pos = tok % period; float freq = powf(theta, -(float)(2 * i) / (float)head_dim); - float angle = (float)tok * freq; + float angle = (float)pos * freq; float c = cosf(angle), sn = sinf(angle); int base = (tok * heads + head) * head_dim; float d0 = dy[base + i], d1 = dy[base + i + half]; @@ -251,10 +259,10 @@ __global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, f dx[base + i + half] = d1 * c - d0 * sn; } void launch_rope_dx_f32(const float* dy, float* dx, int tokens, int heads, - int head_dim, float theta, void* s) { + int head_dim, float theta, int period, void* s) { dim3 grid(tokens, heads); int blk = head_dim / 2; - rope_dx_k<<>>(dy, dx, heads, head_dim, theta); + rope_dx_k<<>>(dy, dx, heads, head_dim, theta, period); } // =====================================================================