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.cpp | src/models/exaone4.cpp | |
|---|---|---|
| per-layer norms created | attn_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_norm | absent | absent |
ffn_norm | absent | absent |
| Q/K/V read | cur = inpL (:92) | cur = inpL (:118) |
| attention output | build_norm(cur, attn_post_norm) (:160-163) | build_norm(cur, attn_post_norm) (:152) |
ffn_inp | add(cur, inpSA) (:165) | add(cur, inpSA) (:155) |
| FFN input | build_ffn(ffn_inp, ...) (:169) | build_ffn(ffn_inp, ...) (:159) |
| FFN output | build_norm(cur, ffn_post_norm) (:177-179) | build_norm(cur, ffn_post_norm) (:166) |
| residual | add(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.