OLMo Family – Repairing TransformerLens residual Identities

docs/source/content/compatibility_mode.md

@@ -64,6 +64,16 @@ After `enable_compatibility_mode()`, these HT hook names fire on the **pre-norm
 - **Post-norm architectures** (OLMo 2, BERT-style) read the **post-attention residual** instead, because the norm semantically lives elsewhere in the block.
 - **MLA blocks** (DeepSeek V2 / V3 / R1) do **not** expose the split-qkv aliases — MLA's compressed K/V doesn't have a clean split.
 
+On post-norm architectures (Gemma 2/3's `ln1_post`/`ln2_post`, OLMo 2/3's `ln1`/`ln2`),
+`blocks.{i}.hook_attn_out` / `hook_mlp_out` fire **after** the post-sublayer norm, so
+that they capture the tensor added to the residual stream and the identities
+`resid_pre + attn_out == resid_mid` and `resid_mid + mlp_out == resid_post` hold
+([issue #1648](https://github.com/TransformerLensOrg/TransformerLens/issues/1648)).
+One consequence for head-level direct logit attribution: per-head contributions from
+`attn.hook_result` no longer sum to `hook_attn_out`, because the norm sits between
+them. This is inherent to post-norm — decompose heads on the pre-norm side
+(`attn.hook_out` for the raw module output) or attribute through the norm explicitly.
+
 An adapter author for a new post-norm or MLA-style architecture must handle these carve-outs in `setup_hook_compatibility`. The Gemma1/Gemma2 adapters are exemplars of when **not** to override `setup_hook_compatibility` — `GemmaTextScaledWordEmbedding` already scales internally, so any added `hook_conversion` would double-scale `embed.hook_out`.
 
 ## The four-quadrant test matrix

tests/integration/model_bridge/test_olmo2_hook_semantics.py

@@ -0,0 +1,193 @@
+"""Integration tests for OLMo 2 residual-branch hook semantics.
+
+OLMo 2 is post-norm: RMSNorm applies to each sublayer output before the
+residual add, so the compatibility aliases hook_attn_out / hook_mlp_out must
+expose the norm outputs (the additive contributions), not the raw module
+outputs. See issue #1648.
+"""
+
+import pytest
+import torch
+
+from transformer_lens.model_bridge import TransformerBridge
+
+MODEL = "hf-internal-testing/tiny-random-Olmo2ForCausalLM"
+
+
+@pytest.fixture(scope="module")
+def olmo2_bridge() -> TransformerBridge:
+    return TransformerBridge.boot_transformers(MODEL, device="cpu", dtype=torch.float32)
+
+
+@pytest.fixture(scope="module")
+def sample_tokens(olmo2_bridge: TransformerBridge) -> torch.Tensor:
+    return olmo2_bridge.to_tokens("The capital of France is Paris.")
+
+
+def test_residual_branch_hooks_decompose_stream(
+    olmo2_bridge: TransformerBridge, sample_tokens: torch.Tensor
+) -> None:
+    with torch.no_grad():
+        _, cache = olmo2_bridge.run_with_cache(sample_tokens)
+
+    for layer in range(olmo2_bridge.cfg.n_layers):
+        torch.testing.assert_close(
+            cache[f"blocks.{layer}.hook_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"],
+            cache[f"blocks.{layer}.hook_resid_mid"] + cache[f"blocks.{layer}.hook_mlp_out"],
+        )
+
+
+def test_contribution_hooks_differ_from_raw_module_outputs(
+    olmo2_bridge: TransformerBridge, sample_tokens: torch.Tensor
+) -> None:
+    """The architecture-shaped hooks stay raw; the aliases are post-norm."""
+    with torch.no_grad():
+        _, cache = olmo2_bridge.run_with_cache(sample_tokens)
+
+    for layer in range(olmo2_bridge.cfg.n_layers):
+        assert not torch.allclose(
+            cache[f"blocks.{layer}.attn.hook_out"],
+            cache[f"blocks.{layer}.hook_attn_out"],
+        )
+        assert not torch.allclose(
+            cache[f"blocks.{layer}.mlp.hook_out"],
+            cache[f"blocks.{layer}.hook_mlp_out"],
+        )
+
+
+def test_attn_out_ablation_collapses_residual_step(
+    olmo2_bridge: TransformerBridge, sample_tokens: torch.Tensor
+) -> None:
+    """Zeroing hook_attn_out must yield resid_mid == resid_pre — pins that writes
+    land on the additive contribution, with no norm applied afterwards."""
+    captured = {}
+
+    def grab(key: str):
+        def hook_fn(tensor: torch.Tensor, hook) -> torch.Tensor:
+            captured[key] = tensor.detach().clone()
+            return tensor
+
+        return hook_fn
+
+    with torch.no_grad():
+        baseline = olmo2_bridge(sample_tokens)
+        ablated = olmo2_bridge.run_with_hooks(
+            sample_tokens,
+            fwd_hooks=[
+                ("blocks.0.hook_attn_out", lambda tensor, hook: torch.zeros_like(tensor)),
+                ("blocks.0.hook_resid_pre", grab("resid_pre")),
+                ("blocks.0.hook_resid_mid", grab("resid_mid")),
+            ],
+        )
+
+    assert not torch.equal(ablated, baseline)
+    torch.testing.assert_close(captured["resid_mid"], captured["resid_pre"])
+
+
+def test_attn_out_write_lands_unmodified(
+    olmo2_bridge: TransformerBridge, sample_tokens: torch.Tensor
+) -> None:
+    """Writing v to hook_attn_out must make the contribution exactly v — a
+    post-hook norm would distort it (zero-ablation can't catch this: zero is a
+    fixed point of RMSNorm)."""
+    torch.manual_seed(1)
+    replacement = torch.randn(1, sample_tokens.shape[1], olmo2_bridge.cfg.d_model)
+    captured = {}
+
+    def grab(key: str):
+        def hook_fn(tensor: torch.Tensor, hook) -> torch.Tensor:
+            captured[key] = tensor.detach().clone()
+            return tensor
+
+        return hook_fn
+
+    with torch.no_grad():
+        olmo2_bridge.run_with_hooks(
+            sample_tokens,
+            fwd_hooks=[
+                ("blocks.0.hook_attn_out", lambda tensor, hook: replacement.clone()),
+                ("blocks.0.hook_resid_pre", grab("resid_pre")),
+                ("blocks.0.hook_resid_mid", grab("resid_mid")),
+            ],
+        )
+
+    torch.testing.assert_close(captured["resid_mid"] - captured["resid_pre"], replacement)
+
+
+def test_mlp_out_ablation_collapses_residual_step(
+    olmo2_bridge: TransformerBridge, sample_tokens: torch.Tensor
+) -> None:
+    """Zeroing hook_mlp_out must yield resid_post == resid_mid."""
+    captured = {}
+
+    def grab(key: str):
+        def hook_fn(tensor: torch.Tensor, hook) -> torch.Tensor:
+            captured[key] = tensor.detach().clone()
+            return tensor
+
+        return hook_fn
+
+    with torch.no_grad():
+        baseline = olmo2_bridge(sample_tokens)
+        ablated = olmo2_bridge.run_with_hooks(
+            sample_tokens,
+            fwd_hooks=[
+                ("blocks.0.hook_mlp_out", lambda tensor, hook: torch.zeros_like(tensor)),
+                ("blocks.0.hook_resid_mid", grab("resid_mid")),
+                ("blocks.0.hook_resid_post", grab("resid_post")),
+            ],
+        )
+
+    assert not torch.equal(ablated, baseline)
+    torch.testing.assert_close(captured["resid_post"], captured["resid_mid"])
+
+
+def test_hook_mlp_in_exposes_mlp_input(
+    olmo2_bridge: TransformerBridge, sample_tokens: torch.Tensor
+) -> None:
+    """Post-norm OLMo 2 has no pre-MLP norm, so hook_mlp_in must capture the
+    mid-residual (the true MLP input), not ln2's input (the raw MLP output)."""
+    olmo2_bridge.set_use_hook_mlp_in(True)
+    try:
+        with torch.no_grad():
+            baseline = olmo2_bridge(sample_tokens)
+            _, cache = olmo2_bridge.run_with_cache(sample_tokens)
+            ablated = olmo2_bridge.run_with_hooks(
+                sample_tokens,
+                fwd_hooks=[("blocks.0.hook_mlp_in", lambda tensor, hook: torch.zeros_like(tensor))],
+            )
+
+        for layer in range(olmo2_bridge.cfg.n_layers):
+            torch.testing.assert_close(
+                cache[f"blocks.{layer}.hook_mlp_in"],
+                cache[f"blocks.{layer}.hook_resid_mid"],
+            )
+        # Writing lands on the MLP input, so it must change the output.
+        assert not torch.equal(ablated, baseline)
+    finally:
+        olmo2_bridge.set_use_hook_mlp_in(False)
+
+
+def test_compatibility_mode_preserves_residual_semantics() -> None:
+    """The identities must survive enable_compatibility_mode (no LN folding for
+    post-norm OLMo 2, but processing must not move the aliases)."""
+    bridge = TransformerBridge.boot_transformers(MODEL, device="cpu", dtype=torch.float32)
+    tokens = bridge.to_tokens("The capital of France is Paris.")
+    bridge.enable_compatibility_mode()
+
+    with torch.no_grad():
+        _, cache = bridge.run_with_cache(tokens)
+
+    for layer in range(bridge.cfg.n_layers):
+        torch.testing.assert_close(
+            cache[f"blocks.{layer}.hook_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"],
+            cache[f"blocks.{layer}.hook_resid_mid"] + cache[f"blocks.{layer}.hook_mlp_out"],
+        )

tests/integration/model_bridge/test_olmo_hybrid_adapter.py

@@ -104,6 +104,83 @@ def grab(tensor, hook):
             assert captured.get(name) == shape, f"{name}: {captured.get(name)}"
 
 
+class TestOlmoHybridHookSemantics:
+    """hook_attn_out / hook_mlp_out must expose the tensor added to the residual
+    stream on both layer types: full-attention layers are OLMo2 post-norm
+    (contribution = norm output), linear-attention layers are pre-norm
+    (contribution = raw sublayer output). See issue #1648."""
+
+    def test_residual_contributions_decompose_stream(self, olmo_bridge, sample_tokens):
+        with torch.no_grad():
+            _, cache = olmo_bridge.run_with_cache(sample_tokens)
+
+        for layer in range(olmo_bridge.cfg.n_layers):
+            torch.testing.assert_close(
+                cache[f"blocks.{layer}.hook_resid_post"],
+                cache[f"blocks.{layer}.hook_resid_pre"]
+                + cache[f"blocks.{layer}.hook_attn_out"]
+                + cache[f"blocks.{layer}.hook_mlp_out"],
+            )
+
+    def test_full_attention_contributions_are_post_norm(self, olmo_bridge, sample_tokens):
+        """Layers 1/3 are full attention: aliases must differ from raw module outputs."""
+        with torch.no_grad():
+            _, cache = olmo_bridge.run_with_cache(sample_tokens)
+
+        for layer in (1, 3):
+            assert not torch.allclose(
+                cache[f"blocks.{layer}.attn.hook_out"],
+                cache[f"blocks.{layer}.hook_attn_out"],
+            )
+            assert not torch.allclose(
+                cache[f"blocks.{layer}.mlp.hook_out"],
+                cache[f"blocks.{layer}.hook_mlp_out"],
+            )
+
+    def test_hook_mlp_in_exposes_mid_residual_on_both_layer_types(self, olmo_bridge, sample_tokens):
+        """hook_mlp_in must capture the mid-residual on both layouts: ln2's
+        input on pre-norm linear layers, the MLP's own input on post-norm
+        full-attention layers (where ln2's input is the raw attention output)."""
+        olmo_bridge.set_use_hook_mlp_in(True)
+        try:
+            with torch.no_grad():
+                _, cache = olmo_bridge.run_with_cache(sample_tokens)
+            for layer in range(olmo_bridge.cfg.n_layers):
+                torch.testing.assert_close(
+                    cache[f"blocks.{layer}.hook_mlp_in"],
+                    cache[f"blocks.{layer}.hook_resid_pre"]
+                    + cache[f"blocks.{layer}.hook_attn_out"],
+                )
+        finally:
+            olmo_bridge.set_use_hook_mlp_in(False)
+
+    def test_full_attention_attn_out_write_lands_unmodified(self, olmo_bridge, sample_tokens):
+        """Writing v to a full-attention layer's hook_attn_out must make the
+        contribution exactly v (mlp.hook_in fires on the post-attention stream)."""
+        torch.manual_seed(1)
+        replacement = torch.randn(1, sample_tokens.shape[1], olmo_bridge.cfg.d_model)
+        captured = {}
+
+        def grab(key):
+            def hook_fn(tensor, hook):
+                captured[key] = tensor.detach().clone()
+                return tensor
+
+            return hook_fn
+
+        with torch.no_grad():
+            olmo_bridge.run_with_hooks(
+                sample_tokens,
+                fwd_hooks=[
+                    ("blocks.1.hook_attn_out", lambda tensor, hook: replacement.clone()),
+                    ("blocks.1.hook_resid_pre", grab("resid_pre")),
+                    ("blocks.1.mlp.hook_in", grab("resid_mid")),
+                ],
+            )
+
+        torch.testing.assert_close(captured["resid_mid"] - captured["resid_pre"], replacement)
+
+
 class TestOlmoHybridGeneration:
     def test_generate_with_stateful_cache(self, olmo_bridge):
         text = olmo_bridge.generate("Hello", max_new_tokens=5, do_sample=False, verbose=False)

tests/integration/model_bridge/test_residual_decomposition_identities.py

@@ -13,6 +13,8 @@
 - gpt2: block-level residual adds, the HookedTransformer reference wiring
 - mistral: Llama-style pre-RMSNorm block-level adds
 - bloom: residual added *inside* the HF attention/MLP modules (the #1639 case)
+- olmo2: post-norm inside the residual branch — RMSNorm applies to the sublayer
+  output before the add, so the contributions are the norm outputs (the #1648 case)
 
 Parallel-residual architectures (Falcon, GPT-J, NeoX, Cohere) are out of scope:
 they have no ``hook_resid_mid``.
@@ -27,6 +29,7 @@
     pytest.param("hf-internal-testing/tiny-random-gpt2", id="gpt2"),
     pytest.param("trl-internal-testing/tiny-MistralForCausalLM-0.2", id="mistral"),
     pytest.param("trl-internal-testing/tiny-BloomForCausalLM", id="bloom"),
+    pytest.param("hf-internal-testing/tiny-random-Olmo2ForCausalLM", id="olmo2"),
 ]