Adjusting post norm MLP to be properly differentiated

docs/source/content/model_structure.md

@@ -42,6 +42,8 @@ Each `blocks.{i}` is a `BlockBridge` with subcomponents:
 - `ln2`: normalization before MLP
 - `mlp`: MLP module
 
+Normalizations applied to a sublayer's output, before the residual add, use `ln1_post` (after attention) and `ln2_post` (after the MLP). Keys follow the norm's position, not its Hugging Face name: Gemma 2/3 have all four, while OLMo 2-style post-norm blocks have only `ln1_post` and `ln2_post`.
+
 ## Canonical Hook Names
 
 Use these canonical names when adding hooks or reading from the cache. Every name below exists on every `TransformerBridge` from boot, and the legacy `HookedTransformer` names in the [alias table](#legacy-alias-table) resolve to the same `HookPoint` objects whether or not compatibility mode is enabled.
@@ -108,7 +110,7 @@ Each row is checked against a tiny GPT-2 bridge and a tiny Gemma 3 bridge by `te
 | `hook_pos_embed` | `pos_embed.hook_out` | `rotary_emb.hook_out` on rotary models, where it never fires |
 | `hook_unembed` | `unembed.hook_out` | |
 | `blocks.{i}.hook_resid_pre` | `blocks.{i}.hook_in` | |
-| `blocks.{i}.hook_resid_mid` | `blocks.{i}.ln2.hook_in` | the residual sum, not `blocks.{i}.attn.hook_out` |
+| `blocks.{i}.hook_resid_mid` | `blocks.{i}.ln2.hook_in` | the residual sum, not `blocks.{i}.attn.hook_out`; `blocks.{i}.mlp.hook_in` when there is no `ln2` |
 | `blocks.{i}.hook_resid_post` | `blocks.{i}.hook_out` | |
 | `blocks.{i}.hook_attn_out` | `blocks.{i}.attn.hook_out` | `blocks.{i}.ln1_post.hook_out` when a post-attention norm exists |
 | `blocks.{i}.hook_mlp_out` | `blocks.{i}.mlp.hook_out` | `blocks.{i}.ln2_post.hook_out` when a post-MLP norm exists |

tests/integration/model_bridge/test_exaone4_adapter.py

@@ -61,9 +61,9 @@ def grab(tensor, hook):
 
         hooks = [
             "blocks.0.attn.hook_out",
-            "blocks.0.ln1.hook_out",
+            "blocks.0.ln1_post.hook_out",
             "blocks.0.mlp.hook_out",
-            "blocks.0.ln2.hook_out",
+            "blocks.0.ln2_post.hook_out",
         ]
         with torch.no_grad():
             bridge.run_with_hooks(sample_tokens, fwd_hooks=[(name, grab) for name in hooks])

tests/integration/model_bridge/test_flex_olmo_adapter.py

@@ -57,7 +57,7 @@ def test_bridge_structure(self) -> None:
         assert len(bridge.blocks) == 3
         assert isinstance(bridge.blocks[0].mlp, MoEBridge)
         # OLMo-2 post-norm layout carried through the subclass
-        assert bridge.blocks[0].ln1.original_component is not None
+        assert bridge.blocks[0].ln1_post.original_component is not None
 
     def test_forward_matches_hf(self) -> None:
         bridge, hf_reference = tiny_flex_olmo_bridge()

tests/integration/model_bridge/test_gemma3n_bridge.py

@@ -64,3 +64,38 @@ def test_run_with_cache_text_only(bridge):
         logits, cache = bridge.run_with_cache(IDS)
     assert torch.isfinite(logits).all()
     assert len(cache) > 0
+
+
+def test_contribution_hooks_are_post_norm_outputs(bridge):
+    """hook_attn_out / hook_mlp_out carry the post-sublayer norm outputs, the tensors the
+    layer actually adds (Gemma 2/3 contract), checked against an independent HF run."""
+    from transformers import AutoModelForCausalLM
+
+    hf = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, torch_dtype=torch.float32, attn_implementation="eager"
+    ).eval()
+    modules = dict(hf.named_modules())
+    reference = {}
+    handles = []
+    for layer in range(bridge.cfg.n_layers):
+        for hf_norm, alias in (
+            ("post_attention_layernorm", "hook_attn_out"),
+            ("post_feedforward_layernorm", "hook_mlp_out"),
+        ):
+            path = next(n for n in modules if n.endswith(f"layers.{layer}.{hf_norm}"))
+            key = f"blocks.{layer}.{alias}"
+            handles.append(
+                modules[path].register_forward_hook(
+                    lambda _m, _i, out, key=key: reference.__setitem__(key, out.detach())
+                )
+            )
+    try:
+        with torch.no_grad():
+            hf(IDS)
+            _, cache = bridge.run_with_cache(IDS)
+    finally:
+        for handle in handles:
+            handle.remove()
+
+    for key, expected in reference.items():
+        torch.testing.assert_close(cache[key], expected, atol=1e-5, rtol=1e-5)

tests/integration/model_bridge/test_olmo2_hook_semantics.py

@@ -191,3 +191,14 @@ def test_compatibility_mode_preserves_residual_semantics() -> None:
             cache[f"blocks.{layer}.hook_resid_post"],
             cache[f"blocks.{layer}.hook_resid_mid"] + cache[f"blocks.{layer}.hook_mlp_out"],
         )
+
+
+def test_mid_layer_cached_ln_explains_missing_pre_norm(
+    olmo2_bridge: TransformerBridge, sample_tokens: torch.Tensor
+) -> None:
+    """Post-norm-only blocks have no ln1 scale to cache; the error must say so rather
+    than point the user at a hook that cannot exist."""
+    with torch.no_grad():
+        _, cache = olmo2_bridge.run_with_cache(sample_tokens)
+    with pytest.raises(KeyError, match=r"no ln1 \(pre-attention norm\)"):
+        cache.decompose_resid(layer=1, apply_ln=True)

tests/integration/model_bridge/test_olmo_hybrid_adapter.py

@@ -123,6 +123,21 @@ def test_residual_contributions_decompose_stream(self, bridge, sample_tokens):
                 + cache[f"blocks.{layer}.hook_mlp_out"],
             )
 
+    def test_resid_mid_splits_the_stream_on_both_layer_types(self, bridge, sample_tokens):
+        with torch.no_grad():
+            _, cache = bridge.run_with_cache(sample_tokens)
+
+        for layer in range(bridge.cfg.n_layers):
+            resid_mid = cache[f"blocks.{layer}.hook_resid_mid"]
+            torch.testing.assert_close(
+                resid_mid,
+                cache[f"blocks.{layer}.hook_resid_pre"] + cache[f"blocks.{layer}.hook_attn_out"],
+            )
+            torch.testing.assert_close(
+                cache[f"blocks.{layer}.hook_resid_post"],
+                resid_mid + cache[f"blocks.{layer}.hook_mlp_out"],
+            )
+
     def test_full_attention_contributions_are_post_norm(self, bridge, sample_tokens):
         """Layers 1/3 are full attention: aliases must differ from raw module outputs."""
         with torch.no_grad():

tests/unit/model_bridge/generalized_components/test_block_post_norm_layout.py

@@ -0,0 +1,35 @@
+"""BlockBridge derives post-norm wiring from the canonical norm keys."""
+
+import pytest
+
+from transformer_lens.model_bridge.generalized_components import BlockBridge
+from transformer_lens.model_bridge.generalized_components.base import (
+    GeneralizedComponent,
+)
+
+
+def _block(*norm_keys: str) -> BlockBridge:
+    submodules = {key: GeneralizedComponent(name=key) for key in norm_keys}
+    submodules["attn"] = GeneralizedComponent(name="self_attn")
+    submodules["mlp"] = GeneralizedComponent(name="mlp")
+    return BlockBridge(name="model.layers", submodules=submodules)
+
+
+def test_post_norm_only_block_reads_resid_mid_at_the_mlp() -> None:
+    block = _block("ln1_post", "ln2_post")
+    assert block.hook_aliases["hook_resid_mid"] == "mlp.hook_in"
+    assert block.hook_aliases["hook_attn_out"] == "ln1_post.hook_out"
+    assert block.hook_aliases["hook_mlp_out"] == "ln2_post.hook_out"
+    assert block.mlp_reads_resid_directly is True
+
+
+def test_sandwich_block_keeps_the_pre_mlp_norm_as_resid_mid() -> None:
+    block = _block("ln1", "ln1_post", "ln2", "ln2_post")
+    assert block.hook_aliases["hook_resid_mid"] == "ln2.hook_in"
+    assert block.hook_aliases["hook_mlp_out"] == "ln2_post.hook_out"
+    assert block.mlp_reads_resid_directly is False
+
+
+def test_sequential_block_without_an_mlp_side_norm_is_rejected() -> None:
+    with pytest.raises(ValueError, match="neither 'ln2' nor 'ln2_post'"):
+        _block("ln1")

tests/unit/model_bridge/supported_architectures/test_exaone4_adapter.py

@@ -63,8 +63,9 @@ def test_no_bos_prepending(self, adapter):
 class TestExaone4ComponentMapping:
     def test_post_norm_positions(self, adapter):
         submodules = adapter.component_mapping["blocks"].submodules
-        assert submodules["ln1"].name == "post_attention_layernorm"
-        assert submodules["ln2"].name == "post_feedforward_layernorm"
+        assert submodules["ln1_post"].name == "post_attention_layernorm"
+        assert submodules["ln2_post"].name == "post_feedforward_layernorm"
+        assert "ln1" not in submodules and "ln2" not in submodules
 
     def test_attention_has_per_head_qk_norms(self, adapter):
         attn = adapter.component_mapping["blocks"].submodules["attn"]
@@ -88,8 +89,8 @@ def test_model_type_detection(self):
 
 
 class TestExaone4PostNormResidMid:
-    """Post-norm layout: ln2.hook_in is the MLP-branch output, so the adapter
-    must override hook_resid_mid to mlp.hook_in (olmo2 precedent)."""
+    """Post-norm-only layout: no pre-MLP norm, so hook_resid_mid is mlp.hook_in
+    (olmo2 precedent)."""
 
     def test_hook_resid_mid_points_at_mlp_hook_in(self, adapter) -> None:
         block = adapter.component_mapping["blocks"]
@@ -98,8 +99,8 @@ def test_hook_resid_mid_points_at_mlp_hook_in(self, adapter) -> None:
     def test_contribution_aliases_point_at_post_norm_outputs(self, adapter) -> None:
         """Inherited from olmo2: contributions are the post-norm outputs (#1648)."""
         block = adapter.component_mapping["blocks"]
-        assert block.hook_aliases["hook_attn_out"] == "ln1.hook_out"
-        assert block.hook_aliases["hook_mlp_out"] == "ln2.hook_out"
+        assert block.hook_aliases["hook_attn_out"] == "ln1_post.hook_out"
+        assert block.hook_aliases["hook_mlp_out"] == "ln2_post.hook_out"
 
 
 class TestExaone4NoPEGating:

tests/unit/model_bridge/supported_architectures/test_flex_olmo_adapter.py

@@ -48,18 +48,18 @@ def test_moe_replaces_dense_mlp(self, adapter):
     def test_inherits_olmo2_post_norm_layout(self, adapter):
         assert isinstance(adapter, Olmo2ArchitectureAdapter)
         blocks = adapter.component_mapping["blocks"]
-        assert isinstance(blocks.submodules["ln1"], RMSNormalizationBridge)
-        assert blocks.submodules["ln1"].name == "post_attention_layernorm"
-        assert blocks.submodules["ln2"].name == "post_feedforward_layernorm"
+        assert isinstance(blocks.submodules["ln1_post"], RMSNormalizationBridge)
+        assert blocks.submodules["ln1_post"].name == "post_attention_layernorm"
+        assert blocks.submodules["ln2_post"].name == "post_feedforward_layernorm"
         attn = blocks.submodules["attn"]
         assert attn.submodules["q_norm"].name == "q_norm"
 
     def test_inherits_post_norm_contribution_aliases(self, adapter):
         """hook_attn_out / hook_mlp_out must stay on the post-norm outputs (#1648);
         a future de-inheritance from olmo2 must not silently regress this."""
         block = adapter.component_mapping["blocks"]
-        assert block.hook_aliases["hook_attn_out"] == "ln1.hook_out"
-        assert block.hook_aliases["hook_mlp_out"] == "ln2.hook_out"
+        assert block.hook_aliases["hook_attn_out"] == "ln1_post.hook_out"
+        assert block.hook_aliases["hook_mlp_out"] == "ln2_post.hook_out"
         assert block.hook_aliases["hook_resid_mid"] == "mlp.hook_in"
 
 

tests/unit/model_bridge/supported_architectures/test_gemma3n_adapter.py

@@ -100,14 +100,21 @@ def test_altup_block_decomposition():
         "mlp",
     ):
         assert name in blocks.submodules
-    for norm in (
-        "input_layernorm",
-        "post_attention_layernorm",
-        "pre_feedforward_layernorm",
-        "post_feedforward_layernorm",
-        "post_per_layer_input_norm",
+    for key, hf_name in (
+        ("ln1", "input_layernorm"),
+        ("ln1_post", "post_attention_layernorm"),
+        ("ln2", "pre_feedforward_layernorm"),
+        ("ln2_post", "post_feedforward_layernorm"),
+        ("post_per_layer_input_norm", "post_per_layer_input_norm"),
     ):
-        assert norm in blocks.submodules
+        assert blocks.submodules[key].name == hf_name
+
+
+def test_contribution_aliases_point_at_post_norm_outputs():
+    """The residual contribution is the post-sublayer norm output (Gemma 2/3 contract)."""
+    blocks = _adapter().component_mapping["blocks"]
+    assert blocks.hook_aliases["hook_attn_out"] == "ln1_post.hook_out"
+    assert blocks.hook_aliases["hook_mlp_out"] == "ln2_post.hook_out"
 
 
 def test_kv_shared_submodules_are_optional():

tests/unit/model_bridge/supported_architectures/test_olmo2_adapter.py

@@ -6,11 +6,12 @@
 - Weight conversions: the adapter uses the base QKVO helper, so only the exact
   key set is asserted at the adapter layer.
 - Component-mapping structure, bridge types, and HF module paths.
-- Post-norm block wiring: ln1 maps to post_attention_layernorm and ln2 maps to
-  post_feedforward_layernorm. This is the central arch-specific decision.
+- Post-norm block wiring: ln1_post maps to post_attention_layernorm and ln2_post
+  to post_feedforward_layernorm, with no ln1/ln2. This is the central
+  arch-specific decision.
 - Q/K-norm submodules under attention with the correct HF names.
-- hook_alias_overrides: hook_resid_mid points at mlp.hook_in (the true post-attn
-  pre-mlp residual under post-norm), overriding BlockBridge's default ln2.hook_in.
+- Derived hook aliases: hook_resid_mid points at mlp.hook_in (the true post-attn
+  pre-mlp residual under post-norm), not BlockBridge's default ln2.hook_in.
 - GQA forward hook shapes for Q and K/V with Q/K-norm wired into the fake attention.
 - setup_component_testing: rotary wiring on template and bridge-model attentions,
   plus forcing eager attention on the HF model and per-layer self_attn configs.
@@ -205,8 +206,8 @@ def test_block_mlp_is_gated_mlp_bridge(self, adapter: Olmo2ArchitectureAdapter)
 
     def test_block_norms_are_rms(self, adapter: Olmo2ArchitectureAdapter) -> None:
         block = _mapping(adapter)["blocks"]
-        assert isinstance(block.submodules["ln1"], RMSNormalizationBridge)
-        assert isinstance(block.submodules["ln2"], RMSNormalizationBridge)
+        assert isinstance(block.submodules["ln1_post"], RMSNormalizationBridge)
+        assert isinstance(block.submodules["ln2_post"], RMSNormalizationBridge)
 
     def test_embed_and_unembed_are_correct_bridge_types(
         self, adapter: Olmo2ArchitectureAdapter
@@ -227,26 +228,28 @@ def test_mlp_gate_in_out_are_linear_bridges(self, adapter: Olmo2ArchitectureAdap
 
 
 class TestOlmo2PostNormBlockWiring:
-    """OLMo 2 is post-norm: RMSNorm applies AFTER attention (ln1) and AFTER MLP (ln2).
-    The HF module names diverge from the pre-norm Llama family default, where ln1
-    would be `input_layernorm` and ln2 would be `post_attention_layernorm`. This is
-    the central arch-specific decision and the single test most likely to catch a
-    porting regression."""
+    """OLMo 2 is post-norm only: RMSNorm applies AFTER attention and AFTER MLP, so
+    the norms sit in the ln1_post/ln2_post slots and ln1/ln2 (pre-sublayer norms)
+    are absent. Tools that read ln1/ln2 as the norm feeding a sublayer must not
+    find these. The HF name `post_attention_layernorm` is the pre-MLP norm in the
+    Llama family, which is the porting trap this guards."""
 
-    def test_ln1_maps_to_post_attention_layernorm(self, adapter: Olmo2ArchitectureAdapter) -> None:
+    def test_ln1_post_maps_to_post_attention_layernorm(
+        self, adapter: Olmo2ArchitectureAdapter
+    ) -> None:
         block = _mapping(adapter)["blocks"]
-        assert block.submodules["ln1"].name == "post_attention_layernorm"
+        assert block.submodules["ln1_post"].name == "post_attention_layernorm"
 
-    def test_ln2_maps_to_post_feedforward_layernorm(
+    def test_ln2_post_maps_to_post_feedforward_layernorm(
         self, adapter: Olmo2ArchitectureAdapter
     ) -> None:
         block = _mapping(adapter)["blocks"]
-        assert block.submodules["ln2"].name == "post_feedforward_layernorm"
+        assert block.submodules["ln2_post"].name == "post_feedforward_layernorm"
 
     def test_block_submodule_set(self, adapter: Olmo2ArchitectureAdapter) -> None:
-        """Exact block submodule set: sequential transformer with both ln1 and ln2."""
+        """Exact block submodule set: post-norms only, no pre-sublayer ln1/ln2."""
         block = _mapping(adapter)["blocks"]
-        assert set(block.submodules.keys()) == {"ln1", "ln2", "attn", "mlp"}
+        assert set(block.submodules.keys()) == {"ln1_post", "ln2_post", "attn", "mlp"}
 
 
 class TestOlmo2QKNormStructure:
@@ -276,11 +279,9 @@ def test_k_norm_uses_k_norm_hf_name(self, adapter: Olmo2ArchitectureAdapter) ->
 
 
 class TestOlmo2HookAliasOverrides:
-    """Under post-norm, `ln2.hook_in` no longer captures the residual between
-    attention and MLP (ln2 applies AFTER MLP, so ln2.hook_in is the MLP output).
-    The true post-attn pre-mlp residual is at `mlp.hook_in`. The adapter overrides
-    `hook_resid_mid` accordingly. BlockBridge's default would point at the wrong
-    tensor without this override."""
+    """With no pre-MLP norm, the post-attn pre-mlp residual is at `mlp.hook_in`.
+    BlockBridge derives that from the post-norm-only layout; its default
+    `ln2.hook_in` would not resolve."""
 
     def test_hook_resid_mid_points_at_mlp_hook_in(self, adapter: Olmo2ArchitectureAdapter) -> None:
         block = _mapping(adapter)["blocks"]
@@ -292,8 +293,8 @@ def test_contribution_aliases_point_at_post_norm_outputs(
         """hook_attn_out / hook_mlp_out must expose the tensor added to the
         residual stream, which under post-norm is the norm output (#1648)."""
         block = _mapping(adapter)["blocks"]
-        assert block.hook_aliases["hook_attn_out"] == "ln1.hook_out"
-        assert block.hook_aliases["hook_mlp_out"] == "ln2.hook_out"
+        assert block.hook_aliases["hook_attn_out"] == "ln1_post.hook_out"
+        assert block.hook_aliases["hook_mlp_out"] == "ln2_post.hook_out"
 
     def test_hook_mlp_in_captures_on_mlp(self, adapter: Olmo2ArchitectureAdapter) -> None:
         """No pre-MLP norm exists, so the hook_mlp_in capture sits on the MLP."""

tests/unit/model_bridge/supported_architectures/test_olmo3_adapter.py

@@ -50,6 +50,6 @@ def test_inherits_post_norm_contribution_aliases() -> None:
     """hook_attn_out / hook_mlp_out must stay on the post-norm outputs (#1648);
     a future de-inheritance from olmo2 must not silently regress this."""
     block = Olmo3ArchitectureAdapter(_cfg("Olmo3ForCausalLM")).component_mapping["blocks"]
-    assert block.hook_aliases["hook_attn_out"] == "ln1.hook_out"
-    assert block.hook_aliases["hook_mlp_out"] == "ln2.hook_out"
+    assert block.hook_aliases["hook_attn_out"] == "ln1_post.hook_out"
+    assert block.hook_aliases["hook_mlp_out"] == "ln2_post.hook_out"
     assert block.hook_aliases["hook_resid_mid"] == "mlp.hook_in"