AI Engineer Exhaustive
EasyPlusGLM-4.5

QK Norm

GLM-4.5

Easy

Normalize the query and key with RMSNorm before the dot product to stabilize attention logits.

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Attention logits are the scaled inner products q^\top k / \sqrt{d}. If the typical magnitude of q or k drifts during training, those logits blow up or collapse, softmax becomes peaky or flat, and the residual stream is hard to optimize. Henry et al. (2020), cited as [15] in the GLM-4.5 report, proposed normalizing queries and keys before the product so the scale of each vector is controlled independently of the projection weights.

GLM-4.5 adopts that idea as QK-Norm and implements it with RMSNorm rather than mean-centered LayerNorm. The paper is explicit that the goal is “to stabilize the range of attention logits.” Table 1 marks QK-Norm as Yes for the 355B model and No for GLM-4.5-Air. That split is confirmed by the released configs: "use_qk_norm": true on zai-org/GLM-4.5 and false on Air. The 355B model also uses 96 heads on a 5120 hidden size — 2.5× more heads than a naive d_{\mathrm{head}}=d_{\mathrm{model}}/H allocation with fewer heads — which makes logit scale more sensitive to head-wise variance. QK-Norm is the cheap stabilizer they paired with that head count.

How it works

After the Q and K linear maps, reshape to per-head vectors of width d = 128:

Q \in \mathbb{R}^{B \times S \times H_q \times d}, \qquad K \in \mathbb{R}^{B \times S \times H_{kv} \times d},

with H_q = 96 and H_{kv} = 8. RMSNorm runs independently on each head vector (last axis only):

\mathrm{RMSNorm}(x) = \gamma \odot \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \varepsilon}},

with a learned \gamma \in \mathbb{R}^{d} and \varepsilon = 10^{-5}. There are two separate norms, q\_\mathrm{norm} and k\_\mathrm{norm}, each with its own \gamma. Values are not normalized.

On GLM-4.5 the official order is:

  1. q = W_q h, k = W_k h (these projections do have bias; attention_bias is true in the 355B config).
  2. View as (B, S, H, d) and apply RMSNorm over d.
  3. Transpose to (B, H, S, d).
  4. Apply partial RoPE to the first d_{\mathrm{rope}} channels.
  5. Scaled dot-product attention with scale d^{-1/2}.

So the quantity that actually enters the product is a unit-RMS query/key (up to \gamma), then rotated. Normalizing after RoPE would mix the learned scale \gamma with the position-dependent rotation in a different way than the checkpoint expects.

Henry et al. originally used LayerNorm, which subtracts a mean. GLM’s RMSNorm does not. The two are not interchangeable at inference: a mean-centered QK-Norm would shift every head by a quantity the 355B weights never saw. Air omits the module entirely; its Glm4MoeAttention path skips the if self.use_qk_norm block.

Because RMSNorm is applied per head, a token whose Q projection is large in one head and small in another is rescaled head-by-head. That is the point: each head’s logit contribution stays O(1) after the 1/\sqrt{d} scale, even if W_q or W_k grow.

Official code

The official repo zai-org/GLM-4.5 does not contain the layer; the README defers to Transformers / vLLM / SGLang.

In Hugging Face Transformers, Glm4MoeAttention (src/transformers/models/glm4_moe/modeling_glm4_moe.py) constructs

q_norm = Glm4MoeRMSNorm(head_dim, eps=rms_norm_eps)
k_norm = Glm4MoeRMSNorm(head_dim, eps=rms_norm_eps)

only when config.use_qk_norm is true. The comment in that file calls this “the main diff from Llama.” Forward applies both norms immediately after .view(hidden_shape) and before the transpose / apply_rotary_pos_emb pair.

Glm4MoeConfig defaults use_qk_norm to False (Air-shaped defaults). The 355B checkpoint overrides it. Do not invent a modeling file under zai-org/GLM-4.5/inference; that folder is client/server wrappers.

Watch-outs

Sources