Skip to main content

Module norm

Module norm 

Source
Expand description

The normalisation at ONE site in a decoder: before the attention branch, before the FFN branch, or before the LM head.

Three variants, because llama.cpp’s build_norm (llama-graph.cpp) has three answers at those sites and ferrox used to have one. It takes a norm TYPE (LLM_NORM = LayerNorm, LLM_NORM_RMS = RMSNorm) and a weight that may be null, and applies the multiply only if (mw). Everything below is a reading of that function and of the three graphs that reach its corners.

// the ordinary case, and nearly every architecture on the generic path
ffn_inp = x       + attn(rms(x, attn_norm))
out     = ffn_inp + ffn(rms(ffn_inp, ffn_norm))

Gemma-2 added a sandwich: the same two pre-norms, plus a norm on each branch’s OUTPUT before its residual add. ferrox has carried those two as crate::decoder::AttnWeights::post_attn_norm and crate::decoder::AttnWeights::post_ffn_norm for a long time, and they are not this type – they are Option<Vec<f32>> RMSNorms, and no architecture has ever wanted anything else there.

§NormOp::None: olmo2 and exaone4

The sandwich with the bread taken off. They have the two post-norms and NO pre-norms at all, and both sublayers read the raw residual:

ffn_inp = x       + post_attn_norm(attn(x))
out     = ffn_inp + post_ffn_norm(ffn(ffn_inp))

That claim is a reading of both files, not a family resemblance:

src/models/olmo2.cppsrc/models/exaone4.cpp
per-layer norms createdattn_q_norm, attn_k_norm, attn_post_norm, ffn_post_norm (:45-52)attn_post_norm, attn_q_norm, attn_k_norm, ffn_post_norm (:60-67)
attn_normabsentabsent
ffn_normabsentabsent
Q/K/V readcur = inpL (:92)cur = inpL (:118)
attention outputbuild_norm(cur, attn_post_norm) (:160-163)build_norm(cur, attn_post_norm) (:152)
ffn_inpadd(cur, inpSA) (:165)add(cur, inpSA) (:155)
FFN inputbuild_ffn(ffn_inp, ...) (:169)build_ffn(ffn_inp, ...) (:159)
FFN outputbuild_norm(cur, ffn_post_norm) (:177-179)build_norm(cur, ffn_post_norm) (:166)
residualadd(cur, ffn_inp) (:182)add(cur, ffn_inp) (:169)

Line for line the same graph. So the two rows share ONE implementation – this module – rather than getting one arm each. What they do NOT share is their QK-norm style (olmo2 norms the 2-D projection over its whole width, exaone4 per head after build_qkv has reshaped), which is why each still needs its own fixture: tests/post_norm_only_graphs.rs.

§NormOp::LayerNormNoParams: olmo, and only olmo

OLMo-1 is a THIRD shape, and it is not about the residual at all. src/models/olmo.cpp:27-35 creates Q/K/V, attn_output and gate/up/down and not one norm tensor – no attn_norm, no ffn_norm, no output_norm – and its graph normalises at all three sites with a null weight AND a null bias:

// olmo.cpp:65-67, :104-106, :128-130
cur = build_norm(inpL, NULL, NULL, LLM_NORM, il);

LLM_NORM is ggml_norm (llama-graph.cpp’s build_norm), which subtracts the mean and divides by the standard deviation (ggml/src/ggml-cpu/ops.cpp:3716-3745); with both weight and bias null, build_norm does nothing further. So it is a pre-norm layer like llama, with a different norm FUNCTION and no parameters at all. NormOp::None is no help here and neither is NormOp::Rms.

This is the only architecture in llama.cpp that does it. Scanned over every build_norm call in all 140 src/models/*.cpp graphs, extracting the weight argument: three calls pass a null weight to LLM_NORM, and all three are olmo.cpp. (talkie.cpp passes a null weight to LLM_NORM_RMS at five sites, which is a different function and a different row.) So this variant closes exactly one refusal and the hoped-for shared cause is not there – see capability::NON_PARAMETRIC_LAYER_NORM, which says so where the next person will look.

§Why this is a type and not a bool on ModelConfig

The RMSNorm weights are handed to fused Metal kernels that apply the norm INSIDE the kernel (PrefillDenseLayerMetal::attn_norm_w, MoeLayerMetal::ffn_norm_w, launch_decode_dense_layer). A flag on the config would leave every one of those launches free to keep reading a &[f32] that no longer means anything, and nothing would fail: that is precisely this repo’s dominant bug shape, two structures that must agree with nothing enforcing it, and it is how post_attn_norm was lost from a decode path once already.

Making the slot an enum instead means a fused launch cannot compile until it has said what it does when there is no weight to hand over. NormOp::rms_weights returns None for BOTH non-RMS variants, and every GPU call site turns that into a fall-back to the host body, which computes the right thing. The disagreement is a type error rather than a silent wrong answer. Decoder::final_norm is this type for the same reason: olmo is the first architecture whose FINAL norm is not an RMSNorm either, and the fused stacks that fold final_norm + lm_head + argmax had Some(&self.final_norm) written into them unconditionally.

Enums§

NormOp
The normalisation applied at one norm site.