[Bug Report] `hook_attn_out` / `hook_mlp_out` are pre-transform module outputs on OLMo 2/3 and Granite

[hijohnnylin]

Summary

On OLMo 2, OLMo 3, and Granite, the legacy block-level aliases

blocks.N.hook_attn_out
blocks.N.hook_mlp_out

expose the raw attention/MLP module output before the transform that produces the tensor added to the residual stream:

• OLMo 2/3 apply a post-sublayer RMSNorm before each residual addition.
• Granite multiplies each sublayer output by residual_multiplier before adding it.

As a result, the standard TransformerLens residual identities fail:

resid_mid  = resid_pre + attn_out
resid_post = resid_mid + mlp_out

This report is specifically about the legacy TransformerLens-compatible aliases. Architecture-shaped hooks such as blocks.N.attn.hook_out and blocks.N.mlp.hook_out may reasonably expose the literal Hugging Face module outputs.

The recently fixed BLOOM issue #1639 establishes the same distinction in the opposite direction: the compatibility aliases should be additive residual contributions, not arbitrary module return values.

Versions

The numerical results below were collected with:

• transformer-lens==3.7.0
• transformers==5.15.0
• PyTorch 2.9
• Python 3.12
• float32 on CPU

I also checked transformer-lens==3.7.1. The OLMo and Granite mappings described below are unchanged. Version 3.7.1 fixes the opposite BLOOM failure but not these post-transform/scaled-residual cases.

Expected behavior

For a sequential residual block, the compatibility hooks should satisfy:

torch.testing.assert_close(
    cache[f"blocks.{layer}.hook_resid_pre"]
    + cache[f"blocks.{layer}.hook_attn_out"],
    cache[f"blocks.{layer}.hook_resid_mid"],
)

torch.testing.assert_close(
    cache[f"blocks.{layer}.hook_resid_mid"]
    + cache[f"blocks.{layer}.hook_mlp_out"],
    cache[f"blocks.{layer}.hook_resid_post"],
)

That requires hook_attn_out and hook_mlp_out to be on the residual-add side of any post-sublayer norm or residual multiplier.

Actual behavior

OLMo 2

On allenai/OLMo-2-0425-1B, TransformerLens hooks equal the raw Hugging Face module outputs and not the tensors added to the residual stream:

layer	hook	vs raw module output	vs tensor added to residual
0	hook_attn_out	cos +1.000000, rel 0.000000	cos +0.544204, rel 7.76
0	hook_mlp_out	cos +1.000000, rel 0.000000	cos +0.298442, rel 14.76
8	hook_attn_out	cos +1.000000, rel 0.000000	cos -0.088797, rel 79.00
8	hook_mlp_out	cos +1.000000, rel 0.000000	cos +0.750923, rel 8.19
15	hook_attn_out	cos +1.000000, rel 0.000000	cos +0.275484, rel 11.37
15	hook_mlp_out	cos +1.000000, rel 0.000001	cos +0.761274, rel 12.22

The same values are returned through HookedTransformer and TransformerBridge for this model.

TransformerLens also contradicts itself without any Hugging Face comparison:

allenai/OLMo-2-0425-1B
 layer 0:  resid_pre + attn_out vs resid_mid   cos +0.348623  rel 4.466390
           resid_mid + mlp_out  vs resid_post  cos +0.226747  rel 5.319047
 layer 8:  resid_pre + attn_out vs resid_mid   cos +0.065463  rel 13.629590
           resid_mid + mlp_out  vs resid_post  cos +0.740115  rel 2.605765
 layer 15: resid_pre + attn_out vs resid_mid   cos +0.213075  rel 5.350170
           resid_mid + mlp_out  vs resid_post  cos +0.599600  rel 7.310253

OLMo 3 inherits the OLMo 2 adapter unchanged, so it has the same bridge alias placement:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/supported_architectures/olmo3.py#L1-L15

Granite

On ibm-granite/granite-3.3-2b-instruct, the error is a pure scale:

layer	hook	vs raw module output	vs tensor added to residual
0	hook_attn_out	cos +1.000000, rel 0.000000	cos +1.000000, rel 3.545455
20	hook_attn_out	cos +1.000000, rel 0.000001	cos +1.000000, rel 3.545455
39	hook_mlp_out	cos +1.000000, rel 0.000000	cos +1.000000, rel 3.545455

The checkpoint uses residual_multiplier=0.22, and 3.545455 == (1 - 0.22) / 0.22. Cosine similarity cannot detect this failure because the wrong tensor is collinear with the correct one.

Granite is currently available through TransformerBridge, not the converted HookedTransformer registry.

Counterexample: Gemma 3

Gemma 3 also normalizes each sublayer output before the residual addition, but TransformerLens places the compatibility aliases after the norms:

layer	hook	vs raw module output	vs tensor added to residual
0	hook_attn_out	cos +0.215148, rel 1.15	cos +1.000000, rel 0.000000
13	hook_mlp_out	cos +0.498468, rel 155.35	cos +1.000000, rel 0.000002
25	hook_mlp_out	cos +0.543978, rel 498.20	cos +1.000000, rel 0.000001

Its residual identities hold. This shows that the OLMo behavior is not a general TransformerLens convention for post-norm models.

Minimal TransformerLens-only reproduction

This check uses only TransformerLens's own cached tensors; it does not need a separate reference implementation.

pip install "transformer-lens==3.7.1" "transformers==5.15.0"
python residual_hook_identity.py olmo2
python residual_hook_identity.py granite

# residual_hook_identity.py
from __future__ import annotations

import sys

import torch
from transformers import AutoTokenizer

PROMPT = "The capital of France is Paris."
DEVICE = "cpu"
DTYPE = torch.float32


def cosine(a: torch.Tensor, b: torch.Tensor) -> float:
    a = a.double().flatten()
    b = b.double().flatten()
    return float((a @ b) / (a.norm() * b.norm()))


def relative_error(a: torch.Tensor, b: torch.Tensor) -> float:
    a = a.double()
    b = b.double()
    return float((a - b).norm() / b.norm())


def load_model(model_id: str, bridge: bool):
    if bridge:
        from transformer_lens.model_bridge import TransformerBridge

        return TransformerBridge.boot_transformers(
            model_id, device=DEVICE, dtype=DTYPE
        )

    from transformer_lens import HookedTransformer

    return HookedTransformer.from_pretrained_no_processing(
        model_id, device=DEVICE, dtype=DTYPE
    )


def check(model_id: str, layers: list[int], bridge: bool) -> None:
    model = load_model(model_id, bridge)
    tokens = AutoTokenizer.from_pretrained(model_id)(
        PROMPT,
        return_tensors="pt",
        add_special_tokens=False,
    )["input_ids"].to(DEVICE)

    names = [
        f"blocks.{layer}.hook_{name}"
        for layer in layers
        for name in (
            "resid_pre",
            "attn_out",
            "resid_mid",
            "mlp_out",
            "resid_post",
        )
    ]

    with torch.no_grad():
        _, cache = model.run_with_cache(
            tokens,
            names_filter=names,
            prepend_bos=False,
        )

    print(f"{model_id} ({'TransformerBridge' if bridge else 'HookedTransformer'})")
    for layer in layers:
        pre = cache[f"blocks.{layer}.hook_resid_pre"]
        attn = cache[f"blocks.{layer}.hook_attn_out"]
        mid = cache[f"blocks.{layer}.hook_resid_mid"]
        mlp = cache[f"blocks.{layer}.hook_mlp_out"]
        post = cache[f"blocks.{layer}.hook_resid_post"]

        reconstructed_mid = pre + attn
        reconstructed_post = mid + mlp

        print(
            f" layer {layer}: resid_pre + attn_out vs resid_mid "
            f"cos {cosine(reconstructed_mid, mid):+.6f} "
            f"rel {relative_error(reconstructed_mid, mid):.6f}"
        )
        print(
            f"           resid_mid + mlp_out vs resid_post "
            f"cos {cosine(reconstructed_post, post):+.6f} "
            f"rel {relative_error(reconstructed_post, post):.6f}"
        )


CASES = {
    "olmo2": ("allenai/OLMo-2-0425-1B", [0, 8, 15], False),
    "olmo2-bridge": ("allenai/OLMo-2-0425-1B", [0, 8, 15], True),
    "granite": ("ibm-granite/granite-3.3-2b-instruct", [0, 20, 39], True),
}

if __name__ == "__main__":
    section = sys.argv[1] if len(sys.argv) > 1 else "olmo2"
    if section not in CASES:
        raise SystemExit(f"choose one of: {', '.join(CASES)}")
    check(*CASES[section])

Root cause

BlockBridge defaults the compatibility aliases to the raw submodule outputs:

hook_attn_out -> attn.hook_out
hook_mlp_out  -> mlp.hook_out

It automatically moves them after a post-norm only when the adapter names the norm ln1_post or ln2_post:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/generalized_components/block.py#L29-L74

The OLMo 2 bridge adapter maps its post-attention and post-MLP norms as ln1 and ln2. It overrides hook_resid_mid, but it does not override hook_attn_out or hook_mlp_out, so both bridge aliases remain on the raw module outputs:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/supported_architectures/olmo2.py#L45-L90

The converted HookedTransformer path has the same placement error in a separate special case. It calls hook_attn_out before applying OLMo's ln1, and apply_mlp calls hook_mlp_out before the OLMo branch applies ln2:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/components/transformer_block.py#L141-L186
• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/components/transformer_block.py#L200-L210

Hugging Face applies those norms before the residual additions:

• https://github.com/huggingface/transformers/blob/v5.15.0/src/transformers/models/olmo2/modeling_olmo2.py#L260-L295

The Granite adapter also uses the default aliases and notes that Granite's scaling is handled inside the Hugging Face forward:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/supported_architectures/granite.py#L22-L26
• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/supported_architectures/granite.py#L68-L81

Hugging Face multiplies the module outputs by residual_multiplier after the modules return and before adding them:

• https://github.com/huggingface/transformers/blob/v5.15.0/src/transformers/models/granite/modeling_granite.py#L200-L255

Gemma 2/3 use ln1_post and ln2_post, which triggers BlockBridge's post-norm alias override and produces the expected semantics.

Related issue with the inverse failure mode:

• [Bug Report] BLOOM TransformerBridge maps hook_attn_out / hook_mlp_out to residual-added states, breaking HookedTransformer residual identities #1639

Impact

Anything that treats these aliases as residual contributions is wrong on the affected families:

• residual decomposition and direct logit attribution;
• activation patching or ablation at hook_attn_out / hook_mlp_out;
• checks based only on cosine similarity, which completely miss Granite's scalar error.

The intervention case is especially important: writing to a pre-transform hook means OLMo applies the norm to the replacement value, and Granite applies the multiplier after the replacement. That is not the same intervention as editing the tensor added to the residual stream.

Suggested fix

Preserve the distinction between raw module-output hooks and residual-contribution aliases.

For the OLMo 2/3 TransformerBridge path, explicit alias overrides appear sufficient:

hook_alias_overrides={
    "hook_resid_mid": "mlp.hook_in",
    "hook_attn_out": "ln1.hook_out",
    "hook_mlp_out": "ln2.hook_out",
}

The converted HookedTransformer path needs the equivalent reorder: apply ln1 before hook_attn_out, and apply ln2 before hook_mlp_out. Because apply_mlp currently invokes hook_mlp_out internally, the OLMo special case may need to call the MLP and post-MLP norm before invoking that hook, or the helper may need to accept a post-transform step.

For Granite, there is no existing submodule exactly at the scaled contribution. TransformerLens likely needs an explicit block-level hook after multiplication by residual_multiplier and before the residual addition. A cache-only rescaling would not be enough because hook interventions must also occur at the correct boundary.

Suggested regression tests:

1. The two residual identities hold for OLMo 2, OLMo 3, and Granite.
2. blocks.N.attn.hook_out / blocks.N.mlp.hook_out can remain raw module outputs.
3. Zeroing blocks.N.hook_attn_out removes the attention contribution without an additional norm or multiplier being applied afterward.
4. Zeroing blocks.N.hook_mlp_out analogously removes the MLP contribution.
5. Granite tests compare norms or exact values, not only cosine similarity.