diff --git a/qwen3_headwise_norm.md b/qwen3_headwise_norm.md new file mode 100644 index 0000000..c36327e --- /dev/null +++ b/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). diff --git a/transmla/lora_qkv.py b/transmla/lora_qkv.py index ebbe46e..0f92299 100644 --- a/transmla/lora_qkv.py +++ b/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( diff --git a/transmla/modify_config.py b/transmla/modify_config.py index 61bc41e..73f205e 100644 --- a/transmla/modify_config.py +++ b/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) diff --git a/transmla/partial_rope.py b/transmla/partial_rope.py index 52b76af..2d26439 100644 --- a/transmla/partial_rope.py +++ b/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) diff --git a/transmla/transformers/mla.py b/transmla/transformers/mla.py index 7b060a7..8f1a554 100644 --- a/transmla/transformers/mla.py +++ b/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) diff --git a/transmla/transformers/qwen3/configuration_qwen3mla.py b/transmla/transformers/qwen3/configuration_qwen3mla.py index 3a842b7..0f7cd3b 100644 --- a/transmla/transformers/qwen3/configuration_qwen3mla.py +++ b/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 diff --git a/transmla/utils.py b/transmla/utils.py index 5b8fb77..ec26f56 100644 --- a/transmla/utils.py +++ b/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))