Skip to content

Adjusting post norm MLP to be properly differentiated - #1890

Merged
jlarson4 merged 1 commit into
devfrom
fix/post-norm-component-names
Oct 12, 2026
Merged

jlarson4 merged 1 commit into
devfrom
fix/post-norm-component-names

Conversation

@jlarson4

Copy link
Copy Markdown
Collaborator

Description

As part of reviewing #1878, I recognized that what is listed under ln1/ln2 could vary from model to model. This breaking change is intended to normalize these names, to ensure the same values always appear in the same location, and that post norm values do not appear under pre-norm names.

  • OLMo 2 family post-norms moved to ln1_post/ln2_post – OLMo 2, OLMo 3, EXAONE-4 and FlexOlmo declared their post-attention/post-feedforward norms in the ln1/ln2 slots, so readers of the pre-sublayer norm silently got the post-norm; those blocks now have no ln1/ln2 and such readers fail loudly.
  • BlockBridge derives post-norm-only wiring – a block with ln2_post and no ln2 gets hook_resid_mid = mlp.hook_in and the MLP-side hook_mlp_in capture automatically, replacing the OLMo 2 adapter's hand-written overrides.
  • OLMo Hybrid re-keys its post-attention norm per layer – post_attention_layernorm is the pre-MLP norm on linear-attention layers but follows attention on full-attention layers; full-attention layers now mount it as ln1_post, so ln2 always feeds the MLP, and hook_resid_mid is restored on both layer types.
  • Gemma 3n contribution hooks carry the post-norm outputs – hook_attn_out/hook_mlp_out pointed at the raw sublayer outputs, which are not what the layer adds to the residual; norms now use the canonical keys and both aliases resolve to ln1_post/ln2_post.hook_out, as on Gemma 2/3.
  • Mid-layer cached-scale LN error explains post-norm-only blocks – apply_ln=True at a mid layer on OLMo 2-style blocks previously divided by the post-attention norm's scale; it now raises, saying the block has no pre-sublayer norm and pointing to apply_ln=False or final-layer LN.
  • Stale OLMo 2 split-QKV xfail converted to a regression test – the env-gated strict xfail already passed on dev with the 1B model; after the rename the fork is correct by construction (no ln1 to wire) rather than by capture timing.
  • Adapter guide and model_structure.md state the rule – norm keys name the norm's position, not the HF attribute; post-norms never go in ln1/ln2.

Breaking changes

  • Hook and cache names on OLMo 2, OLMo 3, EXAONE-4 and FlexOlmo change from blocks.{i}.ln1.*/ln2.* to blocks.{i}.ln1_post.*/ln2_post.*, and on OLMo Hybrid full-attention layers from blocks.{i}.ln2.* to blocks.{i}.ln1_post.*; the canonical aliases (hook_attn_out, hook_mlp_out, hook_resid_mid) keep their meaning.
  • Gemma 3n hook_attn_out/hook_mlp_out change values (raw sublayer output → post-norm output), and its norm hooks are renamed from the HF names to ln1, ln1_post, ln2, ln2_post.

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4
jlarson4 merged commit 0b3ff8f into dev Oct 12, 2026
28 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant