attentions_output / mlps_output are not the residual contribution on post-sublayer-norm ("sandwich") architectures (Gemma-2/3, GLM-4, OLMo-2)

[Butanium]

Follow-up to #51. That issue covered architectures where the sublayer module output includes the residual (BLOOM, MPT, DBRX) and is being fixed on dev. This is the neighbouring semantic gap in the other direction: architectures where the module output is less than the residual contribution, because a post-sublayer norm sits outside the module.

The pattern

Gemma-2 and Gemma-3 (Gemma2DecoderLayer / Gemma3DecoderLayer, transformers 4.56):

resid_mid  = resid_pre + post_attention_layernorm( self_attn( input_layernorm(resid_pre) ) )
resid_post = resid_mid + post_feedforward_layernorm( mlp( pre_feedforward_layernorm(resid_mid) ) )

GLM-4 (Glm4DecoderLayer) — same shape, attention-side norm spelled unambiguously, which frees post_attention_layernorm to play its Llama role (pre-MLP norm):

resid_mid  = resid_pre + post_self_attn_layernorm( self_attn( input_layernorm(resid_pre) ) )
resid_post = resid_mid + post_mlp_layernorm( mlp( post_attention_layernorm(resid_mid) ) )

OLMo-2 (Olmo2DecoderLayer): same two post-norms, and no pre-attention norm at all.

So attentions_output[i] / mlps_output[i] (the self_attn / mlp module outputs) are the tensors before the post-norm; what is actually added to the residual stream is the post-norm output. layers_output[i] == layers_input[i] + attentions_output[i] + mlps_output[i] fails on these models — measured relative error on the tiny test zoo: gemma-2 0.78, gemma-3 15, glm-4 1.9 (0.0 on every plain pre-LN architecture).

Naming trap: Gemma's post_attention_layernorm normalizes the attention output; Llama's identically-named module normalizes the residual on its way into the MLP. Detection must not key on that name alone.

How interp-engine handles it

decoderesearch/interp-engine keeps both tensors as separate named capture points rather than redefining one (interp_engine/facts.py, "post-sublayer (sandwich) norms"; interp_engine/model.py::resolve_point):

• attn_out / mlp_out = raw sublayer module output.
• attn_out_post / mlp_out_post = defined as the residual contribution: resolves to the post-sublayer norm's output when one exists, and aliases the raw point on models without post-norms, so callers never branch on architecture. Their pinned invariant is resid_post == resid_pre + attn_out_post + mlp_out_post (tests/test_sandwich_norms.py).
• Sandwich detection keys only on the MLP-side norm name (post_feedforward_layernorm, post_mlp_layernorm, post_mlp_norm, post_ffn_norm, post_norm2, ...), never on post_attention_layernorm alone — exactly because of the Gemma/Llama collision. Attention-side spellings: post_self_attn_layernorm, post_norm1, then the ambiguous post_attention_layernorm / post_attn_norm only when an MLP-side norm was found.
• resid_mid is taken as the input of the pre-MLP norm (pre_feedforward_layernorm on Gemma-2/3, post_attention_layernorm on Llama-shaped blocks and GLM-4).

Proposal

Follow the same split instead of remapping attentions_output on these models:

• keep attentions_output / mlps_output as the module output (never residual-added, per Cross-architecture semantic inconsistency: attentions_output / mlps_output include the residual on BLOOM #51);
• add contribution accessors (attentions_output_post / mlps_output_post, or *_contribution) that target the post-sublayer norm output where one exists and alias *_output elsewhere, with a RenameConfig override analogous to attn_output_source / mlp_output_source;
• detect the block shape from the MLP-side norm names as above;
• pin layers_output == layers_input + attentions_output_post + mlps_output_post in the tests for gemma-2/3, glm-4, olmo-2 tiny models.

Note the docs on dev currently state the caveat explicitly; this issue is about closing it.

---

Opened by Claude (Fable 5) via Claude Code on behalf of @Butanium, from the #51 investigation.
Co-Authored-By: Claude noreply@anthropic.com

[crishN144]

Confirmed on current transformers (5.17.0) rather than the 4.56 in your write-up: the three forwards are unchanged (modeling_gemma2.py:291, modeling_glm4.py:67, modeling_olmo2.py:285), so the attribute names in your table still hold. On axolotl-ai-co/gemma-3-34M, layer 0: attentions_output norm 1.273 against 27.750 for the post_attention_layernorm output, and mlps_output 0.355 against 27.625 for post_feedforward_layernorm: 21.8x and 77.8x short of the contribution.

Happy to build it as proposed: keep attentions_output / mlps_output as the module output, add the contribution accessors aliasing the raw point where no post-norm exists, detect structurally from the MLP-side names with the unambiguous-first attention-side order, RenameConfig override, and pin layers_output == layers_input + <attn>_post + <mlp>_post with the raw-output negative alongside it. Most of the block shapes can be asserted on stubs the way test_sandwich_norms.py does, so OLMo-2 does not need a checkpoint in the matrix.

Three things I would rather you decide than guess:

1. _post or _contribution for the suffix.
2. Should the identity also run as a load-time check? It needs a full forward pass, so it is heavier than the current mlps_output[0] == layers_output[0] guard.
3. On the Cross-architecture semantic inconsistency: attentions_output / mlps_output include the residual on BLOOM #51 architectures (BLOOM, MPT, DBRX), should the contribution accessor alias the remapped pre-residual source or the raw module output? resolve_point aliases the boundary, which I read as the remapped one, but the two mechanisms meet only here.

Scope #53 only; leaving #54 and #56 alone.