Fix Qwen3 model conversion perplexity issue

qwen3_headwise_norm.md

@@ -0,0 +1,33 @@
+## Qwen3 Head-wise QK Norm Fix
+
+- **Root cause**  
+  Qwen3 applies per-head RMSNorm (`q_norm`/`k_norm`) after the `q_proj`/`k_proj` layers. The original TransMLA pipeline discarded those modules during partial RoPE and low-rank QKV stages, so the calibration data and the final MLA attention operated on unnormalized Q/K activations. This broke the expected scale of dot products, which explains the multi-order-of-magnitude perplexity regression observed after conversion.
+
+- **What changed**  
+  - Preserve and reuse the original `q_norm` / `k_norm` modules in `PartialRope`, `LoraQKV`, and the exported `MLAAttention`.  
+  - Apply those norms to the head-wise (non-RoPE) components before rotary embedding so the converted model matches Qwen3’s behavior.  
+  - Teach `get_qkv_calibrate_outputs` to tap post-norm activations when present and flatten them so PCA still works.  
+  - Surface a `use_qk_head_norm` flag in `Qwen3MLAConfig`/`config.json` so downstream loads know to instantiate the extra norms.
+
+- **How to verify locally**
+  1. Install deps: `pip install -r requirements.txt`.
+  2. Run conversion (example for 8B, requires a CUDA GPU ≥ 48 GB):  
+     ```bash
+     python3 transmla/converter.py \
+       --model-path Qwen/Qwen3-8B \
+       --save-path outputs/qwen3-8B-mla \
+       --dtype bf16 --device cuda \
+       --ppl-eval-batch-size 8 \
+       --freqfold auto --collapse auto \
+       --qk-mqa-dim 64 --q-lora-rank 512 \
+       --kv-lora-rank 512 --use-qkv-norm \
+       --use-original-norm-weights
+     ```
+  3. Compare the reported perplexities before/after conversion—they should now stay within the expected small delta (no more catastrophic blow-up).
+
+- **Current limitation**  
+  This environment has no CUDA devices (`torch.cuda.is_available() == False`), so the end-to-end Qwen3-8B conversion/eval couldn’t be executed here. The commands above can be run on a GPU machine to reproduce the fix.
+
+- **PR checklist**  
+  - [x] Code changes validated via `python3 -m compileall transmla`.  
+  - [ ] Full Qwen3-8B conversion run (blocked by hardware; see limitation above).

transmla/lora_qkv.py

@@ -59,6 +59,8 @@ def __init__(
 
         self.attention_function = ALL_ATTENTION_FUNCTIONS["sdpa"]
         self.scaling = (self.head_dim + self.qk_mqa_dim)**(-0.5)
+        self.q_norm = getattr(self_attn, "q_norm", None)
+        self.k_norm = getattr(self_attn, "k_norm", None)
 
         # -----------------Attributes for the bias-----------------
         q_bias = self_attn.q_proj.bias is not None
@@ -384,6 +386,8 @@ def forward(
 
         query_states = query_states.view(bsz, q_len, self.num_attention_heads, -1).transpose(1,2)
         q_nope, q_rope = query_states.split([self.head_dim, self.qk_mqa_dim], dim=-1)
+        if self.q_norm is not None:
+            q_nope = self.q_norm(q_nope)
 
         # key and value
         compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
@@ -399,6 +403,8 @@ def forward(
             kv_nope = self.kv_a_layernorm(kv_nope)
         kv_nope = self.kv_b_proj(kv_nope).view(bsz, q_len, self.num_attention_heads, self.head_dim * 2).transpose(1, 2)
         k_nope, value_states = kv_nope.split([self.head_dim, self.head_dim],dim=-1)
+        if self.k_norm is not None:
+            k_nope = self.k_norm(k_nope)
         key_states = torch.cat([k_nope, repeat_kv(k_rope, self.num_attention_heads)], dim=-1)
 
         attn_output, attn_weights = self.attention_function(

transmla/modify_config.py

@@ -84,6 +84,7 @@ def modify_config(model, config_path: str, args):
     config["kv_lora_rank"] = args.kv_lora_rank
 
     config["qk_latent_layernorm"] = hasattr(model.model.layers[0].self_attn, "kv_a_layernorm")
+    config["use_qk_head_norm"] = getattr(model.model.layers[0].self_attn, "q_norm", None) is not None
 
     with open(config_path, "w") as f:
         json.dump(config, f, indent=4)

transmla/partial_rope.py

@@ -70,6 +70,8 @@ def __init__(self, self_attn, key_outputs=None, freqfold=1, rope_head=1, collaps
         self.k_proj = self_attn.k_proj
         self.v_proj = self_attn.v_proj
         self.o_proj = self_attn.o_proj
+        self.q_norm = getattr(self_attn, "q_norm", None)
+        self.k_norm = getattr(self_attn, "k_norm", None)
         self._insert_kv_up_proj()
         if key_outputs is not None:
             Rk = self.joint_complex_pca(key_outputs, freqfold)
@@ -158,10 +160,15 @@ def forward(
         bsz, q_len, _ = hidden_states.size()
 
         query_states = self.q_proj(hidden_states)
+        query_states = query_states.view(bsz, q_len, self.num_attention_heads, self.head_dim)
+        if self.q_norm is not None:
+            query_states = self.q_norm(query_states)
         key_states = self.k_proj(hidden_states)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        if self.k_norm is not None:
+            key_states = self.k_norm(key_states)
         value_states = self.v_proj(hidden_states)
 
-        query_states = query_states.view(bsz, q_len, self.num_attention_heads, self.head_dim)
         k_up_weight = self.k_up_proj.weight.view(self.num_attention_heads, self.head_dim, self.latent_dim)
         query_states = torch.einsum("bthd,hdc->bhtc", query_states, k_up_weight)
 

transmla/transformers/mla.py

@@ -42,6 +42,8 @@ def __init__(self, config, layer_idx: int):
         self.qk_head_dim = config.qk_head_dim
 
         self.qk_latent_layernorm = getattr(config, "qk_latent_layernorm", True)
+        self.use_qk_head_norm = getattr(config, "use_qk_head_norm", False)
+        eps = getattr(config, "rms_norm_eps", 1e-6)
 
         self.is_causal = True
         if self.q_lora_rank is None:
@@ -64,6 +66,9 @@ def __init__(self, config, layer_idx: int):
             self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
             bias=False,
         )
+        if self.use_qk_head_norm:
+            self.q_norm = DeepseekV3RMSNorm(self.qk_nope_head_dim, eps=eps)
+            self.k_norm = DeepseekV3RMSNorm(self.qk_nope_head_dim, eps=eps)
 
         self.o_proj = nn.Linear(
             self.num_heads * self.v_head_dim,
@@ -95,6 +100,8 @@ def forward(
             q_states = self.q_b_proj(self.q_a_proj(hidden_states))
         q_states = q_states.view(query_shape).transpose(1, 2)
         q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        if self.use_qk_head_norm:
+            q_pass = self.q_norm(q_pass)
 
         compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
         k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
@@ -104,6 +111,8 @@ def forward(
         else:
             k_pass = self.kv_b_proj(k_pass).view(key_shape).transpose(1, 2)
         k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+        if self.use_qk_head_norm:
+            k_pass = self.k_norm(k_pass)
 
         k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)
 

transmla/transformers/qwen3/configuration_qwen3mla.py

@@ -13,6 +13,7 @@ def __init__(
         qk_nope_head_dim=128,
         v_head_dim=128,
         qk_latent_layernorm=True,
+        use_qk_head_norm=False,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
@@ -24,3 +25,4 @@ def __init__(
         self.qk_head_dim = qk_rope_head_dim + qk_nope_head_dim
         self.v_head_dim = v_head_dim
         self.qk_latent_layernorm = qk_latent_layernorm
+        self.use_qk_head_norm = use_qk_head_norm

transmla/utils.py

@@ -422,12 +422,19 @@ def insert_qkv_hooks(model):
     q_a_proj_outputs = {}
     kv_a_proj_with_mqa_outputs = {}
 
+    def _flatten_attention_output(output: torch.Tensor) -> torch.Tensor:
+        if isinstance(output, torch.Tensor) and output.dim() > 3:
+            return output.reshape(output.shape[0], output.shape[1], -1)
+        return output
+
     def query_hook_fn(module, input, output, index):
+        output = _flatten_attention_output(output)
         if index not in query_outputs:
             query_outputs[index] = []
         query_outputs[index].append(output.to('cpu'))
 
     def key_hook_fn(module, input, output, index):
+        output = _flatten_attention_output(output)
         if index not in key_outputs:
             key_outputs[index] = []
         key_outputs[index].append(output.to('cpu'))
@@ -448,11 +455,13 @@ def kv_a_proj_with_mqa_hook_fn(module, input, output, index):
         kv_a_proj_with_mqa_outputs[index].append(output.to('cpu'))
 
     for idx, layer in enumerate(model.model.layers):
-        if hasattr(layer.self_attn, "q_proj"):
-            query_hook = layer.self_attn.q_proj.register_forward_hook(lambda module, input, output, idx=idx: query_hook_fn(module, input, output, idx))
+        query_source = layer.self_attn.q_norm if hasattr(layer.self_attn, "q_norm") else getattr(layer.self_attn, "q_proj", None)
+        if query_source is not None:
+            query_hook = query_source.register_forward_hook(lambda module, input, output, idx=idx: query_hook_fn(module, input, output, idx))
             query_hooks.append(query_hook)
-        if hasattr(layer.self_attn, "k_proj"):
-            key_hook = layer.self_attn.k_proj.register_forward_hook(lambda module, input, output, idx=idx: key_hook_fn(module, input, output, idx))
+        key_source = layer.self_attn.k_norm if hasattr(layer.self_attn, "k_norm") else getattr(layer.self_attn, "k_proj", None)
+        if key_source is not None:
+            key_hook = key_source.register_forward_hook(lambda module, input, output, idx=idx: key_hook_fn(module, input, output, idx))
             key_hooks.append(key_hook)
         if hasattr(layer.self_attn, "v_proj"):
             value_hook = layer.self_attn.v_proj.register_forward_hook(lambda module, input, output, idx=idx: value_hook_fn(module, input, output, idx))