From 386d80039b587c57b1236b4ca52db9bcbc36ad95 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 22 Jun 2026 14:09:59 -0400 Subject: [PATCH 1/3] feat(inference,tune): LoRA backward surface-B weight grads + forward parity test GDN-layer LoRA weight-gradient surface (gdn_backward) plus the full-model trainer path (train_grad_full, train_grad_layer23) and a LoRA forward parity test. Option-gated: the None path is byte-identical to prior behavior. UNVERIFIED: single/multi-head gradcheck is written but not yet executed on a real model. The surface-B weight grads must pass gradcheck before this merges. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/inference/examples/diff_gdn_layer.rs | 20 +- .../inference/src/attention/gdn_backward.rs | 1007 ++++++++++++++++- crates/inference/src/backward/ops.rs | 18 +- crates/inference/src/backward/tape.rs | 12 +- .../tests/lora_forward_parity_test.rs | 917 +++++++++++++++ crates/tune/src/bin/train_grad_full.rs | 724 +++++++++--- crates/tune/src/bin/train_grad_layer23.rs | 39 +- 7 files changed, 2545 insertions(+), 192 deletions(-) create mode 100644 crates/inference/tests/lora_forward_parity_test.rs diff --git a/crates/inference/examples/diff_gdn_layer.rs b/crates/inference/examples/diff_gdn_layer.rs index 120160b7bb..e8aec5ccc0 100644 --- a/crates/inference/examples/diff_gdn_layer.rs +++ b/crates/inference/examples/diff_gdn_layer.rs @@ -99,7 +99,25 @@ fn main() { eps, ); let mut mat_out = vec![0.0f32; seq_len * hidden]; - gdn_forward_save(&normed, gdn_w, cfg, &mut saved, &mut mat_out); + gdn_forward_save( + &normed, + gdn_w, + cfg, + &mut saved, + &mut mat_out, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + 0, + 0.0, + ); assert_eq!(mat_out.len(), real_out.len(), "output length mismatch"); let (mut max_diff, mut argmax) = (0.0f32, 0usize); diff --git a/crates/inference/src/attention/gdn_backward.rs b/crates/inference/src/attention/gdn_backward.rs index a9394da6d0..18d4616774 100644 --- a/crates/inference/src/attention/gdn_backward.rs +++ b/crates/inference/src/attention/gdn_backward.rs @@ -1,9 +1,10 @@ //! Reverse-mode backward (VJP) for a single GatedDeltaNet token sequence. //! //! Scope: computes grad_input (d_loss/d_input for each token) for a sequence -//! processed through one GDN layer. Parameter (weight) gradients are out of -//! scope here — only input-gradients are needed for LoRA gradient flow through -//! the 18 frozen GDN layers to reach lower GQA layers. +//! processed through one GDN layer, AND (when LoRA is present) the LoRA weight +//! gradients for each of the 5 GDN projections (qkv, z, b/beta, a/alpha, out). +//! This is "surface-B" LoRA: mirrors the GQA surface-A implementation in +//! backward/attention_gqa.rs, using the same lora_vjp primitive from ops.rs. //! //! # Forward recap (matches `gdn_fused.rs` hot path) //! @@ -65,8 +66,34 @@ //! eliminate catastrophic cancellation at eps = 1e-3. use crate::attention::gdn::{sigmoid, softplus}; +use crate::backward::ops::lora_vjp; use crate::model::qwen35_config::Qwen35Config; +/// Gradients of the GDN layer w.r.t. LoRA params for all 5 projections. +/// +/// Naming and layout mirror AttnGrads in backward/attention_gqa.rs. +/// grad_a_* shape: [rank * d_in], grad_b_* shape: [d_out * rank]. +/// Fields are zero-length Vec when LoRA is absent (no-LoRA forward). +pub struct GdnGrads { + /// in_proj_qkv LoRA: d_in=hidden, d_out=qkv_dim + pub grad_a_qkv: Vec, + pub grad_b_qkv: Vec, + /// in_proj_z LoRA: d_in=hidden, d_out=output_dim + pub grad_a_z: Vec, + pub grad_b_z: Vec, + /// in_proj_b (beta sigmoid) LoRA: d_in=hidden, d_out=num_kh + pub grad_a_b: Vec, + pub grad_b_b: Vec, + /// in_proj_a (alpha decay) LoRA: d_in=hidden, d_out=num_kh + pub grad_a_a: Vec, + pub grad_b_a: Vec, + /// out_proj LoRA: d_in=output_dim, d_out=hidden + pub grad_a_out: Vec, + pub grad_b_out: Vec, + /// Input gradient (always present): [seq_len, hidden_size] + pub dx: Vec, +} + /// Saved activations for one forward pass over a sequence. /// /// All vectors are laid out as `[seq_len, ...]` row-major. Per-head buffers @@ -140,6 +167,52 @@ pub struct GdnSaved { /// SiLU(z) per value-head: [seq_len, value_heads, value_dim] /// Stored because backward through gated_rms_norm needs it. pub silu_z: Vec, + + // ---- LoRA caches (populated only when LoRA is present) ---- + // h_* = A_proj @ x for each projection; shape [seq_len * lora_rank]. + // Empty when LoRA absent. + /// LoRA rank used (0 = no LoRA) + pub lora_rank: usize, + /// LoRA scale factor + pub lora_scale: f32, + + /// h_qkv = A_qkv @ x: [seq_len, lora_rank] + pub h_qkv: Vec, + /// h_z = A_z @ x: [seq_len, lora_rank] + pub h_z: Vec, + /// h_b = A_b @ x: [seq_len, lora_rank] (beta sigmoid projection) + pub h_b: Vec, + /// h_a = A_a @ x: [seq_len, lora_rank] (alpha decay projection) + pub h_a: Vec, + /// h_out = A_out @ gated_buf: [seq_len, lora_rank] (output projection input) + pub h_out: Vec, + /// gated_buf (input to out_proj + LoRA_out): [seq_len, output_dim] + /// Needed for the out_proj LoRA backward (x side of that linear). + pub gated_buf: Vec, + + // ---- LoRA weight matrices (cloned from forward params for backward) ---- + // Storing these avoids threading 10 extra args through gdn_backward. + // Empty Vec when LoRA absent. + /// A_qkv: [rank, hidden] + pub lora_a_qkv: Vec, + /// B_qkv: [qkv_dim, rank] + pub lora_b_qkv: Vec, + /// A_z: [rank, hidden] + pub lora_a_z: Vec, + /// B_z: [output_dim, rank] + pub lora_b_z: Vec, + /// A_b: [rank, hidden] + pub lora_a_b: Vec, + /// B_b: [num_kh, rank] + pub lora_b_b: Vec, + /// A_a: [rank, hidden] + pub lora_a_a: Vec, + /// B_a: [num_kh, rank] + pub lora_b_a: Vec, + /// A_out: [rank, output_dim] + pub lora_a_out: Vec, + /// B_out: [hidden, rank] + pub lora_b_out: Vec, } impl GdnSaved { @@ -197,6 +270,26 @@ impl GdnSaved { o_heads: vec![0.0; seq_len * value_heads * value_dim], rms_vals: vec![0.0; seq_len * value_heads], silu_z: vec![0.0; seq_len * value_heads * value_dim], + // LoRA caches: start empty; populated in gdn_forward_save when LoRA present. + lora_rank: 0, + lora_scale: 0.0, + h_qkv: Vec::new(), + h_z: Vec::new(), + h_b: Vec::new(), + h_a: Vec::new(), + h_out: Vec::new(), + gated_buf: Vec::new(), + // LoRA weight matrices: start empty; cloned from params in gdn_forward_save. + lora_a_qkv: Vec::new(), + lora_b_qkv: Vec::new(), + lora_a_z: Vec::new(), + lora_b_z: Vec::new(), + lora_a_b: Vec::new(), + lora_b_b: Vec::new(), + lora_a_a: Vec::new(), + lora_b_a: Vec::new(), + lora_a_out: Vec::new(), + lora_b_out: Vec::new(), } } } @@ -211,15 +304,37 @@ impl GdnSaved { /// `inputs`: [seq_len, hidden_size] /// `weights`: frozen GDN layer weights /// `cfg`: model config -/// `norm_weight`: gamma for gated RMSNorm [value_dim] /// `saved`: output struct, must be pre-allocated via `GdnSaved::new` /// `outputs`: [seq_len, hidden_size] — written in-place +/// +/// Optional LoRA parameters for the 5 GDN projections (qkv, z, b, a, out). +/// When all five pairs are `Some`, the forward adds `scale * (x @ A^T) @ B^T` +/// to each projection output and caches `h = A @ x` for the backward. +/// When any pair is `None`, the forward is byte-identical to the no-LoRA path. +#[allow(clippy::too_many_arguments)] pub fn gdn_forward_save( inputs: &[f32], weights: &crate::attention::gdn::GatedDeltaNetWeights, _cfg: &Qwen35Config, saved: &mut GdnSaved, outputs: &mut [f32], + // LoRA params for in_proj_qkv: a=[rank,hidden], b=[qkv_dim,rank] + lora_a_qkv: Option<&[f32]>, + lora_b_qkv: Option<&[f32]>, + // LoRA params for in_proj_z: a=[rank,hidden], b=[output_dim,rank] + lora_a_z: Option<&[f32]>, + lora_b_z: Option<&[f32]>, + // LoRA params for in_proj_b (beta): a=[rank,hidden], b=[num_kh,rank] + lora_a_b: Option<&[f32]>, + lora_b_b: Option<&[f32]>, + // LoRA params for in_proj_a (alpha): a=[rank,hidden], b=[num_kh,rank] + lora_a_a: Option<&[f32]>, + lora_b_a: Option<&[f32]>, + // LoRA params for out_proj: a=[rank,output_dim], b=[hidden,rank] + lora_a_out: Option<&[f32]>, + lora_b_out: Option<&[f32]>, + lora_rank: usize, + lora_scale: f32, ) { use crate::forward::cpu::matmul_bt; @@ -238,6 +353,67 @@ pub fn gdn_forward_save( let buf_len = kernel_size.saturating_sub(1); let q_total = num_kh * key_dim; + // Bind all ten LoRA slices at once. If any pair is None or rank==0 the entire + // LoRA path is skipped and the forward is byte-identical to the no-LoRA path. + // Option<&[f32]> is Copy, so destructuring copies the refs — no heap allocation. + let lora_bound: Option<( + &[f32], + &[f32], // a_qkv, b_qkv + &[f32], + &[f32], // a_z, b_z + &[f32], + &[f32], // a_b, b_b + &[f32], + &[f32], // a_a, b_a + &[f32], + &[f32], // a_out, b_out + )> = if lora_rank > 0 { + match ( + lora_a_qkv, lora_b_qkv, lora_a_z, lora_b_z, lora_a_b, lora_b_b, lora_a_a, lora_b_a, + lora_a_out, lora_b_out, + ) { + ( + Some(a_qkv), + Some(b_qkv), + Some(a_z), + Some(b_z), + Some(a_b), + Some(b_b), + Some(a_a), + Some(b_a), + Some(a_out), + Some(b_out), + ) => Some((a_qkv, b_qkv, a_z, b_z, a_b, b_b, a_a, b_a, a_out, b_out)), + _ => None, + } + } else { + None + }; + + // Initialise LoRA cache buffers on the saved struct when LoRA is active. + // Also clone the weight matrices so gdn_backward doesn't need them threaded through. + if let Some((a_qkv, b_qkv, a_z, b_z, a_b, b_b, a_a, b_a, a_out, b_out)) = lora_bound { + saved.lora_rank = lora_rank; + saved.lora_scale = lora_scale; + saved.h_qkv = vec![0.0f32; seq_len * lora_rank]; + saved.h_z = vec![0.0f32; seq_len * lora_rank]; + saved.h_b = vec![0.0f32; seq_len * lora_rank]; + saved.h_a = vec![0.0f32; seq_len * lora_rank]; + saved.h_out = vec![0.0f32; seq_len * lora_rank]; + saved.gated_buf = vec![0.0f32; seq_len * output_dim]; + // Clone weight matrices for backward use. + saved.lora_a_qkv = a_qkv.to_vec(); + saved.lora_b_qkv = b_qkv.to_vec(); + saved.lora_a_z = a_z.to_vec(); + saved.lora_b_z = b_z.to_vec(); + saved.lora_a_b = a_b.to_vec(); + saved.lora_b_b = b_b.to_vec(); + saved.lora_a_a = a_a.to_vec(); + saved.lora_b_a = b_a.to_vec(); + saved.lora_a_out = a_out.to_vec(); + saved.lora_b_out = b_out.to_vec(); + } + // Rolling conv buffer (shared across time, updated per step) let mut conv_buf_live = vec![0.0f32; qkv_dim * buf_len]; @@ -252,14 +428,104 @@ pub fn gdn_forward_save( let qkv_out = &mut saved.qkv_proj[t * qkv_dim..(t + 1) * qkv_dim]; matmul_bt(x, &weights.in_proj_qkv, qkv_out, 1, hidden, qkv_dim); + // LoRA for in_proj_qkv: output += scale * (x @ A_qkv^T) @ B_qkv^T + // Cache h_qkv = A_qkv @ x for the backward. + if let Some((a_qkv, b_qkv, _, _, _, _, _, _, _, _)) = lora_bound { + let h = &mut saved.h_qkv[t * lora_rank..(t + 1) * lora_rank]; + for r in 0..lora_rank { + h[r] = a_qkv[r * hidden..(r + 1) * hidden] + .iter() + .zip(x.iter()) + .map(|(a, xi)| a * xi) + .sum(); + } + let qkv_out = &mut saved.qkv_proj[t * qkv_dim..(t + 1) * qkv_dim]; + for i in 0..qkv_dim { + let acc: f32 = lora_scale + * b_qkv[i * lora_rank..(i + 1) * lora_rank] + .iter() + .zip(h.iter()) + .map(|(b, hi)| b * hi) + .sum::(); + qkv_out[i] += acc; + } + } + let z_out = &mut saved.z_proj[t * output_dim..(t + 1) * output_dim]; matmul_bt(x, &weights.in_proj_z, z_out, 1, hidden, output_dim); + // LoRA for in_proj_z + if let Some((_, _, a_z, b_z, _, _, _, _, _, _)) = lora_bound { + let h = &mut saved.h_z[t * lora_rank..(t + 1) * lora_rank]; + for r in 0..lora_rank { + h[r] = a_z[r * hidden..(r + 1) * hidden] + .iter() + .zip(x.iter()) + .map(|(a, xi)| a * xi) + .sum(); + } + let z_out = &mut saved.z_proj[t * output_dim..(t + 1) * output_dim]; + for i in 0..output_dim { + let acc: f32 = lora_scale + * b_z[i * lora_rank..(i + 1) * lora_rank] + .iter() + .zip(h.iter()) + .map(|(b, hi)| b * hi) + .sum::(); + z_out[i] += acc; + } + } + let beta_out = &mut saved.beta_raw[t * num_kh..(t + 1) * num_kh]; matmul_bt(x, &weights.in_proj_b, beta_out, 1, hidden, num_kh); + + // LoRA for in_proj_b (beta, pre-sigmoid linear output) + if let Some((_, _, _, _, a_b, b_b, _, _, _, _)) = lora_bound { + let h = &mut saved.h_b[t * lora_rank..(t + 1) * lora_rank]; + for r in 0..lora_rank { + h[r] = a_b[r * hidden..(r + 1) * hidden] + .iter() + .zip(x.iter()) + .map(|(a, xi)| a * xi) + .sum(); + } + let beta_out = &mut saved.beta_raw[t * num_kh..(t + 1) * num_kh]; + for i in 0..num_kh { + let acc: f32 = lora_scale + * b_b[i * lora_rank..(i + 1) * lora_rank] + .iter() + .zip(h.iter()) + .map(|(b, hi)| b * hi) + .sum::(); + beta_out[i] += acc; + } + } + let alpha_out = &mut saved.alpha_proj[t * num_kh..(t + 1) * num_kh]; matmul_bt(x, &weights.in_proj_a, alpha_out, 1, hidden, num_kh); + // LoRA for in_proj_a (alpha, pre-softplus linear output) + if let Some((_, _, _, _, _, _, a_a, b_a, _, _)) = lora_bound { + let h = &mut saved.h_a[t * lora_rank..(t + 1) * lora_rank]; + for r in 0..lora_rank { + h[r] = a_a[r * hidden..(r + 1) * hidden] + .iter() + .zip(x.iter()) + .map(|(a, xi)| a * xi) + .sum(); + } + let alpha_out = &mut saved.alpha_proj[t * num_kh..(t + 1) * num_kh]; + for i in 0..num_kh { + let acc: f32 = lora_scale + * b_a[i * lora_rank..(i + 1) * lora_rank] + .iter() + .zip(h.iter()) + .map(|(b, hi)| b * hi) + .sum::(); + alpha_out[i] += acc; + } + } + // sigmoid(beta) for kh in 0..num_kh { let raw = saved.beta_raw[t * num_kh + kh]; @@ -409,8 +675,37 @@ pub fn gdn_forward_save( } } + // Save gated_buf when LoRA is active (needed as the x-side input for out_proj LoRA backward). + if lora_bound.is_some() { + saved.gated_buf[t * output_dim..(t + 1) * output_dim].copy_from_slice(&gated_buf); + } + let y_t = &mut outputs[t * hidden..(t + 1) * hidden]; matmul_bt(&gated_buf, &weights.out_proj, y_t, 1, output_dim, hidden); + + // LoRA for out_proj: y_t += scale * (gated_buf @ A_out^T) @ B_out^T + // out_proj is [hidden, output_dim], so d_out=hidden, d_in=output_dim. + // Cache h_out = A_out @ gated_buf for the backward. + if let Some((_, _, _, _, _, _, _, _, a_out, b_out)) = lora_bound { + let h_slot = &mut saved.h_out[t * lora_rank..(t + 1) * lora_rank]; + for r in 0..lora_rank { + h_slot[r] = a_out[r * output_dim..(r + 1) * output_dim] + .iter() + .zip(gated_buf.iter()) + .map(|(a, xi)| a * xi) + .sum(); + } + let y_t = &mut outputs[t * hidden..(t + 1) * hidden]; + for i in 0..hidden { + let acc: f32 = lora_scale + * b_out[i * lora_rank..(i + 1) * lora_rank] + .iter() + .zip(h_slot.iter()) + .map(|(b, hi)| b * hi) + .sum::(); + y_t[i] += acc; + } + } } } @@ -420,16 +715,19 @@ pub fn gdn_forward_save( /// VJP of the GDN sequence forward. /// +/// Returns a `GdnGrads` containing the input gradient (dx, always populated) +/// and — when the saved activations carry LoRA caches (i.e. `gdn_forward_save` +/// was called with LoRA params) — the weight gradients for each of the 5 +/// projections. When LoRA is absent, the grad_a_*/grad_b_* Vecs are empty. +/// /// `grad_outputs`: [seq_len, hidden_size] — upstream gradient of loss w.r.t. output /// `saved`: activations from `gdn_forward_save` /// `weights`: same frozen weights used in forward -/// `grad_inputs`: [seq_len, hidden_size] — output, grad of loss w.r.t. input pub fn gdn_backward( grad_outputs: &[f32], saved: &GdnSaved, weights: &crate::attention::gdn::GatedDeltaNetWeights, - grad_inputs: &mut [f32], -) { +) -> GdnGrads { let seq_len = saved.seq_len; let hidden = saved.hidden_size; let num_kh = saved.num_key_heads; @@ -445,7 +743,63 @@ pub fn gdn_backward( let q_total = num_kh * key_dim; let gamma = &weights.norm_weight[..value_dim]; - grad_inputs.fill(0.0); + let have_lora = saved.lora_rank > 0 && !saved.h_qkv.is_empty(); + let lora_rank = saved.lora_rank; + let lora_scale = saved.lora_scale; + + // Allocate LoRA weight grad accumulators (empty when no LoRA). + let mut grad_a_qkv = if have_lora { + vec![0.0f32; lora_rank * hidden] + } else { + Vec::new() + }; + let mut grad_b_qkv = if have_lora { + vec![0.0f32; qkv_dim * lora_rank] + } else { + Vec::new() + }; + let mut grad_a_z = if have_lora { + vec![0.0f32; lora_rank * hidden] + } else { + Vec::new() + }; + let mut grad_b_z = if have_lora { + vec![0.0f32; output_dim * lora_rank] + } else { + Vec::new() + }; + let mut grad_a_b = if have_lora { + vec![0.0f32; lora_rank * hidden] + } else { + Vec::new() + }; + let mut grad_b_b = if have_lora { + vec![0.0f32; num_kh * lora_rank] + } else { + Vec::new() + }; + let mut grad_a_a = if have_lora { + vec![0.0f32; lora_rank * hidden] + } else { + Vec::new() + }; + let mut grad_b_a = if have_lora { + vec![0.0f32; num_kh * lora_rank] + } else { + Vec::new() + }; + let mut grad_a_out = if have_lora { + vec![0.0f32; lora_rank * output_dim] + } else { + Vec::new() + }; + let mut grad_b_out = if have_lora { + vec![0.0f32; hidden * lora_rank] + } else { + Vec::new() + }; + + let mut grad_inputs = vec![0.0f32; seq_len * hidden]; // Adjoint state dS: accumulated across timesteps, flows backward. // dS[h] is dL/d(S_t) for head h, updated as t decreases. @@ -464,10 +818,17 @@ pub fn gdn_backward( for t in (0..seq_len).rev() { let dy = &grad_outputs[t * hidden..(t + 1) * hidden]; + let x_t = &saved.inputs[t * hidden..(t + 1) * hidden]; // ---- 1. Backward through out_proj: d_gated_buf = W_out^T @ dy ---- // W_out is [hidden, output_dim]. out = gated_buf @ W_out^T means // d_gated_buf[j] = sum_i W_out[i,j] * dy[i] + // + // When LoRA is present: y_t = base_out + scale*(gated_buf@A_out^T)@B_out^T + // The upstream gradient dy is w.r.t. the total y_t. + // lora_vjp(dy, gated_buf_t, h_out_t, a_out, b_out, rank, output_dim, hidden, scale) + // returns (grad_b_out delta, grad_a_out delta, d_gated_buf_lora_delta). + // d_gated_buf = W_out^T dy + lora_vjp dx contribution (d_gated_buf from LoRA path). let mut d_gated_buf = vec![0.0f32; output_dim]; for j in 0..output_dim { let mut acc = 0.0f64; @@ -477,6 +838,35 @@ pub fn gdn_backward( d_gated_buf[j] = acc as f32; } + // LoRA weight grads for out_proj + dx contribution to d_gated_buf. + // out_proj layout: [hidden, output_dim] → d_in=output_dim, d_out=hidden. + // lora_vjp(g=dy, x=gated_buf_t, h=h_out_t, a=lora_a_out, b=lora_b_out, + // rank, d_in=output_dim, d_out=hidden, scale) + if have_lora { + let gated_buf_t = &saved.gated_buf[t * output_dim..(t + 1) * output_dim]; + let h_out_t = &saved.h_out[t * lora_rank..(t + 1) * lora_rank]; + let (gb, ga, dx_lora) = lora_vjp( + dy, + gated_buf_t, + h_out_t, + &saved.lora_a_out, + &saved.lora_b_out, + lora_rank, + output_dim, + hidden, + lora_scale, + ); + for k in 0..ga.len() { + grad_a_out[k] += ga[k]; + } + for k in 0..gb.len() { + grad_b_out[k] += gb[k]; + } + for j in 0..output_dim { + d_gated_buf[j] += dx_lora[j]; + } + } + // ---- 2. Backward through gated RMSNorm for each value-head ---- let z_slice = &saved.z_proj[t * output_dim..(t + 1) * output_dim]; let mut d_o_heads = vec![0.0f32; output_dim]; @@ -525,8 +915,13 @@ pub fn gdn_backward( } } - // ---- 3. Accumulate d_z_proj → d_x via W_z^T ---- - // z_proj = W_z @ x, d_x += W_z^T @ d_z_proj + // ---- 3. Accumulate d_z_proj → d_x via W_z^T, then LoRA weight grads ---- + // z_proj = W_z @ x (+LoRA when present), d_x += W_z^T @ d_z_proj + // + // d_z_proj IS the pre-nonlinearity linear-projection-output gradient for z. + // (z goes through SiLU in the gated RMSNorm; d_z_proj is the gradient + // w.r.t. the linear output of in_proj_z, BEFORE SiLU — that is what step 2 + // computes via the SiLU backward inside the RMSNorm loop.) let dx_t = &mut grad_inputs[t * hidden..(t + 1) * hidden]; for j in 0..output_dim { let dz_j = d_z_proj[j]; @@ -537,12 +932,42 @@ pub fn gdn_backward( dx_t[i] += weights.in_proj_z[j * hidden + i] * dz_j; } } + // LoRA weight grads for in_proj_z. + // g = d_z_proj, x = x_t, h = h_z_t, d_in=hidden, d_out=output_dim + if have_lora { + let h_z_t = &saved.h_z[t * lora_rank..(t + 1) * lora_rank]; + let (gb, ga, dx_lora) = lora_vjp( + &d_z_proj, + x_t, + h_z_t, + &saved.lora_a_z, + &saved.lora_b_z, + lora_rank, + hidden, + output_dim, + lora_scale, + ); + for k in 0..ga.len() { + grad_a_z[k] += ga[k]; + } + for k in 0..gb.len() { + grad_b_z[k] += gb[k]; + } + for j in 0..hidden { + dx_t[j] += dx_lora[j]; + } + } // ---- 4. Per value-head recurrence backward ---- // We process heads in REVERSE order (arbitrary — no inter-head deps). // We work with the per-head adjoint dS[h] which carries across time. let mut d_conv_out = vec![0.0f32; qkv_dim]; + // Per key-head accumulators for alpha and beta scalar grads (multiple + // value-heads share one key-head, so we must sum across value-heads first + // before applying lora_vjp to avoid double-counting the d_in/d_out shape). + let mut d_alpha_kh = vec![0.0f32; num_kh]; // d_loss / d(alpha_proj[t, kh]) + let mut d_beta_raw_kh = vec![0.0f32; num_kh]; // d_loss / d(beta_raw[t, kh]) for h in 0..value_heads { let kh = h / ratio; @@ -709,19 +1134,78 @@ pub fn gdn_backward( let sig = sigmoid(beta_raw_h); let d_beta_raw_h = d_beta_h * sig * (1.0 - sig); - // ---- 4i. Accumulate scalar grads into per-head gradient arrays ---- - // We accumulate per-head d_alpha and d_beta_raw into d_alpha_proj / d_beta_proj - // arrays. Since multiple value-heads share the same key-head, we sum. - // We'll accumulate directly into d_x via W_a^T and W_b^T below. - // Store in temporary scalars indexed by kh (accumulate across h sharing kh). - // Use a local accumulator since the inner-most scope is per-h. - // We immediately push to d_x to avoid extra allocations. - let dx_t = &mut grad_inputs[t * hidden..(t + 1) * hidden]; + // ---- 4i. Accumulate per key-head scalar grads ---- + // Multiple value-heads share the same key-head; accumulate into kh slots. + // These will be fed to W_a^T/W_b^T and lora_vjp after this loop. + d_alpha_kh[kh] += d_alpha_from_g; + d_beta_raw_kh[kh] += d_beta_raw_h; + } // end per-head loop + + // ---- Apply per key-head alpha/beta grads → d_x and LoRA weight grads ---- + // d_alpha_proj[t, kh] = d_alpha_kh[kh] (scalar per head) + // d_beta_raw[t, kh] = d_beta_raw_kh[kh] + // + // For lora_vjp on the beta and alpha projections, we need the full + // d_proj_output vector of shape [num_kh]. We built those above. + // + // g for in_proj_b = d_beta_raw_kh (pre-sigmoid linear-output gradient) + // g for in_proj_a = d_alpha_kh (pre-softplus linear-output gradient) + let dx_t = &mut grad_inputs[t * hidden..(t + 1) * hidden]; + for kh in 0..num_kh { + let d_a = d_alpha_kh[kh]; + let d_b = d_beta_raw_kh[kh]; for i in 0..hidden { - dx_t[i] += weights.in_proj_a[kh * hidden + i] * d_alpha_from_g; - dx_t[i] += weights.in_proj_b[kh * hidden + i] * d_beta_raw_h; + dx_t[i] += weights.in_proj_a[kh * hidden + i] * d_a; + dx_t[i] += weights.in_proj_b[kh * hidden + i] * d_b; + } + } + // LoRA weight grads for in_proj_a and in_proj_b. + // g = d_alpha_kh (full [num_kh] vector), x = x_t, h_a_t, d_in=hidden, d_out=num_kh + if have_lora { + let h_b_t = &saved.h_b[t * lora_rank..(t + 1) * lora_rank]; + let (gb, ga, dx_lora) = lora_vjp( + &d_beta_raw_kh, + x_t, + h_b_t, + &saved.lora_a_b, + &saved.lora_b_b, + lora_rank, + hidden, + num_kh, + lora_scale, + ); + for k in 0..ga.len() { + grad_a_b[k] += ga[k]; + } + for k in 0..gb.len() { + grad_b_b[k] += gb[k]; + } + for j in 0..hidden { + dx_t[j] += dx_lora[j]; } - } // end per-head loop + + let h_a_t = &saved.h_a[t * lora_rank..(t + 1) * lora_rank]; + let (gb, ga, dx_lora) = lora_vjp( + &d_alpha_kh, + x_t, + h_a_t, + &saved.lora_a_a, + &saved.lora_b_a, + lora_rank, + hidden, + num_kh, + lora_scale, + ); + for k in 0..ga.len() { + grad_a_a[k] += ga[k]; + } + for k in 0..gb.len() { + grad_b_a[k] += gb[k]; + } + for j in 0..hidden { + dx_t[j] += dx_lora[j]; + } + } // ---- 5. Backward through conv1d + SiLU ---- // @@ -778,6 +1262,9 @@ pub fn gdn_backward( // d_qkv_proj_all[t] now contains the complete gradient for token t because: // - future-timestep conv contributions were written when t' > t was processed // - current-timestep conv contribution was written in step 5 above + // + // g for in_proj_qkv = d_qkv_proj_all[t] (pre-conv-SiLU linear-output gradient + // summed over time, complete only after all future-t conv contributions land). let dx_t = &mut grad_inputs[t * hidden..(t + 1) * hidden]; for j in 0..qkv_dim { let dq = d_qkv_proj_all[t * qkv_dim + j]; @@ -788,7 +1275,51 @@ pub fn gdn_backward( dx_t[i] += weights.in_proj_qkv[j * hidden + i] * dq; } } + // LoRA weight grads for in_proj_qkv. + // The LoRA forward added scale*(x@A_qkv^T)@B_qkv^T to qkv_proj[t]. + // That delta passed through conv1d+SiLU, whose backward already produced + // d_qkv_proj_all[t] as the full gradient w.r.t. qkv_proj[t] (both base + // and LoRA delta). So d_qkv_proj_all[t] is exactly the g for lora_vjp. + // x = x_t (the original input), h = h_qkv_t. + if have_lora { + let d_qkv_g = &d_qkv_proj_all[t * qkv_dim..(t + 1) * qkv_dim]; + let h_qkv_t = &saved.h_qkv[t * lora_rank..(t + 1) * lora_rank]; + let (gb, ga, dx_lora) = lora_vjp( + d_qkv_g, + x_t, + h_qkv_t, + &saved.lora_a_qkv, + &saved.lora_b_qkv, + lora_rank, + hidden, + qkv_dim, + lora_scale, + ); + for k in 0..ga.len() { + grad_a_qkv[k] += ga[k]; + } + for k in 0..gb.len() { + grad_b_qkv[k] += gb[k]; + } + for j in 0..hidden { + dx_t[j] += dx_lora[j]; + } + } } // end reverse time loop + + GdnGrads { + grad_a_qkv, + grad_b_qkv, + grad_a_z, + grad_b_z, + grad_a_b, + grad_b_b, + grad_a_a, + grad_b_a, + grad_a_out, + grad_b_out, + dx: grad_inputs, + } } // --------------------------------------------------------------------------- @@ -1036,8 +1567,44 @@ mod tests { let mut out_p = vec![0.0f32; seq_len * hidden]; let mut out_m = vec![0.0f32; seq_len * hidden]; - gdn_forward_save(&inp_p_f32, weights, cfg, &mut saved_p, &mut out_p); - gdn_forward_save(&inp_m_f32, weights, cfg, &mut saved_m, &mut out_m); + gdn_forward_save( + &inp_p_f32, + weights, + cfg, + &mut saved_p, + &mut out_p, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + 0, + 0.0, + ); + gdn_forward_save( + &inp_m_f32, + weights, + cfg, + &mut saved_m, + &mut out_m, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + 0, + 0.0, + ); let lp = linear_loss(&out_p, coeffs); let lm = linear_loss(&out_m, coeffs); @@ -1099,9 +1666,27 @@ mod tests { cfg.rms_norm_eps, ); let mut outputs = vec![0.0f32; seq_len * hidden]; - gdn_forward_save(&inputs, &weights, &cfg, &mut saved, &mut outputs); - let mut analytic = vec![0.0f32; seq_len * hidden]; - gdn_backward(&coeffs, &saved, &weights, &mut analytic); + gdn_forward_save( + &inputs, + &weights, + &cfg, + &mut saved, + &mut outputs, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + 0, + 0.0, + ); + let gdn_grads = gdn_backward(&coeffs, &saved, &weights); + let analytic = gdn_grads.dx; // FD let fd = fd_grad_inputs_linear(&inputs, &weights, &cfg, seq_len, hidden, eps, &coeffs); @@ -1212,11 +1797,29 @@ mod tests { cfg.rms_norm_eps, ); let mut outputs = vec![0.0f32; seq_len * hidden]; - gdn_forward_save(&inputs, &weights, &cfg, &mut saved, &mut outputs); + gdn_forward_save( + &inputs, + &weights, + &cfg, + &mut saved, + &mut outputs, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + 0, + 0.0, + ); let grad_out = vec![0.0f32; seq_len * hidden]; - let mut analytic = vec![0.0f32; seq_len * hidden]; - gdn_backward(&grad_out, &saved, &weights, &mut analytic); + let gdn_grads = gdn_backward(&grad_out, &saved, &weights); + let analytic = gdn_grads.dx; for (i, &v) in analytic.iter().enumerate() { assert!( @@ -1225,4 +1828,350 @@ mod tests { ); } } + + // --------------------------------------------------------------------------- + // LoRA weight-grad gradcheck + // --------------------------------------------------------------------------- + // + // Finite-difference check for the 10 LoRA weight gradients (grad_a_*/grad_b_*). + // The existing gradcheck_gdn_backward tests only cover dx with LoRA OFF. + // This test runs gdn_forward_save with LoRA ON and checks that the weight grads + // returned by gdn_backward agree with central-difference estimates. + // + // DO NOT RUN this test manually until the AM gradcheck pass is scheduled. + // (Compile-gate only — verifies the test compiles under --all-targets.) + + /// Helper: run one forward + backward with the given LoRA arrays and return the + /// full GdnGrads. All ten LoRA slices must be Some. + #[allow(clippy::too_many_arguments)] + fn run_lora_forward_backward( + inputs: &[f32], + weights: &GatedDeltaNetWeights, + cfg: &Qwen35Config, + seq_len: usize, + coeffs: &[f32], + lora_rank: usize, + lora_scale: f32, + a_qkv: &[f32], + b_qkv: &[f32], + a_z: &[f32], + b_z: &[f32], + a_b: &[f32], + b_b: &[f32], + a_a: &[f32], + b_a: &[f32], + a_out: &[f32], + b_out: &[f32], + ) -> GdnGrads { + let hidden = cfg.hidden_size; + let num_kh = cfg.linear_num_key_heads; + let num_vh = cfg.linear_num_value_heads(); + let key_dim = cfg.linear_key_head_dim; + let value_dim = cfg.linear_value_head_dim; + let kernel_size = cfg.linear_conv_kernel_dim; + let qkv_dim = cfg.linear_qkv_dim(); + let output_dim = cfg.linear_output_dim(); + let scale = 1.0 / (key_dim as f32).sqrt(); + + let mut saved = GdnSaved::new( + seq_len, + num_kh, + num_vh, + key_dim, + value_dim, + hidden, + qkv_dim, + output_dim, + kernel_size, + scale, + cfg.rms_norm_eps, + ); + let mut outputs = vec![0.0f32; seq_len * hidden]; + gdn_forward_save( + inputs, + weights, + cfg, + &mut saved, + &mut outputs, + Some(a_qkv), + Some(b_qkv), + Some(a_z), + Some(b_z), + Some(a_b), + Some(b_b), + Some(a_a), + Some(b_a), + Some(a_out), + Some(b_out), + lora_rank, + lora_scale, + ); + gdn_backward(coeffs, &saved, weights) + } + + /// Central-difference estimate of d(loss)/d(param[idx]) for a single entry. + #[allow(clippy::too_many_arguments)] + fn fd_lora_weight_entry( + inputs: &[f32], + weights: &GatedDeltaNetWeights, + cfg: &Qwen35Config, + seq_len: usize, + coeffs: &[f32], + lora_rank: usize, + lora_scale: f32, + a_qkv: &[f32], + b_qkv: &[f32], + a_z: &[f32], + b_z: &[f32], + a_b: &[f32], + b_b: &[f32], + a_a: &[f32], + b_a: &[f32], + a_out: &[f32], + b_out: &[f32], + // Which of the 10 arrays to perturb, and which entry: + which: usize, // 0=a_qkv 1=b_qkv 2=a_z 3=b_z 4=a_b 5=b_b 6=a_a 7=b_a 8=a_out 9=b_out + idx: usize, + eps: f32, + ) -> f32 { + // Helper: clone an array, perturb entry [idx] by delta, re-run, return loss. + let perturbed_loss = |delta: f32| -> f32 { + let mut aq = a_qkv.to_vec(); + let mut bq = b_qkv.to_vec(); + let mut az = a_z.to_vec(); + let mut bz = b_z.to_vec(); + let mut ab = a_b.to_vec(); + let mut bb = b_b.to_vec(); + let mut aa = a_a.to_vec(); + let mut ba = b_a.to_vec(); + let mut ao = a_out.to_vec(); + let mut bo = b_out.to_vec(); + let arr: &mut Vec = match which { + 0 => &mut aq, + 1 => &mut bq, + 2 => &mut az, + 3 => &mut bz, + 4 => &mut ab, + 5 => &mut bb, + 6 => &mut aa, + 7 => &mut ba, + 8 => &mut ao, + _ => &mut bo, + }; + arr[idx] += delta; + let hidden = cfg.hidden_size; + let num_kh = cfg.linear_num_key_heads; + let num_vh = cfg.linear_num_value_heads(); + let key_dim = cfg.linear_key_head_dim; + let value_dim = cfg.linear_value_head_dim; + let kernel_size = cfg.linear_conv_kernel_dim; + let qkv_dim = cfg.linear_qkv_dim(); + let output_dim = cfg.linear_output_dim(); + let scale = 1.0 / (key_dim as f32).sqrt(); + let mut saved = GdnSaved::new( + seq_len, + num_kh, + num_vh, + key_dim, + value_dim, + hidden, + qkv_dim, + output_dim, + kernel_size, + scale, + cfg.rms_norm_eps, + ); + let mut outputs = vec![0.0f32; seq_len * hidden]; + gdn_forward_save( + inputs, + weights, + cfg, + &mut saved, + &mut outputs, + Some(&aq), + Some(&bq), + Some(&az), + Some(&bz), + Some(&ab), + Some(&bb), + Some(&aa), + Some(&ba), + Some(&ao), + Some(&bo), + lora_rank, + lora_scale, + ); + outputs + .iter() + .zip(coeffs.iter()) + .map(|(&y, &c)| y * c) + .sum() + }; + let lp = perturbed_loss(eps); + let lm = perturbed_loss(-eps); + (lp - lm) / (2.0 * eps) + } + + // Core helper that exercises weight-grad gradcheck for any fixture. + // Returns (max_rel_err, n_tested). + #[allow(clippy::too_many_arguments)] + fn run_lora_weight_gradcheck( + hidden: usize, + num_kh: usize, + num_vh: usize, + key_dim: usize, + value_dim: usize, + kernel_size: usize, + seq_len: usize, + lora_rank: usize, + lora_scale: f32, + weight_seed: u64, + input_seed: u64, + lora_seed: u64, + coeff_seed: u64, + eps: f32, + ) -> (f32, usize) { + let cfg = tiny_cfg(hidden, num_kh, num_vh, key_dim, value_dim, kernel_size); + let weights = make_tiny_weights( + hidden, + num_kh, + num_vh, + key_dim, + value_dim, + kernel_size, + weight_seed, + ); + let qkv_dim = cfg.linear_qkv_dim(); + let output_dim = cfg.linear_output_dim(); + + // Inputs + let mut rng_in = Rng::new(input_seed); + let mut inputs = vec![0.0f32; seq_len * hidden]; + rng_in.fill(&mut inputs, -0.5, 0.5); + + // Upstream coeffs (linear loss) + let mut rng_c = Rng::new(coeff_seed); + let mut coeffs = vec![0.0f32; seq_len * hidden]; + rng_c.fill(&mut coeffs, -1.0, 1.0); + + // LoRA matrices — small random values so grad_A and grad_B are non-vacuous. + // Shapes: A=[rank, d_in], B=[d_out, rank] for each projection. + let mut rng_l = Rng::new(lora_seed); + let mut a_qkv = vec![0.0f32; lora_rank * hidden]; // [rank, hidden] + let mut b_qkv = vec![0.0f32; qkv_dim * lora_rank]; // [qkv_dim, rank] + let mut a_z = vec![0.0f32; lora_rank * hidden]; // [rank, hidden] + let mut b_z = vec![0.0f32; output_dim * lora_rank]; // [output_dim, rank] + let mut a_b = vec![0.0f32; lora_rank * hidden]; // [rank, hidden] + let mut b_b = vec![0.0f32; num_kh * lora_rank]; // [num_kh, rank] + let mut a_a = vec![0.0f32; lora_rank * hidden]; // [rank, hidden] + let mut b_a = vec![0.0f32; num_kh * lora_rank]; // [num_kh, rank] + let mut a_out = vec![0.0f32; lora_rank * output_dim]; // [rank, output_dim] + let mut b_out = vec![0.0f32; hidden * lora_rank]; // [hidden, rank] + for arr in [ + &mut a_qkv, &mut b_qkv, &mut a_z, &mut b_z, &mut a_b, &mut b_b, &mut a_a, &mut b_a, + &mut a_out, &mut b_out, + ] + .iter_mut() + { + rng_l.fill(arr, -0.05, 0.05); + } + + // Analytic weight grads via forward+backward. + let grads = run_lora_forward_backward( + &inputs, &weights, &cfg, seq_len, &coeffs, lora_rank, lora_scale, &a_qkv, &b_qkv, &a_z, + &b_z, &a_b, &b_b, &a_a, &b_a, &a_out, &b_out, + ); + + // Collect (which_array, analytic_grad_slice, array_len) triples. + let analytic_arrays: [(&[f32], usize); 10] = [ + (&grads.grad_a_qkv, a_qkv.len()), + (&grads.grad_b_qkv, b_qkv.len()), + (&grads.grad_a_z, a_z.len()), + (&grads.grad_b_z, b_z.len()), + (&grads.grad_a_b, a_b.len()), + (&grads.grad_b_b, b_b.len()), + (&grads.grad_a_a, a_a.len()), + (&grads.grad_b_a, b_a.len()), + (&grads.grad_a_out, a_out.len()), + (&grads.grad_b_out, b_out.len()), + ]; + + let mut max_rel = 0.0f32; + let mut n_tested = 0usize; + + for (which, (analytic_slice, arr_len)) in analytic_arrays.iter().enumerate() { + // Sample up to 6 entries spread across each array. + let step = (arr_len / 6).max(1); + let mut idx = 0usize; + while idx < *arr_len { + let analytic_v = analytic_slice[idx]; + let fd_v = fd_lora_weight_entry( + &inputs, &weights, &cfg, seq_len, &coeffs, lora_rank, lora_scale, &a_qkv, + &b_qkv, &a_z, &b_z, &a_b, &b_b, &a_a, &b_a, &a_out, &b_out, which, idx, eps, + ); + let mag = fd_v.abs().max(analytic_v.abs()); + if mag >= 1e-4 { + n_tested += 1; + let rel = (fd_v - analytic_v).abs() / mag; + if rel > max_rel { + max_rel = rel; + } + } + idx += step; + } + } + + (max_rel, n_tested) + } + + #[test] + fn gradcheck_gdn_lora_weight_grads() { + // Tiny GDN: hidden=32, num_kh=1, value_heads=1, key_dim=8, value_dim=8, + // kernel_size=3, seq=4; rank=2; scale=0.5. + let (max_rel, n_tested) = run_lora_weight_gradcheck( + 32, 1, 1, 8, 8, 3, 4, // hidden, num_kh, num_vh, key_dim, value_dim, kernel, seq + 2, // lora_rank + 0.5, // lora_scale + 42, // weight_seed + 1337, // input_seed + 777, // lora_seed + 999, // coeff_seed + 1e-3, // eps + ); + assert!( + n_tested >= 10, + "lora weight gradcheck: too few testable entries ({n_tested}); \ + increase array sizes or tighten threshold" + ); + assert!( + max_rel < 1e-2, + "gdn lora weight gradcheck failed: max_rel={max_rel:.3e} \ + (n_tested={n_tested})" + ); + } + + #[test] + fn gradcheck_gdn_lora_weight_grads_multi_head() { + // Multi-head variant: num_kh=2, value_heads=4 exercises alpha/beta head-sharing reduction. + let (max_rel, n_tested) = run_lora_weight_gradcheck( + 32, 2, 4, 8, 8, 3, + 4, // hidden, num_kh=2, num_vh=4, key_dim=8, value_dim=8, kernel, seq + 2, // lora_rank + 0.5, // lora_scale + 99, // weight_seed + 2024, // input_seed + 555, // lora_seed + 111, // coeff_seed + 1e-3, // eps + ); + assert!( + n_tested >= 10, + "lora weight gradcheck (multi-head): too few testable entries ({n_tested})" + ); + assert!( + max_rel < 1e-2, + "gdn lora weight gradcheck (multi-head) failed: max_rel={max_rel:.3e} \ + (n_tested={n_tested})" + ); + } } diff --git a/crates/inference/src/backward/ops.rs b/crates/inference/src/backward/ops.rs index e2314b9af7..76a76b250c 100644 --- a/crates/inference/src/backward/ops.rs +++ b/crates/inference/src/backward/ops.rs @@ -133,9 +133,15 @@ pub fn lora_vjp( /// RMSNorm backward (weight frozen). /// -/// Forward: y_i = (x_i / rms) * w_i where rms = sqrt(mean(x^2) + eps) +/// Forward: y_i = x_i * (1 + w_i) * inv_rms where inv_rms = 1/sqrt(mean(x^2) + eps) +/// Shifted-gamma convention — matches `qwen35_rms_norm` in norm.rs. +/// /// Backward: given g = dL/dy, -/// dL/dx_i = (w_i * g_i) / rms - x_i / (D * rms^3) * Σ_j x_j * w_j * g_j +/// dL/dx_j = (1 + w_j) * g_j * inv_rms - x_j * inv_rms^3/D * Σ_i g_i*(1+w_i)*x_i +/// +/// Derivation: y_i = (1+w_i)*x_i*r, r = (mean(x^2)+eps)^(-1/2). +/// ∂r/∂x_j = -x_j * r^3 / D +/// dL/dx_j = (1+w_j)*r*g_j + Σ_i g_i*(1+w_i)*x_i * (-x_j*r^3/D) /// /// `x` is the pre-norm input, `w` is the RMSNorm weight (both length D). /// `inv_rms` = 1 / sqrt(mean(x^2) + eps) from the forward. @@ -145,14 +151,14 @@ pub fn rmsnorm_backward(x: &[f32], w: &[f32], inv_rms: f32, g: &[f32]) -> Vec (Vec, f32) { let d = x.len(); let mean_sq: f32 = x.iter().map(|xi| xi * xi).sum::() / d as f32; @@ -74,7 +77,7 @@ pub fn rms_norm_forward(x: &[f32], w: &[f32], eps: f32) -> (Vec, f32) { let normed: Vec = x .iter() .zip(w.iter()) - .map(|(xi, wi)| xi * wi * inv_rms) + .map(|(xi, wi)| xi * (1.0 + wi) * inv_rms) .collect(); (normed, inv_rms) } @@ -132,12 +135,15 @@ mod tests { #[test] fn rms_norm_forward_roundtrip() { let x = vec![1.0f32, 2.0, 3.0, 4.0]; - let w = vec![1.0f32; 4]; + // Use w=0.5 (nonzero, non-unity) so shifted vs plain gamma produce distinct results. + // Shifted convention: y_i = x_i * (1 + w_i) * inv_rms = x_i * 1.5 * inv_rms + let w = vec![0.5f32; 4]; let (normed, inv_rms) = rms_norm_forward(&x, &w, 1e-6); let mean_sq: f32 = x.iter().map(|xi| xi * xi).sum::() / 4.0; let expected_inv = 1.0 / (mean_sq + 1e-6f32).sqrt(); assert!((inv_rms - expected_inv).abs() < 1e-6, "inv_rms mismatch"); - let expected_norm: Vec = x.iter().map(|xi| xi * expected_inv).collect(); + // Shifted gamma: expected = x_i * (1 + w_i) * inv_rms = x_i * 1.5 * inv_rms + let expected_norm: Vec = x.iter().map(|xi| xi * 1.5 * expected_inv).collect(); for (a, b) in normed.iter().zip(expected_norm.iter()) { assert!((a - b).abs() < 1e-5, "normed mismatch {a} vs {b}"); } diff --git a/crates/inference/tests/lora_forward_parity_test.rs b/crates/inference/tests/lora_forward_parity_test.rs new file mode 100644 index 0000000000..db5aed3237 --- /dev/null +++ b/crates/inference/tests/lora_forward_parity_test.rs @@ -0,0 +1,917 @@ +//! Forward-parity differential tests for the LoRA backward surface. +//! +//! These tests MEASURE divergence between the materialised forward used +//! by the backward tape and the real Qwen3.5 model primitives. They do +//! NOT modify production code. The numbers reported here are the gate for +//! any on-par claim (Leo's bar: forward must match the real model under +//! 1e-3 before an on-par claim is valid). +//! +//! Run: +//! cargo test -p lattice-inference --test lora_forward_parity_test \ +//! --features train-backward -- --nocapture +//! +//! TEST 1: RMSNorm convention diff +//! tape.rs uses plain gamma: xi * wi * inv_rms +//! qwen35_rms_norm uses shifted gamma: xi * inv_rms * (1 + gi) +//! Report: elementwise max-diff between the two on the same inputs. +//! +//! TEST 2: GQA single-layer forward parity, materialised vs oracle +//! Path A: gqa_forward_with_cache (materialised backward forward) +//! Path B: oracle built from the REAL primitives (replicated inline): +//! - real_rms_norm_shifted (matches qwen35_rms_norm exactly, shifted gamma) +//! - deinterleave_q_gate + apply_sigmoid_gate (same functions from crate) +//! - same RoPE (stride-half convention, matches real model) +//! - same attention/o_proj arithmetic +//! This exposes convention bugs in the materialised forward. +//! +//! TEST 3: Hypothetical shifted-gamma tape variant +//! What max-diff WOULD be if tape.rs rms_norm_forward used shifted gamma, +//! computed inline in the test without modifying production code. + +#[cfg(feature = "train-backward")] +mod parity_tests { + use lattice_inference::attention::gated::{apply_sigmoid_gate, deinterleave_q_gate}; + use lattice_inference::backward::attention_gqa::gqa_forward_with_cache; + use lattice_inference::backward::tape::rms_norm_forward; + + // =================================================================== + // Inline oracle: exact copy of qwen35_rms_norm (pub(crate) so not + // importable from integration tests — but replicating 10 lines is + // correct and avoids a visibility-change to production code). + // + // Source: crates/inference/src/model/qwen35/norm.rs:5-22 + // Convention: v * inv_rms * (1 + g) (shifted gamma) + // =================================================================== + + fn real_rms_norm_shifted(x: &mut [f32], gamma: &[f32], hidden: usize, eps: f32) { + let num_tokens = x.len() / hidden; + for t in 0..num_tokens { + let row = &mut x[t * hidden..(t + 1) * hidden]; + let mut sum_sq = 0.0f32; + for &v in row.iter() { + sum_sq += v * v; + } + let inv_rms = 1.0 / (sum_sq / hidden as f32 + eps).sqrt(); + for (v, &g) in row.iter_mut().zip(gamma.iter()) { + *v = *v * inv_rms * (1.0 + g); + } + } + } + + // =================================================================== + // Utilities + // =================================================================== + + fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 { + assert_eq!( + a.len(), + b.len(), + "max_abs_diff: length mismatch {} vs {}", + a.len(), + b.len() + ); + a.iter() + .zip(b.iter()) + .map(|(x, y)| (x - y).abs()) + .fold(0.0f32, f32::max) + } + + /// xorshift64 PRNG matching existing test conventions in this codebase. + fn rand_vec(state: &mut u64, n: usize, scale: f32) -> Vec { + let mut out = Vec::with_capacity(n); + for _ in 0..n { + let mut x = *state; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + *state = x; + out.push(((x >> 32) as u32 as f32 / u32::MAX as f32 * 2.0 - 1.0) * scale); + } + out + } + + // =================================================================== + // TEST 1: RMSNorm convention diff + // + // tape.rs rms_norm_forward: normed[i] = xi * wi * inv_rms + // (plain gamma: w is multiplied directly) + // + // qwen35_rms_norm (real model): normed[i] = xi * inv_rms * (1 + gi) + // (shifted gamma: adds 1 to gamma before multiplying) + // + // Theoretical diff = real - tape = xi * inv_rms * (1+gi) - xi*gi*inv_rms + // = xi * inv_rms (the "+1" term, independent of gamma) + // =================================================================== + + #[test] + fn test1_rms_norm_convention_diff() { + let hidden = 64usize; + let eps = 1e-6f32; + let mut rng = 0xCAFE_BABE_u64; + + // Nonzero weights (important: w=0 would collapse to the same when + // gamma_shifted=0 → (1+0)=1 → same as plain gamma with w=1, so we + // need w != 0 and w != 1 to clearly see the convention divergence) + let x = rand_vec(&mut rng, hidden, 1.0); + let w = rand_vec(&mut rng, hidden, 0.5); // gamma values in [-0.5, 0.5] + + // Path A: tape.rs rms_norm_forward — plain gamma xi*wi*inv_rms + let (tape_out, tape_inv_rms) = rms_norm_forward(&x, &w, eps); + + // Path B: real shifted gamma xi*(1+wi)*inv_rms + let mut real_out = x.clone(); + real_rms_norm_shifted(&mut real_out, &w, hidden, eps); + + let diff = max_abs_diff(&tape_out, &real_out); + + let mean_sq: f32 = x.iter().map(|xi| xi * xi).sum::() / hidden as f32; + let expected_inv_rms = 1.0 / (mean_sq + eps).sqrt(); + + // Analytical diff = xi * inv_rms (the constant +1 term) + let theoretical_diff: Vec = x.iter().map(|xi| xi * tape_inv_rms).collect(); + let actual_diff: Vec = real_out + .iter() + .zip(tape_out.iter()) + .map(|(r, t)| r - t) + .collect(); + let meta_diff = max_abs_diff(&theoretical_diff, &actual_diff); + + println!("=== TEST 1: RMSNorm convention diff ==="); + println!(" hidden={hidden}, eps={eps}"); + println!(" inv_rms (tape): {tape_inv_rms:.8}"); + println!(" inv_rms (expected): {expected_inv_rms:.8}"); + println!(" Sample w[0..4]: {:?}", &w[..4]); + println!(" Sample x[0..4]: {:?}", &x[..4]); + println!(" tape_out[0..4]: {:?}", &tape_out[..4]); + println!(" real_out[0..4]: {:?}", &real_out[..4]); + println!(" MEASURED max-diff(tape vs real): {diff:.6e}"); + println!(" Theoretical diff = xi*inv_rms (the +1 term):"); + println!(" theoretical_diff[0..4]: {:?}", &theoretical_diff[..4]); + println!(" actual_diff[0..4]: {:?}", &actual_diff[..4]); + println!(" max-diff(theoretical vs actual diff): {meta_diff:.6e}"); + println!(); + } + + // =================================================================== + // TEST 2: GQA single-layer forward parity — materialised vs oracle + // + // Structural fixture mirrors Qwen3.5-0.8B layer 23 (the last full- + // attention layer). We use small dims to keep the test fast: + // hidden=64, num_q_heads=4, num_kv_heads=2, head_dim=16, rank=8 + // + // Both paths receive the same weights, same LoRA, same input x. + // + // PATH A: gqa_forward_with_cache (materialised backward forward) + // This function uses SHIFTED q_norm/k_norm (1+gamma) as of the + // fix in attention_gqa.rs:445-466. So the q/k-norm convention is + // already correct. Any remaining divergence comes from other places. + // + // PATH B: Oracle from real primitives (NOT driving the loaded model) + // - real_rms_norm_shifted (exact qwen35_rms_norm replica) for q/k norm + // - same deinterleave_q_gate, apply_sigmoid_gate from crate + // - same stride-half RoPE + // - same causal softmax attention (seq_len=1 → trivial: single pos) + // - same o_proj matmul + // + // The oracle uses seq_len=1 for simplicity (position 0 only). The + // real model's decode step also processes one token at a time. + // =================================================================== + + #[test] + fn test2_gqa_layer_forward_parity_materialised_vs_oracle() { + let seq_len = 1usize; + let hidden = 64usize; + let num_q_heads = 4usize; + let num_kv_heads = 2usize; + let head_dim = 16usize; + let rope_dim = 16usize; // full rope on all head dims + let lora_rank = 8usize; + let lora_scale = 1.0f32; + + let q_dim = num_q_heads * head_dim; + let kv_dim = num_kv_heads * head_dim; + let half = rope_dim / 2; + let eps = 1e-6f32; + + let mut rng = 0xDEAD_C0DE_u64; + + let w_q = rand_vec(&mut rng, 2 * q_dim * hidden, 0.1); + let w_k = rand_vec(&mut rng, kv_dim * hidden, 0.1); + let w_v = rand_vec(&mut rng, kv_dim * hidden, 0.1); + let w_o = rand_vec(&mut rng, hidden * q_dim, 0.1); + let q_norm_w = rand_vec(&mut rng, head_dim, 0.3); // nonzero gamma + let k_norm_w = rand_vec(&mut rng, head_dim, 0.3); + let x = rand_vec(&mut rng, seq_len * hidden, 0.5); + + // LoRA (nonzero B so delta is nontrivial) + let lora_a_q = rand_vec(&mut rng, lora_rank * hidden, 0.05); + let lora_b_q = rand_vec(&mut rng, 2 * q_dim * lora_rank, 0.05); + let lora_a_v = rand_vec(&mut rng, lora_rank * hidden, 0.05); + let lora_b_v = rand_vec(&mut rng, kv_dim * lora_rank, 0.05); + + // RoPE tables (position=0) + let cos_table: Vec = (0..seq_len * half) + .map(|i| { + let pos = i / half; + let dim_i = i % half; + let theta = (pos as f32) / 10000f32.powf(2.0 * dim_i as f32 / rope_dim as f32); + theta.cos() + }) + .collect(); + let sin_table: Vec = (0..seq_len * half) + .map(|i| { + let pos = i / half; + let dim_i = i % half; + let theta = (pos as f32) / 10000f32.powf(2.0 * dim_i as f32 / rope_dim as f32); + theta.sin() + }) + .collect(); + + // --------------------------------------------------------------- + // PATH A: materialised forward + // --------------------------------------------------------------- + let (path_a_out, _cache_a) = gqa_forward_with_cache( + &x, + &w_q, + &w_k, + &w_v, + &w_o, + &q_norm_w, + &k_norm_w, + Some(&lora_a_q), + Some(&lora_b_q), + Some(&lora_a_v), + Some(&lora_b_v), + lora_rank, + lora_scale, + seq_len, + hidden, + num_q_heads, + num_kv_heads, + head_dim, + rope_dim, + &cos_table, + &sin_table, + eps, + ); + + // --------------------------------------------------------------- + // PATH B: Oracle from real Qwen3.5 primitives (replicated inline) + // + // Mirrors full_attention_step_from_attn_out + project_qkv from + // crates/inference/src/model/qwen35/forward.rs:170-246 + // --------------------------------------------------------------- + let _scale_attn = 1.0f32 / (head_dim as f32).sqrt(); + let groups = num_q_heads / num_kv_heads; + + let mut oracle_out = vec![0.0f32; seq_len * hidden]; + + // Capture intermediates for localisation + let mut oracle_q_normed = vec![0.0f32; q_dim]; + let mut oracle_k_normed = vec![0.0f32; kv_dim]; + + for t in 0..seq_len { + let x_t = &x[t * hidden..(t + 1) * hidden]; + + // 1. q+gate projection: w_q [2*q_dim, hidden], row-major + let mut q_and_gate = vec![0.0f32; 2 * q_dim]; + for i in 0..2 * q_dim { + let row = &w_q[i * hidden..(i + 1) * hidden]; + q_and_gate[i] = row.iter().zip(x_t.iter()).map(|(a, b)| a * b).sum(); + } + // LoRA Q delta + { + let mut h_q = vec![0.0f32; lora_rank]; + for r in 0..lora_rank { + h_q[r] = lora_a_q[r * hidden..(r + 1) * hidden] + .iter() + .zip(x_t.iter()) + .map(|(a, b)| a * b) + .sum(); + } + for i in 0..2 * q_dim { + let acc: f32 = lora_scale + * lora_b_q[i * lora_rank..(i + 1) * lora_rank] + .iter() + .zip(h_q.iter()) + .map(|(b, hi)| b * hi) + .sum::(); + q_and_gate[i] += acc; + } + } + + // 2. deinterleave: per-head [Q_h | gate_h] layout + let mut q_buf = vec![0.0f32; q_dim]; + let mut gate_z = vec![0.0f32; q_dim]; + deinterleave_q_gate(&q_and_gate, &mut q_buf, &mut gate_z, num_q_heads, head_dim); + + // 3. k_proj + let mut k_buf = vec![0.0f32; kv_dim]; + for i in 0..kv_dim { + let row = &w_k[i * hidden..(i + 1) * hidden]; + k_buf[i] = row.iter().zip(x_t.iter()).map(|(a, b)| a * b).sum(); + } + + // 4. v_proj + LoRA V + let mut v_buf = vec![0.0f32; kv_dim]; + for i in 0..kv_dim { + let row = &w_v[i * hidden..(i + 1) * hidden]; + v_buf[i] = row.iter().zip(x_t.iter()).map(|(a, b)| a * b).sum(); + } + { + let mut h_v = vec![0.0f32; lora_rank]; + for r in 0..lora_rank { + h_v[r] = lora_a_v[r * hidden..(r + 1) * hidden] + .iter() + .zip(x_t.iter()) + .map(|(a, b)| a * b) + .sum(); + } + for i in 0..kv_dim { + let acc: f32 = lora_scale + * lora_b_v[i * lora_rank..(i + 1) * lora_rank] + .iter() + .zip(h_v.iter()) + .map(|(b, hi)| b * hi) + .sum::(); + v_buf[i] += acc; + } + } + + // 5. q_norm: real shifted gamma per head + for qh in 0..num_q_heads { + let start = qh * head_dim; + real_rms_norm_shifted( + &mut q_buf[start..start + head_dim], + &q_norm_w, + head_dim, + eps, + ); + } + + // 6. k_norm: real shifted gamma per head + for kvh in 0..num_kv_heads { + let start = kvh * head_dim; + real_rms_norm_shifted( + &mut k_buf[start..start + head_dim], + &k_norm_w, + head_dim, + eps, + ); + } + + // Capture for localisation + if t == 0 { + oracle_q_normed.copy_from_slice(&q_buf); + oracle_k_normed.copy_from_slice(&k_buf); + } + + // 7. RoPE stride-half (matches gqa_forward_with_cache:481-506) + let cos_t = &cos_table[t * half..(t + 1) * half]; + let sin_t = &sin_table[t * half..(t + 1) * half]; + for qh in 0..num_q_heads { + let start = qh * head_dim; + for i in 0..half { + let c = cos_t[i]; + let s = sin_t[i]; + let x0 = q_buf[start + i]; + let x1 = q_buf[start + half + i]; + q_buf[start + i] = x0 * c - x1 * s; + q_buf[start + half + i] = x0 * s + x1 * c; + } + } + for kvh in 0..num_kv_heads { + let start = kvh * head_dim; + let x0_vals: Vec = (0..half).map(|i| k_buf[start + i]).collect(); + let x1_vals: Vec = (0..half).map(|i| k_buf[start + half + i]).collect(); + for i in 0..half { + let c = cos_t[i]; + let s = sin_t[i]; + k_buf[start + i] = x0_vals[i] * c - x1_vals[i] * s; + k_buf[start + half + i] = x0_vals[i] * s + x1_vals[i] * c; + } + } + + // 8. Causal attention (seq_len=1: position 0 can only attend to 0) + // Softmax over 1 element = 1.0, context = v[0] + let mut context = vec![0.0f32; q_dim]; + for qh in 0..num_q_heads { + let kvh = qh / groups; + let q_off = qh * head_dim; + let kv_off = kvh * head_dim; + let q = &q_buf[q_off..q_off + head_dim]; + let k = &k_buf[kv_off..kv_off + head_dim]; + let v = &v_buf[kv_off..kv_off + head_dim]; + + // Single-element softmax: attention score = exp(dot*scale)/exp(dot*scale) = 1.0 + // so context = 1.0 * v regardless of the score value + let _dot: f32 = q.iter().zip(k.iter()).map(|(qi, ki)| qi * ki).sum(); + let ctx_off = qh * head_dim; + context[ctx_off..ctx_off + head_dim].copy_from_slice(&v[..head_dim]); + } + + // 9. Sigmoid gate: context[i] *= sigmoid(gate_z[i]) + apply_sigmoid_gate(&mut context, &gate_z); + + // 10. o_proj: w_o [hidden, q_dim] + let out_t = &mut oracle_out[t * hidden..(t + 1) * hidden]; + for i in 0..hidden { + let row = &w_o[i * q_dim..(i + 1) * q_dim]; + out_t[i] = row.iter().zip(context.iter()).map(|(a, b)| a * b).sum(); + } + } + + // --------------------------------------------------------------- + // Compare overall + // --------------------------------------------------------------- + let max_diff_overall = max_abs_diff(&path_a_out, &oracle_out); + + // --------------------------------------------------------------- + // Localise: capture materialised q_norm for comparison + // Re-derive the materialised q_norm output from the same weights. + // attention_gqa.rs:448-456 uses the SHIFTED form (1+gamma)*inv_rms + // --------------------------------------------------------------- + let x_t = &x[..hidden]; + let mut q_and_gate_mat = vec![0.0f32; 2 * q_dim]; + for i in 0..2 * q_dim { + let row = &w_q[i * hidden..(i + 1) * hidden]; + q_and_gate_mat[i] = row.iter().zip(x_t.iter()).map(|(a, b)| a * b).sum(); + } + { + let mut h_q = vec![0.0f32; lora_rank]; + for r in 0..lora_rank { + h_q[r] = lora_a_q[r * hidden..(r + 1) * hidden] + .iter() + .zip(x_t.iter()) + .map(|(a, b)| a * b) + .sum(); + } + for i in 0..2 * q_dim { + let acc: f32 = lora_scale + * lora_b_q[i * lora_rank..(i + 1) * lora_rank] + .iter() + .zip(h_q.iter()) + .map(|(b, hi)| b * hi) + .sum::(); + q_and_gate_mat[i] += acc; + } + } + let mut q_mat_raw = vec![0.0f32; q_dim]; + let mut gate_z_mat = vec![0.0f32; q_dim]; + deinterleave_q_gate( + &q_and_gate_mat, + &mut q_mat_raw, + &mut gate_z_mat, + num_q_heads, + head_dim, + ); + + // Materialised q_norm: SHIFTED (1+gamma)*inv_rms per head + // (matches attention_gqa.rs:448-456 which was fixed to use shifted gamma) + let mut q_mat_normed = q_mat_raw.clone(); + for qh in 0..num_q_heads { + let start = qh * head_dim; + let q_head = &mut q_mat_normed[start..start + head_dim]; + let mean_sq: f32 = q_head.iter().map(|xi| xi * xi).sum::() / head_dim as f32; + let inv_rms = 1.0 / (mean_sq + eps).sqrt(); + for (j, qj) in q_head.iter_mut().enumerate() { + *qj *= (1.0 + q_norm_w[j]) * inv_rms; // SHIFTED — matches attention_gqa.rs + } + } + + let qnorm_diff = max_abs_diff(&q_mat_normed, &oracle_q_normed); + + // Pre-attn RMSNorm check (tape.rs rms_norm_forward vs real) + // This is what tape.rs uses for the pre-attention norm in the full training path. + let w_all_ones = vec![1.0f32; hidden]; + let (tape_normed_ones, _) = rms_norm_forward(x_t, &w_all_ones, eps); + let mut real_normed_ones = x_t.to_vec(); + real_rms_norm_shifted(&mut real_normed_ones, &w_all_ones, hidden, eps); + let pre_attn_norm_diff_ones = max_abs_diff(&tape_normed_ones, &real_normed_ones); + + // With nonzero gamma + let mut rng2 = 0xABCD_1234_u64; + let w_nonzero = rand_vec(&mut rng2, hidden, 0.3); + let (tape_normed_nonzero, _) = rms_norm_forward(x_t, &w_nonzero, eps); + let mut real_normed_nonzero = x_t.to_vec(); + real_rms_norm_shifted(&mut real_normed_nonzero, &w_nonzero, hidden, eps); + let pre_attn_norm_diff_nonzero = max_abs_diff(&tape_normed_nonzero, &real_normed_nonzero); + + println!("=== TEST 2: GQA layer forward parity (materialised vs oracle) ==="); + println!(" Approach: Path B = Oracle from real Qwen3.5 primitives (replicated inline)"); + println!( + " Dims: seq_len={seq_len} hidden={hidden} q_heads={num_q_heads} kv_heads={num_kv_heads}" + ); + println!(" head_dim={head_dim} rope_dim={rope_dim} lora_rank={lora_rank}"); + println!(" path_a_out [0..4]: {:?}", &path_a_out[..4.min(hidden)]); + println!(" oracle_out [0..4]: {:?}", &oracle_out[..4.min(hidden)]); + println!(); + println!(" --- Overall ---"); + println!(" MEASURED max-diff(materialised vs oracle): {max_diff_overall:.6e}"); + println!(); + println!(" --- Localisation breakdown ---"); + println!(" post q_norm max-diff (materialised vs oracle): {qnorm_diff:.6e}"); + println!( + " pre-attn RMSNorm diff (tape vs real, w=all-ones): {pre_attn_norm_diff_ones:.6e}" + ); + println!( + " pre-attn RMSNorm diff (tape vs real, w=nonzero): {pre_attn_norm_diff_nonzero:.6e}" + ); + println!(); + println!(" Interpretation:"); + if max_diff_overall < 1e-5 { + println!(" [PASS] Overall diff < 1e-5: materialised forward matches oracle."); + } else if max_diff_overall < 1e-3 { + println!( + " [WARN] Overall diff in (1e-5, 1e-3): small divergence, not on-par threshold." + ); + } else { + println!(" [FAIL] Overall diff >= 1e-3: materialised forward does NOT match oracle."); + } + } + + // =================================================================== + // TEST 3: Hypothetical shifted-gamma tape variant + // + // What would the max-diff be if tape.rs rms_norm_forward used shifted + // gamma xi * (1+wi) * inv_rms instead of xi * wi * inv_rms? + // Computed INLINE IN THIS TEST — production tape.rs is NOT modified. + // =================================================================== + + #[test] + fn test3_hypothetical_shifted_gamma_tape() { + let hidden = 64usize; + let eps = 1e-6f32; + let mut rng = 0xF00D_CAFE_u64; + + let x = rand_vec(&mut rng, hidden, 1.0); + let w = rand_vec(&mut rng, hidden, 0.5); + + // Real (baseline truth): shifted gamma xi*(1+wi)*inv_rms + let mut real_out = x.clone(); + real_rms_norm_shifted(&mut real_out, &w, hidden, eps); + + // Current tape.rs: plain gamma xi*wi*inv_rms + let (current_tape_out, _) = rms_norm_forward(&x, &w, eps); + + // Hypothetical shifted tape: xi*(1+wi)*inv_rms — computed inline + let mean_sq: f32 = x.iter().map(|xi| xi * xi).sum::() / hidden as f32; + let inv_rms = 1.0 / (mean_sq + eps).sqrt(); + let shifted_tape_out: Vec = x + .iter() + .zip(w.iter()) + .map(|(xi, wi)| xi * (1.0 + wi) * inv_rms) + .collect(); + + let diff_current = max_abs_diff(¤t_tape_out, &real_out); + let diff_shifted = max_abs_diff(&shifted_tape_out, &real_out); + + println!("=== TEST 3: Hypothetical shifted-gamma tape ==="); + println!(" hidden={hidden}, eps={eps}"); + println!(" MEASURED max-diff(current tape vs real): {diff_current:.6e}"); + println!(" MEASURED max-diff(shifted tape vs real): {diff_shifted:.6e}"); + if diff_shifted > 0.0 { + println!(" Reduction from fix: {:.1}x", diff_current / diff_shifted); + } else { + println!(" Reduction from fix: inf (shifted tape IS the real formula)"); + } + println!(" NOTE: production tape.rs was NOT modified; this is a what-if measurement."); + println!( + " Verification: shifted_tape[0..4]: {:?}", + &shifted_tape_out[..4] + ); + println!(" Verification: real_out[0..4]: {:?}", &real_out[..4]); + println!(); + } + + // =================================================================== + // TEST 4: Real-model LoRA forward parity + // + // Proves that the materialised GQA forward (gqa_forward_with_cache) + // matches the REAL loaded Qwen3.5-0.8B model at layer 23, WITH an + // identical nonzero LoRA on q_proj and v_proj injected into both paths. + // + // This closes the "self-consistent gradcheck" blind spot: prior + // gradchecks compared the materialised forward to itself; this + // compares it to the actual loaded model running real weights. + // + // Design: + // 1. Load the model from disk (skip gracefully if absent). + // 2. Build a nonzero LoRA for layer 23 (q_proj + v_proj only). + // 3. Inject LoRA into the real forward via a test-local LoraHook. + // 4. Run one token at position 0 through the real model, capturing + // h_in (pre-input-layernorm residual) and attn_out (gated o_proj). + // 5. Apply input RMSNorm to h_in using the real layer-23 weights, + // then call gqa_forward_with_cache with the same LoRA and identity + // RoPE tables (position 0: cos=1, sin=0). + // 6. Assert max-abs-diff < 1e-3. + // + // Skip condition: model dir absent → print SKIP and return without fail. + // Feature gate: #[cfg(all(feature = "train-backward", feature = "f16"))] + // (f16 enables BF16 dequant in from_safetensors). + // =================================================================== + + #[cfg(feature = "f16")] + #[test] + fn test4_real_model_lora_forward_parity_layer23() { + use lattice_inference::lora_hook::LoraHook; + use lattice_inference::model::qwen35::Qwen35Model; + + // ------------------------------------------------------------------ + // 1. Resolve model directory (env-var override with HF cache fallback) + // ------------------------------------------------------------------ + let model_dir = std::env::var("LATTICE_QWEN35_MODEL_DIR") + .map(std::path::PathBuf::from) + .unwrap_or_else(|_| { + let home = std::env::var("HOME").unwrap_or_else(|_| "/root".to_string()); + std::path::PathBuf::from(home) + .join(".cache/huggingface/hub") + .join("models--Qwen--Qwen3.5-0.8B") + .join("snapshots") + .join("2fc06364715b967f1860aea9cf38778875588b17") + }); + + if !model_dir.is_dir() { + println!( + "SKIP: model not found at {} — set LATTICE_QWEN35_MODEL_DIR to run", + model_dir.display() + ); + return; + } + + // ------------------------------------------------------------------ + // 2. Load model and read real dims/eps from config + // ------------------------------------------------------------------ + let mut model = + Qwen35Model::from_safetensors(&model_dir).expect("failed to load Qwen3.5-0.8B"); + + let cfg = model.config(); + let hidden = cfg.hidden_size; + let num_q_heads = cfg.num_attention_heads; + let num_kv_heads = cfg.num_key_value_heads; + let head_dim = cfg.head_dim; + let rope_dim = cfg.rope_dim(); + let eps = cfg.rms_norm_eps; + let q_dim = num_q_heads * head_dim; + let kv_dim = num_kv_heads * head_dim; + + println!( + "=== TEST 4: real-model LoRA forward parity (materialised vs loaded Qwen3.5 layer-23) ===" + ); + println!( + " dims: hidden={hidden} q_heads={num_q_heads} kv_heads={num_kv_heads} head_dim={head_dim}" + ); + println!(" q_dim={q_dim} kv_dim={kv_dim} rope_dim={rope_dim}"); + println!(" rms_norm_eps={eps:.2e} (from model config — not hardcoded)"); + + // ------------------------------------------------------------------ + // 3. Build nonzero LoRA for layer 23 only: q_proj + v_proj + // + // q_proj LoRA shape: A=[rank, hidden], B=[2*q_dim, rank] (q_proj + // output is 2*q_dim due to per-head interleaved [Q|gate] layout) + // v_proj LoRA shape: A=[rank, hidden], B=[kv_dim, rank] + // ------------------------------------------------------------------ + let lora_rank = 4usize; + let lora_scale = 0.03f32; + let mut rng = 0x1234_ABCD_5678_EF01_u64; + + let lora_a_q = rand_vec(&mut rng, lora_rank * hidden, 0.03); + let lora_b_q = rand_vec(&mut rng, 2 * q_dim * lora_rank, 0.03); + let lora_a_v = rand_vec(&mut rng, lora_rank * hidden, 0.03); + let lora_b_v = rand_vec(&mut rng, kv_dim * lora_rank, 0.03); + + // ------------------------------------------------------------------ + // 4. Test-local LoraHook: injects LoRA for (layer=23, q_proj) + // and (layer=23, v_proj); no-op for everything else. + // + // Convention matches gqa_forward_with_cache and production + // forward.rs: output += scale * B @ (A @ x), row-major, A then B. + // ------------------------------------------------------------------ + struct TestLoraHook { + layer: usize, + rank: usize, + scale: f32, + a_q: Vec, + b_q: Vec, + a_v: Vec, + b_v: Vec, + } + + impl LoraHook for TestLoraHook { + fn apply(&self, layer_idx: usize, module: &str, x: &[f32], output: &mut [f32]) { + if layer_idx != self.layer { + return; + } + let (a, b, out_dim) = match module { + "q_proj" => (&self.a_q, &self.b_q, output.len()), + "v_proj" => (&self.a_v, &self.b_v, output.len()), + _ => return, + }; + let rank = self.rank; + let in_dim = x.len(); + // h = A @ x [rank] + let mut h = vec![0.0f32; rank]; + for r in 0..rank { + h[r] = a[r * in_dim..(r + 1) * in_dim] + .iter() + .zip(x.iter()) + .map(|(ai, xi)| ai * xi) + .sum(); + } + // output += scale * B @ h [out_dim] + for i in 0..out_dim { + let acc: f32 = b[i * rank..(i + 1) * rank] + .iter() + .zip(h.iter()) + .map(|(bi, hi)| bi * hi) + .sum(); + output[i] += self.scale * acc; + } + } + } + + model.set_lora(Box::new(TestLoraHook { + layer: 23, + rank: lora_rank, + scale: lora_scale, + a_q: lora_a_q.clone(), + b_q: lora_b_q.clone(), + a_v: lora_a_v.clone(), + b_v: lora_b_v.clone(), + })); + + // ------------------------------------------------------------------ + // 5. Run real forward for a single token at position 0. + // capture_attn_io runs forward_step for each token in sequence; + // with a single token [100], position=0, RoPE is identity (cos=1, + // sin=0 for all dims since theta*0=0). + // ------------------------------------------------------------------ + let tokens: Vec = vec![100]; + let (h_in, captured_attn_out) = model + .capture_attn_io(&tokens, 23) + .expect("capture_attn_io failed for layer 23"); + + assert_eq!( + h_in.len(), + hidden, + "h_in length mismatch: got {} expected {hidden}", + h_in.len() + ); + assert_eq!( + captured_attn_out.len(), + hidden, + "captured_attn_out length mismatch: got {} expected {hidden}", + captured_attn_out.len() + ); + + // ------------------------------------------------------------------ + // 6. Get real layer-23 weights + // ------------------------------------------------------------------ + let (w_q, w_k, w_v, w_o, q_norm_w, k_norm_w, pre_attn_norm, _post, _g, _u, _d) = model + .gqa_layer_weights(23) + .expect("layer 23 is a Full+Dense GQA layer"); + + // ------------------------------------------------------------------ + // 7. Recompute input-layernorm: normed = rms_norm(h_in, pre_attn_norm) + // using real shifted gamma (1+gamma)*inv_rms, matching qwen35_rms_norm. + // The materialised forward expects already-normed input. + // ------------------------------------------------------------------ + let mut normed_input = h_in.clone(); + real_rms_norm_shifted(&mut normed_input, pre_attn_norm, hidden, eps); + + // ------------------------------------------------------------------ + // 8. Build identity RoPE tables for position 0: cos=1, sin=0. + // rope_dim/2 entries per position; seq_len=1. + // ------------------------------------------------------------------ + let half_rope = rope_dim / 2; + let cos_table = vec![1.0f32; half_rope]; // cos(0) = 1 + let sin_table = vec![0.0f32; half_rope]; // sin(0) = 0 + + // ------------------------------------------------------------------ + // 9. Run materialised forward (gqa_forward_with_cache) + // with the same LoRA arrays and real weights. + // ------------------------------------------------------------------ + let (materialised_out, _cache) = gqa_forward_with_cache( + &normed_input, + w_q, + w_k, + w_v, + w_o, + q_norm_w, + k_norm_w, + Some(&lora_a_q), + Some(&lora_b_q), + Some(&lora_a_v), + Some(&lora_b_v), + lora_rank, + lora_scale, + 1, // seq_len + hidden, + num_q_heads, + num_kv_heads, + head_dim, + rope_dim, + &cos_table, + &sin_table, + eps, + ); + + // ------------------------------------------------------------------ + // 10. Localisation baseline: run materialised forward with zero LoRA + // to isolate the LoRA-injection contribution. + // ------------------------------------------------------------------ + let zero_a_q = vec![0.0f32; lora_rank * hidden]; + let zero_b_q = vec![0.0f32; 2 * q_dim * lora_rank]; + let zero_a_v = vec![0.0f32; lora_rank * hidden]; + let zero_b_v = vec![0.0f32; kv_dim * lora_rank]; + + // Materialised forward with zero LoRA (M0). + let (materialised_no_lora, _) = gqa_forward_with_cache( + &normed_input, + w_q, + w_k, + w_v, + w_o, + q_norm_w, + k_norm_w, + Some(&zero_a_q), + Some(&zero_b_q), + Some(&zero_a_v), + Some(&zero_b_v), + lora_rank, + lora_scale, + 1, + hidden, + num_q_heads, + num_kv_heads, + head_dim, + rope_dim, + &cos_table, + &sin_table, + eps, + ); + + // True no-LoRA real baseline (R0): re-capture the real model with a + // zero-LoRA hook so d_no_lora compares materialised-no-LoRA (M0) vs + // real-no-LoRA (R0), isolating base f32 divergence from the LoRA delta. + // Comparing M0 against the with-LoRA capture (R1) would conflate the + // base gap with the LoRA-delta magnitude (they coincide numerically). + model.set_lora(Box::new(TestLoraHook { + layer: 23, + rank: lora_rank, + scale: lora_scale, + a_q: zero_a_q, + b_q: zero_b_q, + a_v: zero_a_v, + b_v: zero_b_v, + })); + let (_h_in_no_lora, captured_attn_out_no_lora) = model + .capture_attn_io(&tokens, 23) + .expect("capture_attn_io (no-LoRA baseline) failed for layer 23"); + + // ------------------------------------------------------------------ + // 11. Measure and report + // ------------------------------------------------------------------ + let d = max_abs_diff(&materialised_out, &captured_attn_out); + let d_no_lora = max_abs_diff(&materialised_no_lora, &captured_attn_out_no_lora); + let lora_contribution = max_abs_diff(&materialised_out, &materialised_no_lora); + + println!(); + println!( + " Sample materialised_out[0..4]: {:?}", + &materialised_out[..4.min(hidden)] + ); + println!( + " Sample captured_attn_out[0..4]: {:?}", + &captured_attn_out[..4.min(hidden)] + ); + println!(); + println!(" Localisation:"); + println!(" max-diff(materialised vs real), NO LoRA: {d_no_lora:.6e}"); + println!(" LoRA delta contribution (mat_lora - mat_no_lora): {lora_contribution:.6e}"); + println!(" MEASURED max-diff(materialised vs real model): {d:.6e}"); + + if d < 1e-3 { + println!( + " [PASS] max-diff < 1e-3: materialised GQA forward matches real model with LoRA." + ); + } else { + println!(" [FAIL] max-diff >= 1e-3: divergence detected."); + println!(" Diagnosis hints:"); + if d_no_lora < 1e-3 && d >= 1e-3 { + println!(" LoRA injection mismatch (no-LoRA passes, LoRA fails)."); + } else if d_no_lora >= 1e-3 { + println!( + " Base forward divergence (no-LoRA fails); check norm eps, o_proj LoRA." + ); + } + } + println!(); + + assert!( + d < 1e-3, + "TEST 4 FAILED: max-abs-diff(materialised vs real-model) = {d:.6e} >= 1e-3. \ + No-LoRA baseline diff = {d_no_lora:.6e}. LoRA contribution = {lora_contribution:.6e}." + ); + } +} diff --git a/crates/tune/src/bin/train_grad_full.rs b/crates/tune/src/bin/train_grad_full.rs index e5d957eca4..7834432b44 100644 --- a/crates/tune/src/bin/train_grad_full.rs +++ b/crates/tune/src/bin/train_grad_full.rs @@ -26,7 +26,7 @@ // gradchecks cannot see. // // Qwen3.5 RMSNorm is SHIFTED (x·inv·(1+gamma)); rms_norm_forward/rmsnorm_backward -// take plain gamma, so layer/final norms get (1+gamma) precomputed weights. +// apply (1+gamma) internally, so layer/final norm fields carry raw gamma. // q_norm/k_norm are shifted inside gqa_forward_with_cache, so they stay raw. // // Usage: train_grad_full --model-dir --data-dir [--first-layer 19] @@ -38,14 +38,16 @@ use std::path::{Path, PathBuf}; use std::time::Instant; use lattice_inference::attention::gdn::GatedDeltaNetWeights; -use lattice_inference::attention::gdn_backward::{GdnSaved, gdn_backward, gdn_forward_save}; +use lattice_inference::attention::gdn_backward::{ + GdnGrads, GdnSaved, gdn_backward, gdn_forward_save, +}; use lattice_inference::backward::attention_gqa::{AttnCache, gqa_backward, gqa_forward_with_cache}; use lattice_inference::backward::ops::{linear_vjp, rmsnorm_backward, swiglu_backward}; use lattice_inference::backward::tape::{rms_norm_forward, swiglu_forward}; use lattice_inference::model::qwen35::Qwen35Model; use lattice_inference::model::qwen35_config::Qwen35Config; use lattice_inference::tokenizer::Tokenizer; -use lattice_tune::lora::AdamState; +use lattice_tune::lora::{AdamState, LoraAdapter, LoraConfig, LoraLayer}; const TOP_LAYER: usize = 23; @@ -78,6 +80,8 @@ Options: --log-every Print NLL every N steps (default: 5) --gradcheck Run finite-difference gradcheck instead of training --probe Gradcheck entries probed per array per layer (default: 6) + --json Emit @@lattice JSON events to stdout alongside human output + --save After training, write a PEFT safetensors adapter to PATH -h, --help Print this help" ); } @@ -212,15 +216,17 @@ impl GdnDims { } } +#[derive(Clone, Copy, PartialEq, Eq)] enum MixerKind { Gqa, Gdn, } /// Borrowed frozen weights for one materialised layer (lifetime tied to &model). -/// Norm `*_shift` fields are the `(1 + gamma)` shifted layer norms; `q_norm`/ -/// `k_norm` are raw (gqa shifts internally). `lora_slot` indexes the mutable -/// LoRA param array for GQA layers; `None` for frozen GDN layers. +/// Norm `pre_norm`/`post_norm`/`final_norm` fields carry raw gamma; the shifted +/// primitives (rms_norm_forward/rmsnorm_backward) apply (1+gamma) internally. +/// `q_norm`/`k_norm` are raw (gqa shifts internally). `lora_slot` indexes the +/// mutable LoRA param array for GQA layers; `None` for frozen GDN layers. struct LayerW<'a> { kind: MixerKind, // GQA mixer (valid iff kind == Gqa) @@ -233,8 +239,8 @@ struct LayerW<'a> { // GDN mixer (valid iff kind == Gdn) gdn: Option<&'a GatedDeltaNetWeights>, // common - pre_shift: Vec, - post_shift: Vec, + pre_norm: Vec, + post_norm: Vec, w_gate: &'a [f32], w_up: &'a [f32], w_down: &'a [f32], @@ -243,30 +249,80 @@ struct LayerW<'a> { struct Head<'a> { lm_head: &'a [f32], - final_shift: &'a [f32], + final_norm: &'a [f32], } -/// Mutable LoRA factors for one GQA layer (q_proj + v_proj). +/// Mutable LoRA factors for one layer. +/// +/// GQA layers populate a_q/b_q/a_v/b_v; GDN fields are empty Vec. +/// GDN layers populate a_qkv/b_qkv/a_z/b_z/a_b/b_b/a_a/b_a/a_out/b_out; +/// GQA fields are empty Vec. #[derive(Clone)] struct LoraParams { + // GQA fields (empty for GDN slots) a_q: Vec, b_q: Vec, a_v: Vec, b_v: Vec, + // GDN fields (empty for GQA slots) + a_qkv: Vec, + b_qkv: Vec, + a_z: Vec, + b_z: Vec, + a_b: Vec, + b_b: Vec, + a_a: Vec, + b_a: Vec, + a_out: Vec, + b_out: Vec, } impl LoraParams { - fn zeros(rank: usize, d: &Dims) -> Self { + fn zeros_gqa(rank: usize, d: &Dims) -> Self { Self { a_q: vec![0.0; rank * d.hidden], b_q: vec![0.0; 2 * d.q_dim * rank], a_v: vec![0.0; rank * d.hidden], b_v: vec![0.0; d.kv_dim * rank], + a_qkv: Vec::new(), + b_qkv: Vec::new(), + a_z: Vec::new(), + b_z: Vec::new(), + a_b: Vec::new(), + b_b: Vec::new(), + a_a: Vec::new(), + b_a: Vec::new(), + a_out: Vec::new(), + b_out: Vec::new(), + } + } + + fn zeros_gdn(rank: usize, d: &Dims, gd: &GdnDims) -> Self { + Self { + a_q: Vec::new(), + b_q: Vec::new(), + a_v: Vec::new(), + b_v: Vec::new(), + a_qkv: vec![0.0; rank * d.hidden], + b_qkv: vec![0.0; gd.qkv_dim * rank], + a_z: vec![0.0; rank * d.hidden], + b_z: vec![0.0; gd.output_dim * rank], + a_b: vec![0.0; rank * d.hidden], + b_b: vec![0.0; gd.num_kh * rank], + a_a: vec![0.0; rank * d.hidden], + b_a: vec![0.0; gd.num_kh * rank], + a_out: vec![0.0; rank * gd.output_dim], + b_out: vec![0.0; d.hidden * rank], } } + + /// Backward-compatibility alias used in TBV and eval paths. + fn zeros(rank: usize, d: &Dims) -> Self { + Self::zeros_gqa(rank, d) + } } -/// LoRA gradients for one GQA layer (same shapes as LoraParams). +/// LoRA gradients for one layer (same shapes as LoraParams). type Grads = LoraParams; /// One sample's frozen context entering `first_layer`. @@ -279,6 +335,7 @@ struct SeqCtx { seq_len: usize, } +#[allow(clippy::large_enum_variant)] enum MixerCache { Gqa(AttnCache), Gdn(GdnSaved), @@ -358,7 +415,7 @@ fn forward_full( for lw in layers { let h_layer_in = h.clone(); - let (normed_pre, inv_pre) = rmsnorm_seq(&h, &lw.pre_shift, hidden, seq, d.eps); + let (normed_pre, inv_pre) = rmsnorm_seq(&h, &lw.pre_norm, hidden, seq, d.eps); let (mixer_out, mixer_cache) = match lw.kind { MixerKind::Gqa => { @@ -404,12 +461,25 @@ fn forward_full( d.eps, ); let mut out = vec![0.0f32; seq * hidden]; + let lora = &loras[lw.lora_slot.expect("GDN layer must have a LoRA slot")]; gdn_forward_save( &normed_pre, lw.gdn.expect("GDN layer must have GDN weights"), cfg, &mut saved, &mut out, + Some(&lora.a_qkv), + Some(&lora.b_qkv), + Some(&lora.a_z), + Some(&lora.b_z), + Some(&lora.a_b), + Some(&lora.b_b), + Some(&lora.a_a), + Some(&lora.b_a), + Some(&lora.a_out), + Some(&lora.b_out), + rank, + scale, ); (out, MixerCache::Gdn(saved)) } @@ -420,7 +490,7 @@ fn forward_full( *a += *b; } - let (normed_ffn, inv_ffn) = rmsnorm_seq(&h_mid, &lw.post_shift, hidden, seq, d.eps); + let (normed_ffn, inv_ffn) = rmsnorm_seq(&h_mid, &lw.post_norm, hidden, seq, d.eps); let mut gate_pre = Vec::with_capacity(seq); let mut up_pre = Vec::with_capacity(seq); let mut h_next = h_mid.clone(); @@ -459,7 +529,7 @@ fn forward_full( let mut positions = Vec::new(); for t in (ctx.completion_start - 1)..seq - 1 { let (final_normed, inv_final) = - rms_norm_forward(&h[t * hidden..(t + 1) * hidden], head.final_shift, d.eps); + rms_norm_forward(&h[t * hidden..(t + 1) * hidden], head.final_norm, d.eps); let logits = lm_head_logits(head.lm_head, &final_normed, hidden, d.vocab); positions.push(HeadPos { t, @@ -516,14 +586,21 @@ fn nll_and_grads( loras: &[LoraParams], head: &Head, d: &Dims, + gdn_dims: &GdnDims, num_slots: usize, + slot_kinds: &[MixerKind], rank: usize, scale: f32, ) -> (f32, usize, Vec) { let hidden = d.hidden; let seq = fwd.h_final.len() / hidden; let n_comp = fwd.positions.len().max(1) as f32; - let mut grads: Vec = (0..num_slots).map(|_| LoraParams::zeros(rank, d)).collect(); + let mut grads: Vec = (0..num_slots) + .map(|s| match slot_kinds[s] { + MixerKind::Gqa => LoraParams::zeros_gqa(rank, d), + MixerKind::Gdn => LoraParams::zeros_gdn(rank, d, gdn_dims), + }) + .collect(); // ---- Head + CE backward ---- let mut d_h = vec![0.0f32; seq * hidden]; @@ -541,7 +618,7 @@ fn nll_and_grads( let d_final = linear_vjp(head.lm_head, &d_logits, hidden, d.vocab); let d_h_t = rmsnorm_backward( &fwd.h_final[p.t * hidden..(p.t + 1) * hidden], - head.final_shift, + head.final_norm, p.inv_final, &d_final, ); @@ -572,7 +649,7 @@ fn nll_and_grads( ); let d_hm = rmsnorm_backward( &lf.h_mid[t * hidden..(t + 1) * hidden], - &lw.post_shift, + &lw.post_norm, lf.inv_ffn[t], &d_normed_ffn, ); @@ -608,18 +685,45 @@ fn nll_and_grads( b_q: g.grad_b_q, a_v: g.grad_a_v, b_v: g.grad_b_v, + // GQA layer has no GDN LoRA factors → empty gradient (mirrors the + // GDN branch, which leaves the GQA fields empty). + a_qkv: Vec::new(), + b_qkv: Vec::new(), + a_z: Vec::new(), + b_z: Vec::new(), + a_b: Vec::new(), + b_b: Vec::new(), + a_a: Vec::new(), + b_a: Vec::new(), + a_out: Vec::new(), + b_out: Vec::new(), }; g.dx } MixerCache::Gdn(saved) => { - let mut dx = vec![0.0f32; seq * hidden]; - gdn_backward( + let slot = lw.lora_slot.expect("GDN layer must have a LoRA slot"); + let g = gdn_backward( &d_h_mid, saved, lw.gdn.expect("GDN layer must have GDN weights"), - &mut dx, ); - dx + grads[slot] = Grads { + a_q: Vec::new(), + b_q: Vec::new(), + a_v: Vec::new(), + b_v: Vec::new(), + a_qkv: g.grad_a_qkv, + b_qkv: g.grad_b_qkv, + a_z: g.grad_a_z, + b_z: g.grad_b_z, + a_b: g.grad_a_b, + b_b: g.grad_b_b, + a_a: g.grad_a_a, + b_a: g.grad_b_a, + a_out: g.grad_a_out, + b_out: g.grad_b_out, + }; + g.dx } }; @@ -629,7 +733,7 @@ fn nll_and_grads( for t in 0..seq { let d_hl = rmsnorm_backward( &lf.h_layer_in[t * hidden..(t + 1) * hidden], - &lw.pre_shift, + &lw.pre_norm, lf.inv_pre[t], &d_normed_pre[t * hidden..(t + 1) * hidden], ); @@ -643,10 +747,6 @@ fn nll_and_grads( (nll_sum as f32, fwd.positions.len(), grads) } -fn shifted(gamma: &[f32]) -> Vec { - gamma.iter().map(|g| 1.0 + g).collect() -} - /// xorshift small-random fill in [-amp, amp]. fn rand_fill(rng: &mut u64, n: usize, amp: f32) -> Vec { (0..n) @@ -740,6 +840,8 @@ fn main() -> Result<(), Box> { let fd_eps: f32 = parse_arg(&args, "--fd-eps") .and_then(|s| s.parse().ok()) .unwrap_or(4e-3); + let emit_json = parse_flag(&args, "--json"); + let save_path: Option = parse_arg(&args, "--save").map(PathBuf::from); if first_layer > TOP_LAYER { return Err(format!("--first-layer {first_layer} must be <= {TOP_LAYER}").into()); @@ -785,22 +887,24 @@ fn main() -> Result<(), Box> { let (lm_head_s, final_norm_s, _embed) = model.head_weights(); let lm_head = lm_head_s.to_vec(); - let final_shift = shifted(final_norm_s); + let final_norm = final_norm_s.to_vec(); // raw gamma; rms_norm_forward applies (1+gamma) let head = Head { lm_head: &lm_head, - final_shift: &final_shift, + final_norm: &final_norm, }; // Build the materialised layer stack [first_layer ..= 23], assigning a LoRA - // slot to each GQA layer. All weight slices are borrowed from `model`. + // slot to each GQA or GDN layer. All weight slices are borrowed from `model`. let mut layers: Vec = Vec::new(); let mut slot_layers: Vec = Vec::new(); // global layer index per slot + let mut slot_kinds: Vec = Vec::new(); // GQA or GDN per slot for layer_idx in first_layer..=TOP_LAYER { if let Some((w_q, w_k, w_v, w_o, q_norm, k_norm, pre, post, gate, up, down)) = model.gqa_layer_weights(layer_idx) { let slot = slot_layers.len(); slot_layers.push(layer_idx); + slot_kinds.push(MixerKind::Gqa); layers.push(LayerW { kind: MixerKind::Gqa, w_q, @@ -810,14 +914,17 @@ fn main() -> Result<(), Box> { q_norm, k_norm, gdn: None, - pre_shift: shifted(pre), - post_shift: shifted(post), + pre_norm: pre.to_vec(), + post_norm: post.to_vec(), w_gate: gate, w_up: up, w_down: down, lora_slot: Some(slot), }); } else if let Some((gdn, pre, post, gate, up, down)) = model.gdn_layer_weights(layer_idx) { + let slot = slot_layers.len(); + slot_layers.push(layer_idx); + slot_kinds.push(MixerKind::Gdn); layers.push(LayerW { kind: MixerKind::Gdn, w_q: &[], @@ -827,12 +934,12 @@ fn main() -> Result<(), Box> { q_norm: &[], k_norm: &[], gdn: Some(gdn), - pre_shift: shifted(pre), - post_shift: shifted(post), + pre_norm: pre.to_vec(), + post_norm: post.to_vec(), w_gate: gate, w_up: up, w_down: down, - lora_slot: None, + lora_slot: Some(slot), }); } else { return Err( @@ -849,7 +956,7 @@ fn main() -> Result<(), Box> { }) .collect(); println!( - " materialised {} layers [{}]: {kinds} ({} GQA LoRA slots at layers {:?})", + " materialised {} layers [{}]: {kinds} ({} LoRA slots at layers {:?})", layers.len(), (first_layer..=TOP_LAYER) .map(|i| i.to_string()) @@ -859,7 +966,7 @@ fn main() -> Result<(), Box> { slot_layers ); if num_slots == 0 { - return Err("no GQA layers in range — nothing to train".into()); + return Err("no trainable layers in range — nothing to train".into()); } let tokenizer = model.tokenizer().clone(); @@ -929,7 +1036,10 @@ fn main() -> Result<(), Box> { // model's own compute_token_nlls — validates the whole assembled forward // (every layer's shifted norms, GQA/GDN mixer, FFN, head) against the model. let zero_loras: Vec = (0..num_slots) - .map(|_| LoraParams::zeros(rank, &dims)) + .map(|s| match slot_kinds[s] { + MixerKind::Gqa => LoraParams::zeros_gqa(rank, &dims), + MixerKind::Gdn => LoraParams::zeros_gdn(rank, &dims, &gdn_dims), + }) .collect(); { let s0 = &train_samples[0]; @@ -966,44 +1076,122 @@ fn main() -> Result<(), Box> { // Non-zero A AND B so grad_A and the gate path are non-vacuous. let mut rng = 0x1234_5678u64; let mut loras: Vec = (0..num_slots) - .map(|_| LoraParams { - a_q: rand_fill(&mut rng, rank * dims.hidden, 0.05), - b_q: rand_fill(&mut rng, 2 * dims.q_dim * rank, 0.05), - a_v: rand_fill(&mut rng, rank * dims.hidden, 0.05), - b_v: rand_fill(&mut rng, dims.kv_dim * rank, 0.05), + .map(|s| match slot_kinds[s] { + MixerKind::Gqa => LoraParams { + a_q: rand_fill(&mut rng, rank * dims.hidden, 0.05), + b_q: rand_fill(&mut rng, 2 * dims.q_dim * rank, 0.05), + a_v: rand_fill(&mut rng, rank * dims.hidden, 0.05), + b_v: rand_fill(&mut rng, dims.kv_dim * rank, 0.05), + a_qkv: Vec::new(), + b_qkv: Vec::new(), + a_z: Vec::new(), + b_z: Vec::new(), + a_b: Vec::new(), + b_b: Vec::new(), + a_a: Vec::new(), + b_a: Vec::new(), + a_out: Vec::new(), + b_out: Vec::new(), + }, + MixerKind::Gdn => LoraParams { + a_q: Vec::new(), + b_q: Vec::new(), + a_v: Vec::new(), + b_v: Vec::new(), + a_qkv: rand_fill(&mut rng, rank * dims.hidden, 0.05), + b_qkv: rand_fill(&mut rng, gdn_dims.qkv_dim * rank, 0.05), + a_z: rand_fill(&mut rng, rank * dims.hidden, 0.05), + b_z: rand_fill(&mut rng, gdn_dims.output_dim * rank, 0.05), + a_b: rand_fill(&mut rng, rank * dims.hidden, 0.05), + b_b: rand_fill(&mut rng, gdn_dims.num_kh * rank, 0.05), + a_a: rand_fill(&mut rng, rank * dims.hidden, 0.05), + b_a: rand_fill(&mut rng, gdn_dims.num_kh * rank, 0.05), + a_out: rand_fill(&mut rng, rank * gdn_dims.output_dim, 0.05), + b_out: rand_fill(&mut rng, dims.hidden * rank, 0.05), + }, }) .collect(); let fwd = forward_full( &caches[0], &layers, &loras, &head, &dims, &gdn_dims, &cfg, rank, scale, ); - let (_, _, analytic) = - nll_and_grads(&fwd, &layers, &loras, &head, &dims, num_slots, rank, scale); + let (_, _, analytic) = nll_and_grads( + &fwd, + &layers, + &loras, + &head, + &dims, + &gdn_dims, + num_slots, + &slot_kinds, + rank, + scale, + ); println!(" fd-eps center {fd_eps:.0e} (per-entry min over 0.25/0.5/1/2x)"); let mut worst = 0.0f64; let mut all_pass = true; for slot in 0..num_slots { let layer_idx = slot_layers[slot]; - for (name, alen) in [ - ("a_q", analytic[slot].a_q.len()), - ("b_q", analytic[slot].b_q.len()), - ("a_v", analytic[slot].a_v.len()), - ("b_v", analytic[slot].b_v.len()), - ] { - let agrad = match name { + // Build the list of (name, len) pairs for this slot's kind. + let arrays: &[(&str, usize)] = match slot_kinds[slot] { + MixerKind::Gqa => &[ + ("a_q", analytic[slot].a_q.len()), + ("b_q", analytic[slot].b_q.len()), + ("a_v", analytic[slot].a_v.len()), + ("b_v", analytic[slot].b_v.len()), + ], + MixerKind::Gdn => &[ + ("a_qkv", analytic[slot].a_qkv.len()), + ("b_qkv", analytic[slot].b_qkv.len()), + ("a_z", analytic[slot].a_z.len()), + ("b_z", analytic[slot].b_z.len()), + ("a_b", analytic[slot].a_b.len()), + ("b_b", analytic[slot].b_b.len()), + ("a_a", analytic[slot].a_a.len()), + ("b_a", analytic[slot].b_a.len()), + ("a_out", analytic[slot].a_out.len()), + ("b_out", analytic[slot].b_out.len()), + ], + }; + for &(name, alen) in arrays { + let agrad: &[f32] = match name { "a_q" => &analytic[slot].a_q, "b_q" => &analytic[slot].b_q, "a_v" => &analytic[slot].a_v, - _ => &analytic[slot].b_v, + "b_v" => &analytic[slot].b_v, + "a_qkv" => &analytic[slot].a_qkv, + "b_qkv" => &analytic[slot].b_qkv, + "a_z" => &analytic[slot].a_z, + "b_z" => &analytic[slot].b_z, + "a_b" => &analytic[slot].a_b, + "b_b" => &analytic[slot].b_b, + "a_a" => &analytic[slot].a_a, + "b_a" => &analytic[slot].b_a, + "a_out" => &analytic[slot].a_out, + _ => &analytic[slot].b_out, }; + if alen == 0 { + // skip arrays that don't belong to this slot's kind + continue; + } let mut idxs = top_k_indices(agrad, probe.min(alen)); - let seed = (slot as u64 * 4 + let seed = (slot as u64 * 10 + match name { "a_q" => 0, "b_q" => 1, "a_v" => 2, - _ => 3, + "b_v" => 3, + "a_qkv" => 0, + "b_qkv" => 1, + "a_z" => 2, + "b_z" => 3, + "a_b" => 4, + "b_b" => 5, + "a_a" => 6, + "b_a" => 7, + "a_out" => 8, + _ => 9, }) ^ 0xDEAD; for p in strided_probes(alen, probe.min(alen), seed) { @@ -1017,20 +1205,40 @@ fn main() -> Result<(), Box> { let a = agrad[k]; // central FD on the chain NLL let save = { - let arr = match name { + let arr: &mut Vec = match name { "a_q" => &mut loras[slot].a_q, "b_q" => &mut loras[slot].b_q, "a_v" => &mut loras[slot].a_v, - _ => &mut loras[slot].b_v, + "b_v" => &mut loras[slot].b_v, + "a_qkv" => &mut loras[slot].a_qkv, + "b_qkv" => &mut loras[slot].b_qkv, + "a_z" => &mut loras[slot].a_z, + "b_z" => &mut loras[slot].b_z, + "a_b" => &mut loras[slot].a_b, + "b_b" => &mut loras[slot].b_b, + "a_a" => &mut loras[slot].a_a, + "b_a" => &mut loras[slot].b_a, + "a_out" => &mut loras[slot].a_out, + _ => &mut loras[slot].b_out, }; arr[k] }; let bump = |loras: &mut [LoraParams], val: f32| { - let arr = match name { + let arr: &mut Vec = match name { "a_q" => &mut loras[slot].a_q, "b_q" => &mut loras[slot].b_q, "a_v" => &mut loras[slot].a_v, - _ => &mut loras[slot].b_v, + "b_v" => &mut loras[slot].b_v, + "a_qkv" => &mut loras[slot].a_qkv, + "b_qkv" => &mut loras[slot].b_qkv, + "a_z" => &mut loras[slot].a_z, + "b_z" => &mut loras[slot].b_z, + "a_b" => &mut loras[slot].a_b, + "b_b" => &mut loras[slot].b_b, + "a_a" => &mut loras[slot].a_a, + "b_a" => &mut loras[slot].b_a, + "a_out" => &mut loras[slot].a_out, + _ => &mut loras[slot].b_out, }; arr[k] = val; }; @@ -1107,11 +1315,39 @@ fn main() -> Result<(), Box> { let init_amp = 1.0 / (dims.hidden as f32).sqrt(); let mut rng = 0xFEED_FACEu64; let mut loras: Vec = (0..num_slots) - .map(|_| LoraParams { - a_q: rand_fill(&mut rng, rank * dims.hidden, init_amp), - b_q: vec![0.0; 2 * dims.q_dim * rank], - a_v: rand_fill(&mut rng, rank * dims.hidden, init_amp), - b_v: vec![0.0; dims.kv_dim * rank], + .map(|s| match slot_kinds[s] { + MixerKind::Gqa => LoraParams { + a_q: rand_fill(&mut rng, rank * dims.hidden, init_amp), + b_q: vec![0.0; 2 * dims.q_dim * rank], + a_v: rand_fill(&mut rng, rank * dims.hidden, init_amp), + b_v: vec![0.0; dims.kv_dim * rank], + a_qkv: Vec::new(), + b_qkv: Vec::new(), + a_z: Vec::new(), + b_z: Vec::new(), + a_b: Vec::new(), + b_b: Vec::new(), + a_a: Vec::new(), + b_a: Vec::new(), + a_out: Vec::new(), + b_out: Vec::new(), + }, + MixerKind::Gdn => LoraParams { + a_q: Vec::new(), + b_q: Vec::new(), + a_v: Vec::new(), + b_v: Vec::new(), + a_qkv: rand_fill(&mut rng, rank * dims.hidden, init_amp), + b_qkv: vec![0.0; gdn_dims.qkv_dim * rank], + a_z: rand_fill(&mut rng, rank * dims.hidden, init_amp), + b_z: vec![0.0; gdn_dims.output_dim * rank], + a_b: rand_fill(&mut rng, rank * dims.hidden, init_amp), + b_b: vec![0.0; gdn_dims.num_kh * rank], + a_a: rand_fill(&mut rng, rank * dims.hidden, init_amp), + b_a: vec![0.0; gdn_dims.num_kh * rank], + a_out: rand_fill(&mut rng, rank * gdn_dims.output_dim, init_amp), + b_out: vec![0.0; dims.hidden * rank], + }, }) .collect(); @@ -1138,9 +1374,20 @@ fn main() -> Result<(), Box> { &caches, &layers, &loras, &head, &dims, &gdn_dims, &cfg, rank, scale, ); let base_valid = eval_valid(&loras); - match base_valid { - Some(v) => println!("\n step 0 train NLL: {base_nll:.4} held-out NLL: {v:.4}"), - None => println!("\n step 0 train NLL: {base_nll:.4}"), + if !emit_json { + match base_valid { + Some(v) => println!("\n step 0 train NLL: {base_nll:.4} held-out NLL: {v:.4}"), + None => println!("\n step 0 train NLL: {base_nll:.4}"), + } + } + if emit_json { + let val_json = match base_valid { + Some(v) => format!("{v:.6}"), + None => "null".to_string(), + }; + println!( + "@@lattice {{\"ev\":\"train_step\",\"step\":0,\"loss\":{base_nll:.6},\"val_loss\":{val_json},\"lr\":{lr:.6}}}" + ); } let tstep = Instant::now(); @@ -1149,70 +1396,208 @@ fn main() -> Result<(), Box> { let fwd = forward_full( ctx, &layers, &loras, &head, &dims, &gdn_dims, &cfg, rank, scale, ); - let (_nll, _n, grads) = - nll_and_grads(&fwd, &layers, &loras, &head, &dims, num_slots, rank, scale); + let (_nll, _n, grads) = nll_and_grads( + &fwd, + &layers, + &loras, + &head, + &dims, + &gdn_dims, + num_slots, + &slot_kinds, + rank, + scale, + ); for slot in 0..num_slots { let li = slot_layers[slot]; - adam.step( - &format!("l{li}_a_q"), - &mut loras[slot].a_q, - &grads[slot].a_q, - lr, - beta1, - beta2, - eps_adam, - 0.0, - false, - ); - adam.step( - &format!("l{li}_b_q"), - &mut loras[slot].b_q, - &grads[slot].b_q, - lr, - beta1, - beta2, - eps_adam, - 0.0, - false, - ); - adam.step( - &format!("l{li}_a_v"), - &mut loras[slot].a_v, - &grads[slot].a_v, - lr, - beta1, - beta2, - eps_adam, - 0.0, - false, - ); - adam.step( - &format!("l{li}_b_v"), - &mut loras[slot].b_v, - &grads[slot].b_v, - lr, - beta1, - beta2, - eps_adam, - 0.0, - false, - ); + match slot_kinds[slot] { + MixerKind::Gqa => { + adam.step( + &format!("l{li}_a_q"), + &mut loras[slot].a_q, + &grads[slot].a_q, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_b_q"), + &mut loras[slot].b_q, + &grads[slot].b_q, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_a_v"), + &mut loras[slot].a_v, + &grads[slot].a_v, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_b_v"), + &mut loras[slot].b_v, + &grads[slot].b_v, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + } + MixerKind::Gdn => { + adam.step( + &format!("l{li}_a_qkv"), + &mut loras[slot].a_qkv, + &grads[slot].a_qkv, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_b_qkv"), + &mut loras[slot].b_qkv, + &grads[slot].b_qkv, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_a_z"), + &mut loras[slot].a_z, + &grads[slot].a_z, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_b_z"), + &mut loras[slot].b_z, + &grads[slot].b_z, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_a_b"), + &mut loras[slot].a_b, + &grads[slot].a_b, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_b_b"), + &mut loras[slot].b_b, + &grads[slot].b_b, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_a_a"), + &mut loras[slot].a_a, + &grads[slot].a_a, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_b_a"), + &mut loras[slot].b_a, + &grads[slot].b_a, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_a_out"), + &mut loras[slot].a_out, + &grads[slot].a_out, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + adam.step( + &format!("l{li}_b_out"), + &mut loras[slot].b_out, + &grads[slot].b_out, + lr, + beta1, + beta2, + eps_adam, + 0.0, + false, + ); + } + } } if step % log_every == 0 || step == steps { let mean_nll = eval_chain_nll( &caches, &layers, &loras, &head, &dims, &gdn_dims, &cfg, rank, scale, ); - match eval_valid(&loras) { - Some(v) => println!( - " step {step:4} train NLL: {mean_nll:.4} held-out NLL: {v:.4} (train d {:+.4})", - mean_nll - base_nll - ), - None => println!( - " step {step:4} train NLL: {mean_nll:.4} (delta from base: {:+.4})", - mean_nll - base_nll - ), + let val_nll = eval_valid(&loras); + if !emit_json { + match val_nll { + Some(v) => println!( + " step {step:4} train NLL: {mean_nll:.4} held-out NLL: {v:.4} (train d {:+.4})", + mean_nll - base_nll + ), + None => println!( + " step {step:4} train NLL: {mean_nll:.4} (delta from base: {:+.4})", + mean_nll - base_nll + ), + } + } + if emit_json { + let val_json = match val_nll { + Some(v) => format!("{v:.6}"), + None => "null".to_string(), + }; + println!( + "@@lattice {{\"ev\":\"train_step\",\"step\":{step},\"loss\":{mean_nll:.6},\"val_loss\":{val_json},\"lr\":{lr:.6}}}" + ); } } } @@ -1221,16 +1606,91 @@ fn main() -> Result<(), Box> { &caches, &layers, &loras, &head, &dims, &gdn_dims, &cfg, rank, scale, ); let secs = tstep.elapsed().as_secs_f64(); - match (base_valid, eval_valid(&loras)) { - (Some(b), Some(f)) => println!( - "\n=== done: train {base_nll:.4}→{final_nll:.4} ({:+.4}) | held-out {b:.4}→{f:.4} ({:+.4}) in {secs:.1}s ===", - final_nll - base_nll, - f - b - ), - _ => println!( - "\n=== done: base NLL {base_nll:.4} → final NLL {final_nll:.4} ({:+.4}) in {secs:.1}s ===", - final_nll - base_nll - ), + let final_valid = eval_valid(&loras); + if !emit_json { + match (base_valid, final_valid) { + (Some(b), Some(f)) => println!( + "\n=== done: train {base_nll:.4}→{final_nll:.4} ({:+.4}) | held-out {b:.4}→{f:.4} ({:+.4}) in {secs:.1}s ===", + final_nll - base_nll, + f - b + ), + _ => println!( + "\n=== done: base NLL {base_nll:.4} → final NLL {final_nll:.4} ({:+.4}) in {secs:.1}s ===", + final_nll - base_nll + ), + } } + + // Save PEFT adapter if --save was specified. + let saved_path: Option = if let Some(ref out_path) = save_path { + let mut lora_layers = std::collections::HashMap::new(); + for (slot, lp) in loras.iter().enumerate() { + let layer_idx = slot_layers[slot]; + // q_proj: A=(rank, hidden), B=(2*q_dim, rank). + // b_q was allocated as vec![0.0; 2 * dims.q_dim * rank], so d_out = 2*q_dim. + lora_layers.insert( + (layer_idx, "q_proj".to_string()), + LoraLayer { + a: lp.a_q.clone(), + b: lp.b_q.clone(), + d_in: dims.hidden, + d_out: 2 * dims.q_dim, + rank, + }, + ); + // v_proj: A=(rank, hidden), B=(kv_dim, rank). + lora_layers.insert( + (layer_idx, "v_proj".to_string()), + LoraLayer { + a: lp.a_v.clone(), + b: lp.b_v.clone(), + d_in: dims.hidden, + d_out: dims.kv_dim, + rank, + }, + ); + } + let adapter = LoraAdapter::new( + LoraConfig { + rank, + alpha, + target_modules: vec!["q_proj".to_string(), "v_proj".to_string()], + }, + lora_layers, + ); + match adapter.save_safetensors(out_path) { + Ok(()) => { + println!(" adapter saved to {}", out_path.display()); + Some(out_path.display().to_string()) + } + Err(e) => { + eprintln!( + "warning: failed to save adapter to {}: {e}", + out_path.display() + ); + None + } + } + } else { + None + }; + + if emit_json { + let best_val_json = match base_valid { + Some(_) => { + let bv = eval_valid(&loras).unwrap_or(f32::INFINITY); + format!("{bv:.6}") + } + None => "null".to_string(), + }; + let saved_json = match &saved_path { + Some(p) => format!("\"{}\"", p.replace('\\', "\\\\").replace('"', "\\\"")), + None => "null".to_string(), + }; + println!( + "@@lattice {{\"ev\":\"train_done\",\"base_nll\":{base_nll:.6},\"final_nll\":{final_nll:.6},\"best_val\":{best_val_json},\"duration_s\":{secs:.3},\"saved\":{saved_json}}}" + ); + } + Ok(()) } diff --git a/crates/tune/src/bin/train_grad_layer23.rs b/crates/tune/src/bin/train_grad_layer23.rs index 5006df3753..23375a5a5d 100644 --- a/crates/tune/src/bin/train_grad_layer23.rs +++ b/crates/tune/src/bin/train_grad_layer23.rs @@ -18,7 +18,7 @@ // every base weight is frozen; only the four LoRA factors move. // // Qwen3.5 RMSNorm is SHIFTED (x·inv·(1+gamma)); rms_norm_forward/rmsnorm_backward -// use plain gamma, so the layer/final norms get (1+gamma) precomputed weights. +// apply (1+gamma) internally, so the layer/final norm fields carry raw gamma. // q_norm/k_norm are shifted inside gqa_forward_with_cache, so they stay raw. // // Usage: train_grad_layer23 --model-dir --data-dir [--steps 25] @@ -140,8 +140,9 @@ struct L23Cache { seq_len: usize, } -/// Borrowed frozen weights for layer 23 + head. Norm `*_shift` fields are the -/// `(1 + gamma)` shifted layer norms; `q_norm`/`k_norm` are raw (gqa shifts). +/// Borrowed frozen weights for layer 23 + head. Norm `pre_norm`/`post_norm`/ +/// `final_norm` fields carry raw gamma; rms_norm_forward/rmsnorm_backward apply +/// (1+gamma) internally. `q_norm`/`k_norm` are raw (gqa shifts internally). struct L23Weights<'a> { w_q: &'a [f32], w_k: &'a [f32], @@ -153,9 +154,9 @@ struct L23Weights<'a> { w_up: &'a [f32], w_down: &'a [f32], lm_head: &'a [f32], - pre_shift: &'a [f32], - post_shift: &'a [f32], - final_shift: &'a [f32], + pre_norm: &'a [f32], + post_norm: &'a [f32], + final_norm: &'a [f32], } struct Lora<'a> { @@ -221,7 +222,7 @@ fn forward_layer23(cache: &L23Cache, w: &L23Weights, lora: &Lora, d: &Dims) -> L let mut normed_pre_attn = vec![0.0f32; seq_len * h]; for t in 0..seq_len { let x_t = &cache.h_in[t * h..(t + 1) * h]; - let (n, _inv) = rms_norm_forward(x_t, w.pre_shift, d.eps); + let (n, _inv) = rms_norm_forward(x_t, w.pre_norm, d.eps); normed_pre_attn[t * h..(t + 1) * h].copy_from_slice(&n); } @@ -261,7 +262,7 @@ fn forward_layer23(cache: &L23Cache, w: &L23Weights, lora: &Lora, d: &Dims) -> L let mut positions = Vec::new(); for t in (cache.completion_start - 1)..seq_len - 1 { let h_mid_t = &post_attn[t * h..(t + 1) * h]; - let (normed_ffn, inv_ffn) = rms_norm_forward(h_mid_t, w.post_shift, d.eps); + let (normed_ffn, inv_ffn) = rms_norm_forward(h_mid_t, w.post_norm, d.eps); let (ffn_out, gate_pre, up_pre) = swiglu_forward(&normed_ffn, w.w_gate, w.w_up, w.w_down, h, d.inter); @@ -269,7 +270,7 @@ fn forward_layer23(cache: &L23Cache, w: &L23Weights, lora: &Lora, d: &Dims) -> L for (o, f) in h_out.iter_mut().zip(ffn_out.iter()) { *o += *f; } - let (final_normed, inv_final) = rms_norm_forward(&h_out, w.final_shift, d.eps); + let (final_normed, inv_final) = rms_norm_forward(&h_out, w.final_norm, d.eps); let logits = lm_head_logits(w.lm_head, &final_normed, h, d.vocab); positions.push(PosFwd { @@ -341,7 +342,7 @@ fn nll_and_grads(cache: &L23Cache, w: &L23Weights, lora: &Lora, d: &Dims) -> (f3 } // lm_head VJP → final_norm backward → into the residual at position t. let d_final = linear_vjp(w.lm_head, &d_logits, h, d.vocab); - let d_h_out = rmsnorm_backward(&p.post_ffn, w.final_shift, p.inv_final, &d_final); + let d_h_out = rmsnorm_backward(&p.post_ffn, w.final_norm, p.inv_final, &d_final); d_post_ffn[p.t * h..(p.t + 1) * h].copy_from_slice(&d_h_out); } @@ -359,7 +360,7 @@ fn nll_and_grads(cache: &L23Cache, w: &L23Weights, lora: &Lora, d: &Dims) -> (f3 h, d.inter, ); - let d_h_mid = rmsnorm_backward(&p.post_attn, w.post_shift, p.inv_ffn, &d_normed_ffn); + let d_h_mid = rmsnorm_backward(&p.post_attn, w.post_norm, p.inv_ffn, &d_normed_ffn); for (j, dv) in d_h_mid.iter().enumerate() { d_attn_out[p.t * h + j] += *dv; } @@ -395,10 +396,6 @@ fn nll_and_grads(cache: &L23Cache, w: &L23Weights, lora: &Lora, d: &Dims) -> (f3 ) } -fn shifted(gamma: &[f32]) -> Vec { - gamma.iter().map(|g| 1.0 + g).collect() -} - fn main() -> Result<(), Box> { let args: Vec = std::env::args().collect(); if parse_flag(&args, "-h") || parse_flag(&args, "--help") { @@ -468,7 +465,7 @@ fn main() -> Result<(), Box> { let (lm_head_s, final_norm_s, _embed) = model.head_weights(); let lm_head = lm_head_s.to_vec(); - let final_shift = shifted(final_norm_s); + let final_norm = final_norm_s.to_vec(); // raw gamma; rms_norm_forward applies (1+gamma) let (w_q, w_k, w_v, w_o, q_norm, k_norm, pre_attn_norm, post_attn_norm, w_gate, w_up, w_down) = model @@ -477,8 +474,8 @@ fn main() -> Result<(), Box> { let (w_q, w_k, w_v, w_o) = (w_q.to_vec(), w_k.to_vec(), w_v.to_vec(), w_o.to_vec()); let (q_norm, k_norm) = (q_norm.to_vec(), k_norm.to_vec()); let (w_gate, w_up, w_down) = (w_gate.to_vec(), w_up.to_vec(), w_down.to_vec()); - let pre_shift = shifted(pre_attn_norm); - let post_shift = shifted(post_attn_norm); + let pre_norm = pre_attn_norm.to_vec(); // raw gamma + let post_norm = post_attn_norm.to_vec(); // raw gamma let weights = L23Weights { w_q: &w_q, @@ -491,9 +488,9 @@ fn main() -> Result<(), Box> { w_up: &w_up, w_down: &w_down, lm_head: &lm_head, - pre_shift: &pre_shift, - post_shift: &post_shift, - final_shift: &final_shift, + pre_norm: &pre_norm, + post_norm: &post_norm, + final_norm: &final_norm, }; let tokenizer = model.tokenizer().clone(); From 3881654321ab7e106587d530804bc96860e718e3 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 22 Jun 2026 14:59:16 -0400 Subject: [PATCH 2/3] fix(tune): address codex review on GDN LoRA backward surface - gdn_forward_save: unconditionally reset saved LoRA state (rank, scale, h_* caches, weight matrices) before the bound-only populate, so a reused GdnSaved from a prior LoRA call cannot leak stale state into a no-LoRA forward (gdn_backward gates its LoRA-grad path on saved.lora_rank) - train_grad_full --save: branch on slot_kinds; GDN slots now save their five real modules (in_proj_qkv/z/b/a, out_proj) with loader-matched names and shapes instead of empty q_proj/v_proj; target_modules reflects the modules actually present (GQA / GDN / mixed) - clippy under --features train-backward (not linted by default CI): factor the 10-slice lora_bound tuple into a LoraBound type alias; drop the unused GdnGrads import and the dead LoraParams::zeros alias Co-Authored-By: Claude Opus 4.8 (1M context) --- .../inference/src/attention/gdn_backward.rs | 54 +++++-- crates/tune/src/bin/train_grad_full.rs | 136 +++++++++++++----- 2 files changed, 146 insertions(+), 44 deletions(-) diff --git a/crates/inference/src/attention/gdn_backward.rs b/crates/inference/src/attention/gdn_backward.rs index 18d4616774..ae22fde6eb 100644 --- a/crates/inference/src/attention/gdn_backward.rs +++ b/crates/inference/src/attention/gdn_backward.rs @@ -312,6 +312,23 @@ impl GdnSaved { /// to each projection output and caches `h = A @ x` for the backward. /// When any pair is `None`, the forward is byte-identical to the no-LoRA path. #[allow(clippy::too_many_arguments)] +/// The ten LoRA weight-matrix slices bound for one GDN forward, in fixed order: +/// (a_qkv, b_qkv, a_z, b_z, a_b, b_b, a_a, b_a, a_out, b_out). Factored out to keep +/// the `lora_bound` binding under clippy's type-complexity threshold (this module is +/// `train-backward`-gated, so default clippy never lints it — see the app-bins gate). +type LoraBound<'a> = ( + &'a [f32], + &'a [f32], + &'a [f32], + &'a [f32], + &'a [f32], + &'a [f32], + &'a [f32], + &'a [f32], + &'a [f32], + &'a [f32], +); + pub fn gdn_forward_save( inputs: &[f32], weights: &crate::attention::gdn::GatedDeltaNetWeights, @@ -356,18 +373,7 @@ pub fn gdn_forward_save( // Bind all ten LoRA slices at once. If any pair is None or rank==0 the entire // LoRA path is skipped and the forward is byte-identical to the no-LoRA path. // Option<&[f32]> is Copy, so destructuring copies the refs — no heap allocation. - let lora_bound: Option<( - &[f32], - &[f32], // a_qkv, b_qkv - &[f32], - &[f32], // a_z, b_z - &[f32], - &[f32], // a_b, b_b - &[f32], - &[f32], // a_a, b_a - &[f32], - &[f32], // a_out, b_out - )> = if lora_rank > 0 { + let lora_bound: Option = if lora_rank > 0 { match ( lora_a_qkv, lora_b_qkv, lora_a_z, lora_b_z, lora_a_b, lora_b_b, lora_a_a, lora_b_a, lora_a_out, lora_b_out, @@ -390,6 +396,30 @@ pub fn gdn_forward_save( None }; + // Reset LoRA cache state unconditionally. `saved` is a reusable buffer; a + // prior LoRA-bound call would otherwise leave stale rank/matrices/h_* here. + // gdn_backward gates its LoRA-gradient path on saved.lora_rank, so a leftover + // nonzero rank after a no-LoRA call computes phantom gradients. The Some branch + // below repopulates these; the None path now stays truly LoRA-free. + saved.lora_rank = 0; + saved.lora_scale = 0.0; + saved.h_qkv.clear(); + saved.h_z.clear(); + saved.h_b.clear(); + saved.h_a.clear(); + saved.h_out.clear(); + saved.gated_buf.clear(); + saved.lora_a_qkv.clear(); + saved.lora_b_qkv.clear(); + saved.lora_a_z.clear(); + saved.lora_b_z.clear(); + saved.lora_a_b.clear(); + saved.lora_b_b.clear(); + saved.lora_a_a.clear(); + saved.lora_b_a.clear(); + saved.lora_a_out.clear(); + saved.lora_b_out.clear(); + // Initialise LoRA cache buffers on the saved struct when LoRA is active. // Also clone the weight matrices so gdn_backward doesn't need them threaded through. if let Some((a_qkv, b_qkv, a_z, b_z, a_b, b_b, a_a, b_a, a_out, b_out)) = lora_bound { diff --git a/crates/tune/src/bin/train_grad_full.rs b/crates/tune/src/bin/train_grad_full.rs index 7834432b44..960d5cb93d 100644 --- a/crates/tune/src/bin/train_grad_full.rs +++ b/crates/tune/src/bin/train_grad_full.rs @@ -38,9 +38,7 @@ use std::path::{Path, PathBuf}; use std::time::Instant; use lattice_inference::attention::gdn::GatedDeltaNetWeights; -use lattice_inference::attention::gdn_backward::{ - GdnGrads, GdnSaved, gdn_backward, gdn_forward_save, -}; +use lattice_inference::attention::gdn_backward::{GdnSaved, gdn_backward, gdn_forward_save}; use lattice_inference::backward::attention_gqa::{AttnCache, gqa_backward, gqa_forward_with_cache}; use lattice_inference::backward::ops::{linear_vjp, rmsnorm_backward, swiglu_backward}; use lattice_inference::backward::tape::{rms_norm_forward, swiglu_forward}; @@ -315,11 +313,6 @@ impl LoraParams { b_out: vec![0.0; d.hidden * rank], } } - - /// Backward-compatibility alias used in TBV and eval paths. - fn zeros(rank: usize, d: &Dims) -> Self { - Self::zeros_gqa(rank, d) - } } /// LoRA gradients for one layer (same shapes as LoraParams). @@ -1626,35 +1619,114 @@ fn main() -> Result<(), Box> { let mut lora_layers = std::collections::HashMap::new(); for (slot, lp) in loras.iter().enumerate() { let layer_idx = slot_layers[slot]; - // q_proj: A=(rank, hidden), B=(2*q_dim, rank). - // b_q was allocated as vec![0.0; 2 * dims.q_dim * rank], so d_out = 2*q_dim. - lora_layers.insert( - (layer_idx, "q_proj".to_string()), - LoraLayer { - a: lp.a_q.clone(), - b: lp.b_q.clone(), - d_in: dims.hidden, - d_out: 2 * dims.q_dim, - rank, - }, - ); - // v_proj: A=(rank, hidden), B=(kv_dim, rank). - lora_layers.insert( - (layer_idx, "v_proj".to_string()), - LoraLayer { - a: lp.a_v.clone(), - b: lp.b_v.clone(), - d_in: dims.hidden, - d_out: dims.kv_dim, - rank, - }, - ); + match slot_kinds[slot] { + MixerKind::Gqa => { + // q_proj: A=(rank, hidden), B=(2*q_dim, rank). + // b_q was allocated as vec![0.0; 2 * dims.q_dim * rank], so d_out = 2*q_dim. + lora_layers.insert( + (layer_idx, "q_proj".to_string()), + LoraLayer { + a: lp.a_q.clone(), + b: lp.b_q.clone(), + d_in: dims.hidden, + d_out: 2 * dims.q_dim, + rank, + }, + ); + // v_proj: A=(rank, hidden), B=(kv_dim, rank). + lora_layers.insert( + (layer_idx, "v_proj".to_string()), + LoraLayer { + a: lp.a_v.clone(), + b: lp.b_v.clone(), + d_in: dims.hidden, + d_out: dims.kv_dim, + rank, + }, + ); + } + MixerKind::Gdn => { + // GDN trains five LoRA modules. Save each under the name and shape + // the loader expects (crates/tune/src/lora/mod.rs). Previously this + // loop wrote empty q_proj/v_proj for GDN slots (a_q/a_v are empty + // for GDN) and omitted every GDN module, yielding a useless adapter. + // in_proj_qkv: A=(rank, hidden), B=(qkv_dim, rank). + lora_layers.insert( + (layer_idx, "in_proj_qkv".to_string()), + LoraLayer { + a: lp.a_qkv.clone(), + b: lp.b_qkv.clone(), + d_in: dims.hidden, + d_out: gdn_dims.qkv_dim, + rank, + }, + ); + // in_proj_z: A=(rank, hidden), B=(output_dim, rank). + lora_layers.insert( + (layer_idx, "in_proj_z".to_string()), + LoraLayer { + a: lp.a_z.clone(), + b: lp.b_z.clone(), + d_in: dims.hidden, + d_out: gdn_dims.output_dim, + rank, + }, + ); + // in_proj_b (beta): A=(rank, hidden), B=(num_kh, rank). + lora_layers.insert( + (layer_idx, "in_proj_b".to_string()), + LoraLayer { + a: lp.a_b.clone(), + b: lp.b_b.clone(), + d_in: dims.hidden, + d_out: gdn_dims.num_kh, + rank, + }, + ); + // in_proj_a (alpha): A=(rank, hidden), B=(num_kh, rank). + lora_layers.insert( + (layer_idx, "in_proj_a".to_string()), + LoraLayer { + a: lp.a_a.clone(), + b: lp.b_a.clone(), + d_in: dims.hidden, + d_out: gdn_dims.num_kh, + rank, + }, + ); + // out_proj: A=(rank, output_dim), B=(hidden, rank). + lora_layers.insert( + (layer_idx, "out_proj".to_string()), + LoraLayer { + a: lp.a_out.clone(), + b: lp.b_out.clone(), + d_in: gdn_dims.output_dim, + d_out: dims.hidden, + rank, + }, + ); + } + } + } + // target_modules lists only the module names actually present in this run, + // so a pure-GQA, pure-GDN, or mixed adapter each declares the right set. + let mut target_modules: Vec = Vec::new(); + if slot_kinds.iter().any(|k| matches!(k, MixerKind::Gqa)) { + target_modules.push("q_proj".to_string()); + target_modules.push("v_proj".to_string()); + } + if slot_kinds.iter().any(|k| matches!(k, MixerKind::Gdn)) { + target_modules.push("in_proj_qkv".to_string()); + target_modules.push("in_proj_z".to_string()); + target_modules.push("in_proj_b".to_string()); + target_modules.push("in_proj_a".to_string()); + target_modules.push("out_proj".to_string()); } let adapter = LoraAdapter::new( LoraConfig { rank, alpha, - target_modules: vec!["q_proj".to_string(), "v_proj".to_string()], + target_modules, }, lora_layers, ); From 5e0bec19fe72ea25f58713c719281bea5090683c Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 22 Jun 2026 15:17:54 -0400 Subject: [PATCH 3/3] fix(tune): validate in_proj_b/in_proj_a GDN LoRA modules train_grad_full --save now emits all five GDN LoRA modules, but validate_against only recognized in_proj_qkv/in_proj_z/out_proj, so a saved full-GDN adapter failed validation before it could load. The forward pass already applies in_proj_b/in_proj_a LoRA (gdn_fused.rs) with d_out = linear_num_key_heads, and safetensors parsing + set_lora accept them; the only gap was these two match arms. Add them so the validation contract matches the forward and the trainer's saved output. Adds a regression test covering all five GDN LoRA modules with config-derived dims, so the contract can't silently drift again. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/tune/src/lora/mod.rs | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/crates/tune/src/lora/mod.rs b/crates/tune/src/lora/mod.rs index 605e9269f3..648499cf98 100644 --- a/crates/tune/src/lora/mod.rs +++ b/crates/tune/src/lora/mod.rs @@ -213,6 +213,8 @@ impl LoraAdapter { ("o_proj", true) => (config.full_q_dim(), config.hidden_size), ("in_proj_qkv", false) => (config.hidden_size, config.linear_qkv_dim()), ("in_proj_z", false) => (config.hidden_size, config.linear_output_dim()), + ("in_proj_b", false) => (config.hidden_size, config.linear_num_key_heads), + ("in_proj_a", false) => (config.hidden_size, config.linear_num_key_heads), ("out_proj", false) => (config.linear_output_dim(), config.hidden_size), ("gate_proj", _) => (config.hidden_size, config.intermediate_size), ("up_proj", _) => (config.hidden_size, config.intermediate_size), @@ -482,5 +484,33 @@ mod tests { let err = adapter.validate_against(&cfg).unwrap_err(); assert!(err.to_string().contains("not a recognised")); } + + #[test] + fn test_validate_against_gdn_all_modules_pass() { + // Regression: train_grad_full --save emits all five GDN LoRA modules + // (in_proj_qkv/z/b/a, out_proj). The forward applies every one of them + // (gdn_fused.rs), so validate_against must accept each with the dims the + // loader and trainer use. Dims are derived from the config, not hardcoded, + // so this stays correct if the reference dims change. + let cfg = Qwen35Config::qwen35_0_8b(); + let gdn_layer = (0..cfg.num_hidden_layers) + .find(|&i| !cfg.is_full_attention(i)) + .expect("0.8b config has linear-attention layers"); + let h = cfg.hidden_size; + let cases = [ + ("in_proj_qkv", h, cfg.linear_qkv_dim()), + ("in_proj_z", h, cfg.linear_output_dim()), + ("in_proj_b", h, cfg.linear_num_key_heads), + ("in_proj_a", h, cfg.linear_num_key_heads), + ("out_proj", cfg.linear_output_dim(), h), + ]; + for (module, d_in, d_out) in cases { + let adapter = make_adapter_for_layer(gdn_layer, module, d_in, d_out); + assert!( + adapter.validate_against(&cfg).is_ok(), + "GDN LoRA module {module} should validate (d_in={d_in}, d_out={d_out})" + ); + } + } } }