CRITICAL FIX: Add missing q_b_norm (unweighted RMSNorm after q_b_proj)
The HuggingFace reference (DeepseekV4ForCausalLM) applies an unweighted RMSNorm after q_b_proj, normalizing Q before attention. Without it, Q magnitudes are too large, causing attention scores to collapse to uniform (entropy ~3.2 with 24 positions) and the model to produce garbage. q_b_norm has no learnable parameters — just q / RMS(q). This explains the nearly-uniform attention weights we've been seeing.
This commit is contained in:
@@ -429,6 +429,13 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
|
||||
w[f"{pre}.q_b_proj.weight_scale"],
|
||||
w[f"{pre}.q_b_proj.weight_scale_2"]) # (T, n_h * hd)
|
||||
|
||||
# q_b_norm — unweighted RMSNorm after q_b_proj (paper §2.3.1)
|
||||
# This is critical: normalizes Q before attention, preventing score collapse.
|
||||
# No learnable parameters — just q / RMS(q).
|
||||
q_f = q.float()
|
||||
q_rms = q_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt()
|
||||
q = (q_f * q_rms).bfloat16()
|
||||
|
||||
# -- KV projection (MQA: 1 KV head) + KV norm --
|
||||
kv = nvfp4_linear(x_normed,
|
||||
w[f"{pre}.kv_proj.weight"],
|
||||
|
||||
Reference in New Issue
Block a user