Add QK-norm support for Qwen3 models
keypair34 committed
Apr 22, 2026 at 19:31 UTC
5c68a814b34991f4d87418e723989f4fb107e399
1 file changed
+30
-2
src/finetune.rs
+30
-2
@@ -291,15 +291,19 @@ fn run_finetune(
291
// ---------------------------------------------------------------------------
292
293
#[derive(serde::Deserialize)]
294
+#[allow(dead_code)]
295
struct ModelConfig {
296
hidden_size: usize,
297
num_hidden_layers: usize,
298
num_attention_heads: usize,
299
num_key_value_heads: Option<usize>,
300
+ head_dim: Option<usize>,
301
intermediate_size: usize,
302
vocab_size: usize,
303
rms_norm_eps: Option<f64>,
304
rope_theta: Option<f64>,
305
+ #[serde(default)]
306
+ tie_word_embeddings: bool,
307
}
308
309
#[derive(serde::Deserialize)]
@@ -483,6 +487,9 @@ struct TransformerLayer {
487
k_proj: FrozenLinear, // frozen
488
v_proj: LoraLinear, // LoRA applied
489
o_proj: FrozenLinear, // frozen
490
+ // QK-norm (Qwen3 only — None for Qwen2/2.5)
491
+ q_norm: Option<Tensor>, // [head_dim]
492
+ k_norm: Option<Tensor>, // [head_dim]
493
// MLP projections (SwiGLU)
494
gate_proj: FrozenLinear,
495
up_proj: FrozenLinear,
@@ -526,6 +533,18 @@ impl TransformerLayer {
533
.reshape((b, seq, self.num_kv_heads, self.head_dim))?
534
.transpose(1, 2)?;
535
536
+ // Qwen3 QK-Norm: apply per-head RMSNorm to Q and K before RoPE.
537
+ let q = if let Some(w) = &self.q_norm {
538
+ rms_norm(&q, w, self.rms_norm_eps)?
539
+ } else {
540
+ q
541
+ };
542
+ let k = if let Some(w) = &self.k_norm {
543
+ rms_norm(&k, w, self.rms_norm_eps)?
544
+ } else {
545
+ k
546
+ };
547
+
548
// Apply RoPE (requires contiguous storage for the custom kernel)
549
let q = apply_rope(&q.contiguous()?, cos, sin)?;
550
let k = apply_rope(&k.contiguous()?, cos, sin)?;
@@ -628,7 +647,7 @@ impl LoraQwenModel {
647
let vocab_size = cfg.vocab_size;
648
let num_heads = cfg.num_attention_heads;
649
let num_kv_heads = cfg.num_key_value_heads.unwrap_or(num_heads);
631
- let head_dim = hidden_size / num_heads;
650
+ let head_dim = cfg.head_dim.unwrap_or(hidden_size / num_heads);
651
let intermediate_size = cfg.intermediate_size;
652
let rms_norm_eps = cfg.rms_norm_eps.unwrap_or(1e-6);
653
let rope_theta = cfg.rope_theta.unwrap_or(10_000.0);
@@ -681,6 +700,10 @@ impl LoraQwenModel {
700
bias: None,
701
};
702
703
+ // --- QK-Norm (Qwen3 only) ---
704
+ let q_norm = vb_attn.pp("q_norm").get((head_dim,), "weight").ok();
705
+ let k_norm = vb_attn.pp("k_norm").get((head_dim,), "weight").ok();
706
+
707
// --- MLP projections (all frozen) ---
708
let vb_mlp = vb_layer.pp("mlp");
709
@@ -707,6 +730,8 @@ impl LoraQwenModel {
730
k_proj,
731
v_proj,
732
o_proj,
733
+ q_norm,
734
+ k_norm,
735
gate_proj: FrozenLinear {
736
weight: gate_weight,
737
bias: None,
@@ -730,7 +755,10 @@ impl LoraQwenModel {
755
756
// Final RMSNorm and language model head
757
let norm = vb.pp("model").pp("norm").get((hidden_size,), "weight")?;
733
- let lm_head = vb.pp("lm_head").get((vocab_size, hidden_size), "weight")?;
758
+ let lm_head = vb
759
+ .pp("lm_head")
760
+ .get((vocab_size, hidden_size), "weight")
761
+ .unwrap_or_else(|_| embed_tokens.clone());
762
763
// Precompute RoPE tables
764
let (rope_cos, rope_sin) = precompute_rope(head_dim, max_seq_len, rope_theta, device)?;