From 5e6d2b9bd8f534f16390ffd1974bc27779eefb01 Mon Sep 17 00:00:00 2001 From: jlarson4 Date: Wed, 12 Aug 2026 10:14:50 -0500 Subject: [PATCH 1/4] Fixed issue 1647 --- .../unit/model_bridge/test_config_mapping.py | 224 ++++++++++++++++++ .../config/transformer_bridge_config.py | 7 + .../position_embeddings_attention.py | 9 +- .../rotary_embedding.py | 11 +- .../model_bridge/sources/_bridge_builder.py | 5 + .../model_bridge/sources/transformers.py | 75 +++++- .../tools/model_registry/verify_models.py | 6 + .../utilities/heterogeneous_config.py | 71 ++++++ 8 files changed, 391 insertions(+), 17 deletions(-) create mode 100644 transformer_lens/utilities/heterogeneous_config.py diff --git a/tests/unit/model_bridge/test_config_mapping.py b/tests/unit/model_bridge/test_config_mapping.py index 1632352e46..716a450a5a 100644 --- a/tests/unit/model_bridge/test_config_mapping.py +++ b/tests/unit/model_bridge/test_config_mapping.py @@ -6,9 +6,16 @@ from types import SimpleNamespace +import torch + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, +) from transformer_lens.model_bridge.sources.transformers import ( map_default_transformer_lens_config, ) +from transformer_lens.utilities.heterogeneous_config import majority_value def _hf_config(**overrides: object) -> SimpleNamespace: @@ -39,3 +46,220 @@ def test_d_head_falls_back_to_derived_when_head_dim_is_none() -> None: """HF configs may carry head_dim=None; treat it the same as absent.""" mapped = map_default_transformer_lens_config(_hf_config(head_dim=None)) assert mapped.d_head == 16 // 4 + + +class _AmbiguousAccessError(Exception): + """Stand-in for transformers>=5.15 AmbiguousGlobalPerLayerAttributeError. + + Deliberately not an AttributeError: neither hasattr() nor getattr(..., default) + suppresses the real exception, which is the crash mode of issue #1647. + """ + + +class _HeterogeneousConfig: + """Minimal transformers>=5.15 heterogeneous-config contract. + + Exposes ``is_heterogeneous`` / ``per_layer_attributes`` / ``per_layer_config`` + and raises a non-AttributeError on any global read of a registered per-layer + attribute, proving the mapping never performs the forbidden access. + """ + + def __init__(self, per_layer: dict, **global_fields: object) -> None: + object.__setattr__(self, "_per_layer", per_layer) + for key, value in global_fields.items(): + object.__setattr__(self, key, value) + + def __getattribute__(self, name: str): + # __dict__ lookup keeps deepcopy working: it reconstructs instances + # attribute-by-attribute, probing before _per_layer exists. + per_layer = object.__getattribute__(self, "__dict__").get("_per_layer", {}) + if name in per_layer: + raise _AmbiguousAccessError(f"'{name}' is a per-layer attribute") + return object.__getattribute__(self, name) + + @property + def is_heterogeneous(self) -> bool: + return True + + @property + def per_layer_attributes(self) -> set: + return set(self._per_layer) + + @property + def per_layer_config(self) -> list: + return [ + SimpleNamespace(**{key: values[i] for key, values in self._per_layer.items()}) + for i in range(self.num_hidden_layers) + ] + + +_GEMMA4_GLOBALS: dict[str, object] = { + "hidden_size": 16, + "num_hidden_layers": 6, + "vocab_size": 32, + "max_position_embeddings": 64, + "intermediate_size": 32, + "sliding_window": 8, +} + + +def test_heterogeneous_head_dim_does_not_read_global() -> None: + """Gemma-4-E2B shape: per-layer head_dim (sliding 256 / full-attention 512).""" + config = _HeterogeneousConfig( + per_layer={"head_dim": [256, 256, 256, 256, 256, 512]}, + num_attention_heads=8, + num_key_value_heads=1, + **_GEMMA4_GLOBALS, + ) + mapped = map_default_transformer_lens_config(config) + assert mapped.d_head == 256 # majority-layer value, matching pre-5.15 behavior + assert mapped.per_layer_head_dim == [256, 256, 256, 256, 256, 512] + assert mapped.n_key_value_heads == 1 # uniform, still read globally + + +def test_heterogeneous_kv_heads_does_not_read_global() -> None: + """Gemma-4-31B shape: head_dim and num_key_value_heads both vary per layer.""" + config = _HeterogeneousConfig( + per_layer={ + "head_dim": [256, 256, 256, 256, 256, 512], + "num_key_value_heads": [16, 16, 16, 16, 16, 4], + }, + num_attention_heads=32, + **_GEMMA4_GLOBALS, + ) + mapped = map_default_transformer_lens_config(config) + assert mapped.d_head == 256 + assert mapped.n_key_value_heads == 16 + assert mapped.per_layer_head_dim == [256, 256, 256, 256, 256, 512] + assert mapped.per_layer_num_key_value_heads == [16, 16, 16, 16, 16, 4] + + +def test_legacy_global_fields_reconstruct_per_layer_geometry() -> None: + """Pre-5.15 Gemma 4: holds the sliding value, global_ the + full-attention value; the per-layer view is rebuilt from layer_types.""" + layer_types = ["sliding_attention"] * 5 + ["full_attention"] + mapped = map_default_transformer_lens_config( + _hf_config( + num_attention_heads=32, + head_dim=256, + global_head_dim=512, + num_key_value_heads=16, + num_global_key_value_heads=4, + layer_types=layer_types, + num_hidden_layers=6, + ) + ) + assert mapped.d_head == 256 # unchanged from the pre-fix scalar + assert mapped.n_key_value_heads == 16 + assert mapped.per_layer_head_dim == [256] * 5 + [512] + assert mapped.per_layer_num_key_value_heads == [16] * 5 + [4] + + +def test_legacy_kv_reconstruction_gated_on_attention_k_eq_v() -> None: + """HF applies num_global_key_value_heads only when attention_k_eq_v is set; + with it False, every layer runs the base KV-head count.""" + mapped = map_default_transformer_lens_config( + _hf_config( + num_attention_heads=32, + head_dim=256, + global_head_dim=512, + num_key_value_heads=16, + num_global_key_value_heads=4, + attention_k_eq_v=False, + layer_types=["sliding_attention"] * 5 + ["full_attention"], + num_hidden_layers=6, + ) + ) + assert not hasattr(mapped, "per_layer_num_key_value_heads") + assert mapped.n_key_value_heads == 16 + # head_dim reconstruction is ungated — global_head_dim applies regardless. + assert mapped.per_layer_head_dim == [256] * 5 + [512] + + +def test_unanticipated_per_layer_field_resolves_to_majority() -> None: + """Any registered per-layer field — not just the ones Gemma 4 registers today — + must resolve to its majority value instead of crashing the probe.""" + config = _HeterogeneousConfig( + per_layer={"sliding_window": [8, 8, 16, 8]}, + num_attention_heads=4, + hidden_size=16, + num_hidden_layers=4, + vocab_size=32, + max_position_embeddings=64, + intermediate_size=32, + ) + mapped = map_default_transformer_lens_config(config) + # Read via __dict__: boot consumes the mapped config that way, and attribute + # access on the (copied) heterogeneous config itself still raises. + assert mapped.__dict__["sliding_window"] == 8 + + +def test_legacy_global_field_equal_to_base_stays_scalar() -> None: + """No per-layer view when the global_* value matches the base field.""" + mapped = map_default_transformer_lens_config( + _hf_config( + head_dim=8, + global_head_dim=8, + layer_types=["sliding_attention", "full_attention"], + ) + ) + assert mapped.d_head == 8 + assert not hasattr(mapped, "per_layer_head_dim") + + +def test_homogeneous_config_gets_no_per_layer_fields() -> None: + mapped = map_default_transformer_lens_config(_hf_config(head_dim=8)) + assert not hasattr(mapped, "per_layer_head_dim") + assert not hasattr(mapped, "per_layer_num_key_value_heads") + + +def test_majority_value_semantics() -> None: + """The scalar collapse is the most-common value — not first, min, or max.""" + assert majority_value([512, 256, 256]) == 256 # majority beats first element + assert majority_value([64, 32, 64]) == 64 # majority beats min + assert majority_value([512, 256, 512, 256]) == 512 # tie breaks to earliest layer + + +def test_heterogeneous_config_through_bridge_config_build() -> None: + """Full boot pipeline (map → from_dict → passthrough) on a heterogeneous config. + + intermediate_size is both per-layer here and in _HF_PASSTHROUGH_ATTRS, so the + passthrough loop must skip it rather than perform the raising global read. + """ + config = _HeterogeneousConfig( + per_layer={ + "head_dim": [32, 32, 32, 64], + "intermediate_size": [128, 64, 64, 64], + "num_attention_heads": [8, 8, 8, 4], + }, + num_key_value_heads=2, + hidden_size=16, + num_hidden_layers=4, + vocab_size=32, + max_position_embeddings=64, + ) + bridge_config = build_bridge_config_from_hf(config, "TestArch", "test-model", torch.float32) + assert bridge_config.n_heads == 8 # majority of per-layer num_attention_heads + assert bridge_config.d_head == 32 + assert bridge_config.per_layer_head_dim == [32, 32, 32, 64] + assert bridge_config.n_key_value_heads == 2 # uniform, still read globally + assert bridge_config.d_mlp == 128 # per-layer intermediate_size collapses to max + # The passthrough loop must skip the per-layer-registered attr, not copy it. + assert not hasattr(bridge_config, "intermediate_size") + + +def test_bridge_config_retains_per_layer_fields() -> None: + """from_dict filters to signature params; the per-layer fields must survive.""" + bridge_config = TransformerBridgeConfig.from_dict( + { + "d_model": 16, + "d_head": 256, + "n_layers": 6, + "n_ctx": 64, + "n_heads": 8, + "per_layer_head_dim": [256] * 5 + [512], + "per_layer_num_key_value_heads": [16] * 5 + [4], + } + ) + assert bridge_config.per_layer_head_dim == [256] * 5 + [512] + assert bridge_config.per_layer_num_key_value_heads == [16] * 5 + [4] diff --git a/transformer_lens/config/transformer_bridge_config.py b/transformer_lens/config/transformer_bridge_config.py index ddfa446de5..e09d55e75a 100644 --- a/transformer_lens/config/transformer_bridge_config.py +++ b/transformer_lens/config/transformer_bridge_config.py @@ -74,6 +74,11 @@ def __init__( num_experts: Optional[int] = None, experts_per_token: Optional[int] = None, n_key_value_heads: Optional[int] = None, + # Heterogeneous attention geometry (e.g. Gemma 4): per-layer values when + # they vary across layers; d_head / n_key_value_heads then hold the + # majority-layer scalar and attention math is delegated to HF. + per_layer_head_dim: Optional[list] = None, + per_layer_num_key_value_heads: Optional[list] = None, relative_attention_max_distance: Optional[int] = None, relative_attention_num_buckets: Optional[int] = None, decoder_start_token_id: Optional[int] = None, @@ -169,6 +174,8 @@ def __init__( self.num_experts = num_experts self.experts_per_token = experts_per_token self.n_key_value_heads = n_key_value_heads + self.per_layer_head_dim = per_layer_head_dim + self.per_layer_num_key_value_heads = per_layer_num_key_value_heads self.relative_attention_max_distance = relative_attention_max_distance self.relative_attention_num_buckets = relative_attention_num_buckets self.decoder_start_token_id = decoder_start_token_id diff --git a/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py b/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py index 565d787ac0..839b869bae 100644 --- a/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py +++ b/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py @@ -22,6 +22,7 @@ from transformer_lens.model_bridge.generalized_components.position_embedding_hooks_mixin import ( PositionEmbeddingHooksMixin, ) +from transformer_lens.utilities.heterogeneous_config import safe_config_get from transformer_lens.utilities.hf_utils import get_rotary_pct_from_config # Global registry mapping HF attention modules to their bridge instances @@ -669,12 +670,8 @@ def get_random_inputs( inputs: Dict[str, Any] = { "hidden_states": torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype) } - num_heads = ( - self.config.num_attention_heads - if self.config and hasattr(self.config, "num_attention_heads") - else 4 - ) - head_dim = self.config.head_dim if self.config and hasattr(self.config, "head_dim") else 256 + num_heads = safe_config_get(self.config, "num_attention_heads", 4) if self.config else 4 + head_dim = safe_config_get(self.config, "head_dim", 256) if self.config else 256 dummy_qk = torch.randn(1, seq_len, num_heads, head_dim, device=device, dtype=dtype) position_ids = torch.arange(seq_len, device=device).unsqueeze(0) if self._rotary_emb is not None: diff --git a/transformer_lens/model_bridge/generalized_components/rotary_embedding.py b/transformer_lens/model_bridge/generalized_components/rotary_embedding.py index c560fff960..717024821d 100644 --- a/transformer_lens/model_bridge/generalized_components/rotary_embedding.py +++ b/transformer_lens/model_bridge/generalized_components/rotary_embedding.py @@ -10,6 +10,7 @@ from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) +from transformer_lens.utilities.heterogeneous_config import safe_config_get class RotaryEmbeddingBridge(GeneralizedComponent): @@ -62,14 +63,8 @@ def get_random_inputs( device = torch.device("cpu") if dtype is None: dtype = torch.float32 - if self.config and hasattr(self.config, "num_attention_heads"): - num_heads = self.config.num_attention_heads - else: - num_heads = 4 - if self.config and hasattr(self.config, "head_dim"): - head_dim = self.config.head_dim - else: - head_dim = 256 + num_heads = safe_config_get(self.config, "num_attention_heads", 4) if self.config else 4 + head_dim = safe_config_get(self.config, "head_dim", 256) if self.config else 256 x = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype) position_ids = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1) args: tuple = (x, position_ids) diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 0e56748606..b8cbecd225 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -13,6 +13,7 @@ ) from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.utilities.heterogeneous_config import per_layer_attr_names # Architecture-agnostic; do not extend per-architecture. _HF_PASSTHROUGH_ATTRS = [ @@ -180,7 +181,11 @@ def build_bridge_config_from_hf( bridge_config.dtype = dtype effective_config = get_effective_text_config(hf_config) + # Per-layer-registered attrs would raise on global access (transformers>=5.15). + _het_attrs = per_layer_attr_names(effective_config) | per_layer_attr_names(hf_config) for attr in _HF_PASSTHROUGH_ATTRS: + if attr in _het_attrs: + continue val = getattr(effective_config, attr, None) if val is None and effective_config is not hf_config: val = getattr(hf_config, attr, None) diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index 0200aabbbf..76d97ec501 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -28,6 +28,12 @@ from transformer_lens.model_bridge.sources._bridge_builder import _HF_PASSTHROUGH_ATTRS from transformer_lens.supported_models import MODEL_ALIASES from transformer_lens.utilities import get_device, get_tokenizer_with_bos +from transformer_lens.utilities.heterogeneous_config import ( + het_safe_view, + majority_value, + per_layer_attr_names, + per_layer_values, +) # Suppress transformers warnings that go to stderr # This prevents notebook tests from failing due to unexpected stderr output @@ -62,9 +68,31 @@ def map_default_transformer_lens_config(hf_config): A copy of hf_config with additional TransformerLens fields """ # Extract language model config from text_config for multimodal models - source_config = get_effective_text_config(hf_config) + raw_source_config = get_effective_text_config(hf_config) tl_config = copy.deepcopy(hf_config) + + # transformers>=5.15 heterogeneous configs (e.g. Gemma 4) refuse global reads of + # registered per-layer attributes — the raise is not an AttributeError, so hasattr() + # does not suppress it. The view resolves any registered field to its majority-layer + # value, keeping every probe below safe regardless of which fields a future + # architecture registers; geometry fields needing full per-layer detail are + # handled explicitly first. + het_attrs = per_layer_attr_names(raw_source_config) + source_config = het_safe_view(raw_source_config) + + def legacy_per_layer(base_name: str, global_name: str) -> Any: + """Pre-5.15 Gemma 4 split geometry across (sliding-attention layers) + and global_ (full-attention layers); rebuild the per-layer view.""" + if base_name in het_attrs or "layer_types" in het_attrs: + return None + base = getattr(source_config, base_name, None) + full = getattr(source_config, global_name, None) + layer_types = getattr(source_config, "layer_types", None) + if base is None or full is None or full == base or not layer_types: + return None + return [full if t == "full_attention" else base for t in layer_types] + if hasattr(source_config, "n_embd"): tl_config.d_model = source_config.n_embd elif hasattr(source_config, "hidden_size"): @@ -90,7 +118,29 @@ def map_default_transformer_lens_config(hf_config): source_config.num_query_heads, list ): tl_config.n_heads = max(source_config.num_query_heads) - if ( + if "num_key_value_heads" in het_attrs: + per_layer_kv = per_layer_values(source_config, "num_key_value_heads") + elif getattr(source_config, "attention_k_eq_v", True): + # HF applies num_global_key_value_heads only on attention_k_eq_v models; + # the 5.15 per-layer migration gates identically, defaulting True when absent. + per_layer_kv = legacy_per_layer("num_key_value_heads", "num_global_key_value_heads") + else: + per_layer_kv = None + kv_values = [v for v in per_layer_kv if v is not None] if per_layer_kv else [] + if kv_values: + # Heterogeneous KV geometry (e.g. Gemma 4 31B: 16 KV heads on sliding layers, + # 4 on full-attention layers). The scalar keeps the majority-layer value — + # attention math is delegated to HF for these architectures — and the + # per-layer truth is preserved alongside it. + tl_config.per_layer_num_key_value_heads = per_layer_kv + try: + num_kv_heads = int(majority_value(kv_values)) + num_heads = int(getattr(tl_config, "n_heads", 0)) + if num_kv_heads != num_heads: + tl_config.n_key_value_heads = num_kv_heads + except (TypeError, ValueError): + pass + elif ( hasattr(source_config, "num_key_value_heads") and source_config.num_key_value_heads is not None ): @@ -178,6 +228,10 @@ def map_default_transformer_lens_config(hf_config): tl_config.n_ctx = 2048 if hasattr(source_config, "n_inner"): tl_config.d_mlp = source_config.n_inner + elif "intermediate_size" in het_attrs: + # Same max collapse as the per-layer-list case below. + mlp_values = [v for v in per_layer_values(source_config, "intermediate_size") if v] + tl_config.d_mlp = max(mlp_values) if mlp_values else None elif hasattr(source_config, "intermediate_size"): intermediate_size = source_config.intermediate_size # Gemma 3n exposes a per-layer intermediate_size list (the MatFormer design permits @@ -191,7 +245,18 @@ def map_default_transformer_lens_config(hf_config): tl_config.d_mlp = source_config.mlp_hidden_size elif hasattr(tl_config, "d_model"): tl_config.d_mlp = getattr(source_config, "n_inner", 4 * tl_config.d_model) - if hasattr(source_config, "head_dim") and source_config.head_dim is not None: + if "head_dim" in het_attrs: + per_layer_hd = per_layer_values(source_config, "head_dim") + else: + per_layer_hd = legacy_per_layer("head_dim", "global_head_dim") + hd_values = [v for v in per_layer_hd if v is not None] if per_layer_hd else [] + if hd_values: + # Heterogeneous head_dim (e.g. Gemma 4: 256 on sliding layers, 512 on + # full-attention layers). Scalar d_head keeps the majority-layer value; + # the per-layer truth is preserved alongside it. + tl_config.per_layer_head_dim = per_layer_hd + tl_config.d_head = majority_value(hd_values) + elif hasattr(source_config, "head_dim") and source_config.head_dim is not None: tl_config.d_head = source_config.head_dim elif hasattr(tl_config, "d_model") and hasattr(tl_config, "n_heads"): tl_config.d_head = tl_config.d_model // tl_config.n_heads @@ -624,7 +689,11 @@ def boot( # Propagate HF-specific config attributes that adapters may need. # Canonical list lives in sources/_bridge_builder.py (architecture-agnostic). effective_config = get_effective_text_config(hf_config) + # Per-layer-registered attrs would raise on global access (transformers>=5.15). + _het_attrs = per_layer_attr_names(effective_config) | per_layer_attr_names(hf_config) for attr in _HF_PASSTHROUGH_ATTRS: + if attr in _het_attrs: + continue val = getattr(effective_config, attr, None) if val is None and effective_config is not hf_config: val = getattr(hf_config, attr, None) diff --git a/transformer_lens/tools/model_registry/verify_models.py b/transformer_lens/tools/model_registry/verify_models.py index 22541888ba..5e6a092831 100644 --- a/transformer_lens/tools/model_registry/verify_models.py +++ b/transformer_lens/tools/model_registry/verify_models.py @@ -35,6 +35,8 @@ from pathlib import Path from typing import Optional +from transformer_lens.utilities.heterogeneous_config import het_safe_view + # Exit code used for graceful interrupts (Ctrl+C). The wrapper script # recognises this and stops without marking the in-flight model as failed. _EXIT_GRACEFUL_INTERRUPT = 42 @@ -270,6 +272,10 @@ def estimate_model_params(model_id: str) -> int: lang_config = _subcfg break + # Heterogeneous configs (transformers>=5.15 Gemma 4) raise on global reads of + # per-layer fields like head_dim; the view resolves them to majority values. + lang_config = het_safe_view(lang_config) + # Extract dimensions from config (different models use different attribute names) d_model = ( getattr(lang_config, "hidden_size", None) diff --git a/transformer_lens/utilities/heterogeneous_config.py b/transformer_lens/utilities/heterogeneous_config.py new file mode 100644 index 0000000000..686e895949 --- /dev/null +++ b/transformer_lens/utilities/heterogeneous_config.py @@ -0,0 +1,71 @@ +"""Safe attribute access for transformers>=5.15 heterogeneous configs. + +Configs whose attention geometry varies across layers (e.g. Gemma 4) register +those fields as per-layer: reading one on the global config raises +``AmbiguousGlobalPerLayerAttributeError`` from ``__getattribute__``, which +neither ``hasattr()`` nor ``getattr(..., default)`` suppresses. These helpers +let config-probing code stay safe on any transformers version — the per-layer +machinery is simply absent pre-5.15, where every helper degrades to plain +attribute access. +""" + +from typing import Any + + +def per_layer_attr_names(config: Any) -> frozenset: + """Fields a heterogeneous config refuses to serve globally. + + Empty for homogeneous configs and pre-5.15 transformers. + """ + if not getattr(config, "is_heterogeneous", False): + return frozenset() + return frozenset(getattr(config, "per_layer_attributes", None) or ()) + + +def per_layer_values(config: Any, name: str) -> list: + """Collect a per-layer attribute from a heterogeneous config's layer configs.""" + per_layer_config = config.per_layer_config + return [getattr(per_layer_config[i], name, None) for i in range(len(per_layer_config))] + + +def majority_value(values: list) -> Any: + """Most common value in a per-layer list; ties break toward the earliest layer.""" + counts: dict = {} + for v in values: + counts[v] = counts.get(v, 0) + 1 + return max(counts, key=lambda v: (counts[v], -values.index(v))) + + +def safe_config_get(config: Any, name: str, default: Any = None) -> Any: + """getattr that resolves per-layer-registered fields to their majority value.""" + if name in per_layer_attr_names(config): + values = [v for v in per_layer_values(config, name) if v is not None] + return majority_value(values) if values else default + return getattr(config, name, default) + + +class HetSafeConfigView: + """Read-only getattr proxy: per-layer-registered fields resolve to their + majority-layer value instead of raising; everything else passes through. + + Wrap a config once and downstream ``hasattr``/``getattr`` probes need no + per-field awareness of heterogeneity. + """ + + def __init__(self, config: Any) -> None: + object.__setattr__(self, "_config", config) + object.__setattr__(self, "_het_attrs", per_layer_attr_names(config)) + + def __getattr__(self, name: str) -> Any: + config = object.__getattribute__(self, "_config") + if name in object.__getattribute__(self, "_het_attrs"): + values = [v for v in per_layer_values(config, name) if v is not None] + if not values: + raise AttributeError(name) + return majority_value(values) + return getattr(config, name) + + +def het_safe_view(config: Any) -> Any: + """Wrap heterogeneous configs in a :class:`HetSafeConfigView`; pass others through.""" + return HetSafeConfigView(config) if per_layer_attr_names(config) else config From 9679f5c597843a5c239c0cdc5df036721a07bfff Mon Sep 17 00:00:00 2001 From: jlarson4 Date: Wed, 12 Aug 2026 12:22:27 -0500 Subject: [PATCH 2/4] OLMo fixes --- docs/source/content/compatibility_mode.md | 10 + .../model_bridge/test_olmo2_hook_semantics.py | 193 ++++++++++++++++++ .../model_bridge/test_olmo_hybrid_adapter.py | 79 +++++++ .../test_residual_decomposition_identities.py | 3 + .../test_transformer_block_olmo_post_norm.py | 164 +++++++++++++++ .../test_exaone4_adapter.py | 6 + .../test_flex_olmo_adapter.py | 8 + .../test_olmo2_adapter.py | 14 ++ .../test_olmo3_adapter.py | 9 + .../test_olmo_hybrid_adapter.py | 29 +++ .../components/transformer_block.py | 11 +- transformer_lens/model_bridge/bridge.py | 50 ++++- .../generalized_components/block.py | 26 ++- .../supported_architectures/olmo2.py | 12 +- .../supported_architectures/olmo_hybrid.py | 38 +++- 15 files changed, 634 insertions(+), 18 deletions(-) create mode 100644 tests/integration/model_bridge/test_olmo2_hook_semantics.py create mode 100644 tests/unit/components/test_transformer_block_olmo_post_norm.py diff --git a/docs/source/content/compatibility_mode.md b/docs/source/content/compatibility_mode.md index de6eb49ab9..083c497888 100644 --- a/docs/source/content/compatibility_mode.md +++ b/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 diff --git a/tests/integration/model_bridge/test_olmo2_hook_semantics.py b/tests/integration/model_bridge/test_olmo2_hook_semantics.py new file mode 100644 index 0000000000..eb768f3f03 --- /dev/null +++ b/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"], + ) diff --git a/tests/integration/model_bridge/test_olmo_hybrid_adapter.py b/tests/integration/model_bridge/test_olmo_hybrid_adapter.py index ad0023b377..e69d66de57 100644 --- a/tests/integration/model_bridge/test_olmo_hybrid_adapter.py +++ b/tests/integration/model_bridge/test_olmo_hybrid_adapter.py @@ -104,6 +104,85 @@ 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) diff --git a/tests/integration/model_bridge/test_residual_decomposition_identities.py b/tests/integration/model_bridge/test_residual_decomposition_identities.py index 0eac7339cc..6ee50b8d91 100644 --- a/tests/integration/model_bridge/test_residual_decomposition_identities.py +++ b/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"), ] diff --git a/tests/unit/components/test_transformer_block_olmo_post_norm.py b/tests/unit/components/test_transformer_block_olmo_post_norm.py new file mode 100644 index 0000000000..6a5a36a042 --- /dev/null +++ b/tests/unit/components/test_transformer_block_olmo_post_norm.py @@ -0,0 +1,164 @@ +"""OLMo 2/3 post-norm hook placement in HookedTransformer's TransformerBlock. + +The residual identities are weight-independent, so a random-weight model with +original_architecture="Olmo2ForCausalLM" pins the hook ordering (ln1/ln2 must +apply before hook_attn_out / hook_mlp_out). See issue #1648. +""" + +import pytest +import torch + +from transformer_lens import HookedTransformer, HookedTransformerConfig + +D_VOCAB = 100 + + +@pytest.fixture(scope="module", params=["Olmo2ForCausalLM", "Olmo3ForCausalLM"]) +def olmo_model(request) -> HookedTransformer: + torch.manual_seed(0) + cfg = HookedTransformerConfig( + n_layers=2, + d_model=64, + n_ctx=32, + d_head=16, + n_heads=4, + d_mlp=128, + d_vocab=D_VOCAB, + act_fn="silu", + gated_mlp=True, + normalization_type="RMS", + positional_embedding_type="rotary", + rotary_dim=16, + original_architecture=request.param, + ) + model = HookedTransformer(cfg) + model.init_weights() + return model + + +@pytest.fixture(scope="module") +def sample_tokens() -> torch.Tensor: + torch.manual_seed(0) + return torch.randint(0, D_VOCAB, (1, 10)) + + +def test_residual_branch_hooks_decompose_stream( + olmo_model: HookedTransformer, sample_tokens: torch.Tensor +) -> None: + with torch.no_grad(): + _, cache = olmo_model.run_with_cache(sample_tokens) + + for layer in range(olmo_model.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_attn_out_ablation_collapses_residual_step( + olmo_model: HookedTransformer, 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 = olmo_model(sample_tokens) + ablated = olmo_model.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( + olmo_model: HookedTransformer, 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], olmo_model.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(): + olmo_model.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_hook_mlp_in_exposes_mid_residual( + olmo_model: HookedTransformer, sample_tokens: torch.Tensor +) -> None: + """With use_hook_mlp_in, hook_mlp_in must equal resid_mid — post-norm OLMo + has no pre-MLP norm, so the MLP input is the mid-residual itself.""" + olmo_model.cfg.use_hook_mlp_in = True + try: + with torch.no_grad(): + _, cache = olmo_model.run_with_cache(sample_tokens) + for layer in range(olmo_model.cfg.n_layers): + torch.testing.assert_close( + cache[f"blocks.{layer}.hook_mlp_in"], + cache[f"blocks.{layer}.hook_resid_mid"], + ) + finally: + olmo_model.cfg.use_hook_mlp_in = False + + +def test_mlp_out_ablation_collapses_residual_step( + olmo_model: HookedTransformer, 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 = olmo_model(sample_tokens) + ablated = olmo_model.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"]) diff --git a/tests/unit/model_bridge/supported_architectures/test_exaone4_adapter.py b/tests/unit/model_bridge/supported_architectures/test_exaone4_adapter.py index 8c5a701338..f7f67d3be5 100644 --- a/tests/unit/model_bridge/supported_architectures/test_exaone4_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_exaone4_adapter.py @@ -95,6 +95,12 @@ def test_hook_resid_mid_points_at_mlp_hook_in(self, adapter) -> None: block = adapter.component_mapping["blocks"] assert block.hook_aliases["hook_resid_mid"] == "mlp.hook_in" + 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" + class TestExaone4NoPEGating: """Hybrid checkpoints skip RoPE on full-attention layers.""" diff --git a/tests/unit/model_bridge/supported_architectures/test_flex_olmo_adapter.py b/tests/unit/model_bridge/supported_architectures/test_flex_olmo_adapter.py index 5cf89ca1fe..abfe4737fd 100644 --- a/tests/unit/model_bridge/supported_architectures/test_flex_olmo_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_flex_olmo_adapter.py @@ -54,6 +54,14 @@ def test_inherits_olmo2_post_norm_layout(self, adapter): 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_resid_mid"] == "mlp.hook_in" + def test_factory_registration(): assert SUPPORTED_ARCHITECTURES["FlexOlmoForCausalLM"] is FlexOlmoArchitectureAdapter diff --git a/tests/unit/model_bridge/supported_architectures/test_olmo2_adapter.py b/tests/unit/model_bridge/supported_architectures/test_olmo2_adapter.py index bc83a2b514..1b0c67552d 100644 --- a/tests/unit/model_bridge/supported_architectures/test_olmo2_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_olmo2_adapter.py @@ -302,6 +302,20 @@ def test_hook_resid_mid_points_at_mlp_hook_in(self, adapter: Olmo2ArchitectureAd block = _mapping(adapter)["blocks"] assert block.hook_aliases["hook_resid_mid"] == "mlp.hook_in" + def test_contribution_aliases_point_at_post_norm_outputs( + self, adapter: Olmo2ArchitectureAdapter + ) -> None: + """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" + + 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.""" + block = _mapping(adapter)["blocks"] + assert block.mlp_reads_resid_directly is True + class TestOlmo2GQAHookShapes: """Wire a fake attention module into the bridge and verify GQA hook shapes. diff --git a/tests/unit/model_bridge/supported_architectures/test_olmo3_adapter.py b/tests/unit/model_bridge/supported_architectures/test_olmo3_adapter.py index 16e2064ca9..39730d75a3 100644 --- a/tests/unit/model_bridge/supported_architectures/test_olmo3_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_olmo3_adapter.py @@ -44,3 +44,12 @@ def test_builds_same_component_surface_as_olmo2() -> None: olmo2 = Olmo2ArchitectureAdapter(_cfg("Olmo2ForCausalLM")) assert isinstance(olmo3, Olmo2ArchitectureAdapter) assert set(olmo3.component_mapping) == set(olmo2.component_mapping) + + +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_resid_mid"] == "mlp.hook_in" diff --git a/tests/unit/model_bridge/supported_architectures/test_olmo_hybrid_adapter.py b/tests/unit/model_bridge/supported_architectures/test_olmo_hybrid_adapter.py index 8342e486a2..a7a29d5c9a 100644 --- a/tests/unit/model_bridge/supported_architectures/test_olmo_hybrid_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_olmo_hybrid_adapter.py @@ -74,3 +74,32 @@ class TestOlmoHybridResidMidDropped: def test_hook_resid_mid_alias_absent(self, adapter) -> None: block = adapter.component_mapping["blocks"] assert "hook_resid_mid" not in block.hook_aliases + + +class TestOlmoHybridPerLayerContributionAliases: + """set_original_component selects hook_attn_out / hook_mlp_out per layer + type: full-attention layers are post-norm (contribution = norm output), + linear-attention layers are pre-norm (contribution = raw output). #1648.""" + + def _bind(self, adapter, hf_layer): + import copy + + block = copy.deepcopy(adapter.component_mapping["blocks"]) + block.set_original_component(hf_layer) + return block + + def test_full_attention_layer_uses_post_norm_outputs(self, adapter) -> None: + import torch.nn as nn + + hf_layer = nn.Module() + hf_layer.post_feedforward_layernorm = nn.LayerNorm(8) + block = self._bind(adapter, hf_layer) + assert block.hook_aliases["hook_attn_out"] == "ln2.hook_out" + assert block.hook_aliases["hook_mlp_out"] == "ln2_post.hook_out" + + def test_linear_attention_layer_uses_raw_outputs(self, adapter) -> None: + import torch.nn as nn + + block = self._bind(adapter, nn.Module()) + assert block.hook_aliases["hook_attn_out"] == "linear_attn.hook_out" + assert block.hook_aliases["hook_mlp_out"] == "mlp.hook_out" diff --git a/transformer_lens/components/transformer_block.py b/transformer_lens/components/transformer_block.py index b80289ae6d..7017d0acba 100644 --- a/transformer_lens/components/transformer_block.py +++ b/transformer_lens/components/transformer_block.py @@ -182,10 +182,11 @@ def forward( # and before the hook. We do it before the hook so hook_attn_out captures "that which # is added to the residual stream" attn_out = self.ln1_post(attn_out) - attn_out = self.hook_attn_out(attn_out) - if self.cfg.original_architecture in ("Olmo2ForCausalLM", "Olmo3ForCausalLM"): + # OLMo 2/3 post-norm: ln1 applies before the residual add, so it must + # precede the hook for hook_attn_out to capture the additive contribution. attn_out = self.ln1(attn_out) + attn_out = self.hook_attn_out(attn_out) if resid_pre.device != attn_out.device: resid_pre = resid_pre.to(attn_out.device) @@ -196,8 +197,8 @@ def forward( resid_mid if not self.cfg.use_hook_mlp_in else self.hook_mlp_in(resid_mid.clone()) ) if self.cfg.original_architecture in ("Olmo2ForCausalLM", "Olmo3ForCausalLM"): + # Post-norm: apply_mlp applies ln2 before hook_mlp_out internally. mlp_out = self.apply_mlp(mlp_in) - mlp_out = self.ln2(mlp_out) else: normalized_resid_mid = self.ln2(mlp_in) mlp_out = self.apply_mlp(normalized_resid_mid) @@ -227,4 +228,8 @@ def apply_mlp( mlp_out = self.mlp(normalized_resid) # [batch, pos, d_model] if self.cfg.use_normalization_before_and_after: mlp_out = self.ln2_post(mlp_out) + if self.cfg.original_architecture in ("Olmo2ForCausalLM", "Olmo3ForCausalLM"): + # OLMo 2/3 post-norm: ln2 applies before the residual add, so it must + # precede the hook for hook_mlp_out to capture the additive contribution. + mlp_out = self.ln2(mlp_out) return self.hook_mlp_out(mlp_out) diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 3cec4a623f..f1d758fe93 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -61,6 +61,9 @@ _BLOCK_PATTERN = re.compile("blocks\\.(\\d+)") +# Block-list container attributes a bridge may expose. +_BLOCK_LIST_ATTRS = ("blocks", "encoder_blocks", "decoder_blocks", "L_blocks", "H_blocks") + def _resolve_attr_path(obj: nn.Module, attr_path: str) -> torch.Tensor: """Walk a dot-separated attribute path and return the final tensor.""" @@ -440,7 +443,7 @@ def _set_processed_weight_attributes(self) -> None: d_head = self.cfg.d_head d_model = self.cfg.d_model blocks_iter = [] - for bl_name in ("blocks", "encoder_blocks", "decoder_blocks", "L_blocks", "H_blocks"): + for bl_name in _BLOCK_LIST_ATTRS: if hasattr(self, bl_name): blocks_iter.append(getattr(self, bl_name)) if not blocks_iter: @@ -578,9 +581,52 @@ def _collect_hook_aliases_from_registry(self): aliases_tuple = self._compute_hook_aliases_cached( hook_names_tuple, component_aliases_tuple ) - return dict(aliases_tuple) + aliases = dict(aliases_tuple) + aliases.update(self._collect_block_instance_aliases()) + return aliases return {} + def _collect_block_instance_aliases(self) -> Dict[str, str]: + """Collect per-block-instance aliases, overriding template-derived ones. + + The template collection above reads ``adapter.component_mapping`` and so + cannot see aliases a block rebinds per layer at bind time (heterogeneous + architectures like OlmoHybrid) or prunes for absent optional submodules. + """ + aliases: Dict[str, str] = {} + unresolved: List[str] = [] + for bl_name in _BLOCK_LIST_ATTRS: + block_list = getattr(self, bl_name, None) + if block_list is None: + continue + for i, block in enumerate(block_list): + block_aliases = getattr(block, "hook_aliases", None) + if not block_aliases: + continue + # A block with no registered hooks means the registry hasn't + # scanned it yet — unresolved aliases there are timing, not drops. + block_prefix = f"{bl_name}.{i}." + if f"{block_prefix}hook_in" not in self._hook_registry: + continue + for alias_name, target in block_aliases.items(): + targets = target if isinstance(target, list) else [target] + for single_target in targets: + full_target = f"{block_prefix}{single_target}" + if full_target in self._hook_registry: + aliases[f"{block_prefix}{alias_name}"] = full_target + break + else: + unresolved.append(f"{block_prefix}{alias_name}") + if unresolved: + # Surface drops instead of silently swallowing, mirroring + # GeneralizedComponent._register_aliases. + warnings.warn( + f"{len(unresolved)} block hook alias(es) did not resolve to a " + f"registered hook (e.g. '{unresolved[0]}').", + stacklevel=2, + ) + return aliases + def _add_aliases_to_hooks(self, hooks: Dict[str, HookPoint]) -> None: """Add aliases to hooks in place.""" component_aliases = self._collect_hook_aliases_from_registry() diff --git a/transformer_lens/model_bridge/generalized_components/block.py b/transformer_lens/model_bridge/generalized_components/block.py index a216f8743d..e983dcd6c1 100644 --- a/transformer_lens/model_bridge/generalized_components/block.py +++ b/transformer_lens/model_bridge/generalized_components/block.py @@ -58,6 +58,7 @@ def __init__( config: Optional[Any] = None, submodules: Optional[Dict[str, GeneralizedComponent]] = None, hook_alias_overrides: Optional[Dict[str, str]] = None, + mlp_reads_resid_directly: bool = False, ): """Initialize the block bridge. @@ -68,6 +69,10 @@ def __init__( hook_alias_overrides: Optional dictionary to override default hook aliases. For example, {"hook_attn_out": "ln1_post.hook_out"} will make hook_attn_out point to ln1_post.hook_out instead of the default attn.hook_out. + mlp_reads_resid_directly: True for post-norm blocks where the MLP consumes + the mid-residual with no pre-MLP norm (OLMo 2 layout). Moves the + hook_mlp_in capture from ln2 (whose input there is the raw MLP output) + to the MLP itself. """ # ln1_post/ln2_post redirect attn_out/mlp_out to match HookedTransformer's # placement (hook fires after the post-norm, not before). @@ -108,7 +113,9 @@ def __init__( self._pre_ln_capture_handles: list[torch.utils.hooks.RemovableHandle] = [] # Fallback for _read_use_hook_mlp_in when block.config is None. self._use_hook_mlp_in: bool = False - # Fires pre-ln2 when use_hook_mlp_in is set. See #1317. + self.mlp_reads_resid_directly = mlp_reads_resid_directly + # Fires on the MLP-branch entry (pre-ln2, or the MLP input on post-norm + # blocks) when use_hook_mlp_in is set. See #1317. self.hook_mlp_in = HookPoint() def _maybe_wire_pre_ln_capture(self) -> None: @@ -143,12 +150,21 @@ def _capture_pre_ln1(_module: torch.nn.Module, args: tuple) -> None: self._pre_ln_capture_handles.append(handle) attn._ln1_module = ln1.original_component - ln2 = self.submodules.get("ln2") if self.submodules else None - if ln2 is not None and getattr(ln2, "original_component", None) is not None: + # hook_mlp_in must capture the MLP-branch entry point: ln2's input on + # pre-norm blocks, the MLP's own input on post-norm blocks (where ln2 + # follows the MLP and its input is the raw MLP output). + if self.mlp_reads_resid_directly: + capture_target = self.submodules.get("mlp") if self.submodules else None + else: + capture_target = self.submodules.get("ln2") if self.submodules else None + if ( + capture_target is not None + and getattr(capture_target, "original_component", None) is not None + ): hook_mlp_in = self.hook_mlp_in block_ref = weakref.proxy(self) - def _capture_pre_ln2(_module: torch.nn.Module, args: tuple) -> Any: + def _capture_mlp_in(_module: torch.nn.Module, args: tuple) -> Any: if not block_ref._read_use_hook_mlp_in(): return None if args and isinstance(args[0], torch.Tensor): @@ -156,7 +172,7 @@ def _capture_pre_ln2(_module: torch.nn.Module, args: tuple) -> Any: return (hooked,) + args[1:] return None - handle = ln2.register_forward_pre_hook(_capture_pre_ln2) + handle = capture_target.register_forward_pre_hook(_capture_mlp_in) self._pre_ln_capture_handles.append(handle) self._pre_ln_capture_wired = True diff --git a/transformer_lens/model_bridge/supported_architectures/olmo2.py b/transformer_lens/model_bridge/supported_architectures/olmo2.py index cde00ed5ab..ff604e7171 100644 --- a/transformer_lens/model_bridge/supported_architectures/olmo2.py +++ b/transformer_lens/model_bridge/supported_architectures/olmo2.py @@ -90,12 +90,18 @@ def __init__(self, cfg: Any) -> None: ), "mlp": self._build_mlp_bridge(), }, - # Post-norm override: ln2 is post_feedforward_layernorm applied AFTER - # MLP, so "ln2.hook_in" captures the MLP output (wrong mid-point). - # The true residual mid-point (between attention and MLP) is mlp.hook_in. + # Post-norm overrides: ln1/ln2 are applied AFTER attention/MLP and + # BEFORE the residual add, so the residual mid-point is mlp.hook_in + # and the additive contributions are the norm outputs, not the raw + # module outputs (attn.hook_out / mlp.hook_out stay raw). hook_alias_overrides={ "hook_resid_mid": "mlp.hook_in", + "hook_attn_out": "ln1.hook_out", + "hook_mlp_out": "ln2.hook_out", }, + # No pre-MLP norm: the MLP consumes the mid-residual directly, so + # the hook_mlp_in capture must sit on the MLP, not on ln2. + mlp_reads_resid_directly=True, ), "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), diff --git a/transformer_lens/model_bridge/supported_architectures/olmo_hybrid.py b/transformer_lens/model_bridge/supported_architectures/olmo_hybrid.py index 47366fb90b..ff1a971860 100644 --- a/transformer_lens/model_bridge/supported_architectures/olmo_hybrid.py +++ b/transformer_lens/model_bridge/supported_architectures/olmo_hybrid.py @@ -13,6 +13,8 @@ from typing import Any +import torch.nn as nn + from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( AttentionBridge, @@ -28,11 +30,18 @@ class _OlmoHybridBlockBridge(BlockBridge): - """BlockBridge without the hook_resid_mid alias. - - No single target fits both layer types (ln2.hook_in is the mid-point on - linear-attention layers but the raw attn-branch output on full-attention - layers); dropped type-visibly, as on ParallelBlockBridge. + """BlockBridge with per-layer-type hook aliases and no hook_resid_mid. + + hook_resid_mid: no single target fits both layer types (ln2.hook_in is the + mid-point on linear-attention layers but the raw attn-branch output on + full-attention layers); dropped type-visibly, as on ParallelBlockBridge. + + hook_attn_out / hook_mlp_out must expose the tensor added to the residual + stream, which also differs by layer type: full-attention layers are OLMo2 + post-norm (contribution = norm output), linear-attention layers are pre-norm + (contribution = raw sublayer output). Both candidate targets exist on both + layer types, so alias fallback lists cannot discriminate; instead the + aliases are selected per layer at bind time in set_original_component. """ def __init__(self, *args: Any, **kwargs: Any): @@ -41,6 +50,25 @@ def __init__(self, *args: Any, **kwargs: Any): self.hook_aliases = dict(self.hook_aliases) self.hook_aliases.pop("hook_resid_mid", None) + def set_original_component(self, original_component: nn.Module) -> None: + super().set_original_component(original_component) + if self.hook_aliases is BlockBridge.hook_aliases: + self.hook_aliases = dict(self.hook_aliases) + if getattr(original_component, "post_feedforward_layernorm", None) is not None: + # Full-attention (post-norm) layer: ln2 = post_attention_layernorm + # applied after attention, ln2_post = post_feedforward_layernorm. + # The MLP consumes the mid-residual directly, so the hook_mlp_in + # capture must sit on the MLP, not on ln2 (whose input here is the + # raw attention output). + self.hook_aliases["hook_attn_out"] = "ln2.hook_out" + self.hook_aliases["hook_mlp_out"] = "ln2_post.hook_out" + self.mlp_reads_resid_directly = True + else: + # Linear-attention (pre-norm) layer. + self.hook_aliases["hook_attn_out"] = "linear_attn.hook_out" + self.hook_aliases["hook_mlp_out"] = "mlp.hook_out" + self.mlp_reads_resid_directly = False + class OlmoHybridArchitectureAdapter(ArchitectureAdapter): """Architecture adapter for OlmoHybridForCausalLM models.""" From 9ffcd8e3d7e81fa66f5a4c07f1e43782e72ccb01 Mon Sep 17 00:00:00 2001 From: jlarson4 Date: Wed, 12 Aug 2026 13:06:34 -0500 Subject: [PATCH 3/4] Granite resolution --- .../test_granite_hook_semantics.py | 264 ++++++++++++++++++ .../test_granite_moe_hybrid_adapter.py | 29 ++ .../model_bridge/test_olmo_hybrid_adapter.py | 4 +- .../test_granite_adapter.py | 25 ++ .../test_granite_moe_adapter.py | 12 + .../test_granite_moe_hybrid_adapter.py | 18 ++ .../test_hook_alias_resolution.py | 1 - transformer_lens/model_bridge/bridge.py | 4 +- .../generalized_components/__init__.py | 2 + .../generalized_components/block.py | 164 ++++++++++- .../model_bridge/sources/_bridge_builder.py | 1 + .../supported_architectures/granite.py | 8 +- .../supported_architectures/granite_moe.py | 7 +- .../granite_moe_hybrid.py | 15 +- .../supported_architectures/llada.py | 8 +- 15 files changed, 533 insertions(+), 29 deletions(-) create mode 100644 tests/integration/model_bridge/test_granite_hook_semantics.py diff --git a/tests/integration/model_bridge/test_granite_hook_semantics.py b/tests/integration/model_bridge/test_granite_hook_semantics.py new file mode 100644 index 0000000000..ea6a3f54bf --- /dev/null +++ b/tests/integration/model_bridge/test_granite_hook_semantics.py @@ -0,0 +1,264 @@ +"""Integration tests for Granite residual-branch hook semantics. + +Granite's HF blocks compute ``residual + sublayer_out * residual_multiplier``, +so hook_attn_out / hook_mlp_out must expose the scaled contribution (issue +#1648). The hub tiny-random Granite ships residual_multiplier=1.0, which cannot +catch a missing scale, so the fixtures build local checkpoints with 0.22 (the +granite-3.3 value). Norm-based comparisons throughout — the wrong tensor is +collinear with the right one, so cosine checks are blind here. +""" + +import pytest +import torch + +from transformer_lens.model_bridge import TransformerBridge + +RESIDUAL_MULTIPLIER = 0.22 +TOKENIZER_SOURCE = "hf-internal-testing/tiny-random-GraniteForCausalLM" + + +@pytest.fixture(scope="module") +def granite_path(tmp_path_factory): + from transformers import AutoTokenizer, GraniteConfig, GraniteForCausalLM + + tok = AutoTokenizer.from_pretrained(TOKENIZER_SOURCE) + torch.manual_seed(0) + cfg = GraniteConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=len(tok), + max_position_embeddings=512, + pad_token_id=tok.pad_token_id or 0, + residual_multiplier=RESIDUAL_MULTIPLIER, + ) + path = tmp_path_factory.mktemp("granite") / "tiny-granite" + model = GraniteForCausalLM(cfg).to(torch.float32) + model.save_pretrained(path) + tok.save_pretrained(path) + return str(path) + + +@pytest.fixture(scope="module") +def granite_bridge(granite_path) -> TransformerBridge: + return TransformerBridge.boot_transformers(granite_path, device="cpu", dtype=torch.float32) + + +@pytest.fixture(scope="module") +def sample_tokens(granite_bridge: TransformerBridge) -> torch.Tensor: + return granite_bridge.to_tokens("The capital of France is Paris.") + + +def test_forward_matches_fresh_hf(granite_bridge, granite_path, sample_tokens) -> None: + from transformers import AutoModelForCausalLM + + fresh = AutoModelForCausalLM.from_pretrained( + granite_path, dtype=torch.float32, attn_implementation="eager" + ) + fresh.eval() + with torch.no_grad(): + bridge_out = granite_bridge(sample_tokens) + hf_out = fresh(input_ids=sample_tokens).logits + max_diff = (bridge_out - hf_out).abs().max().item() + assert max_diff < 1e-5, f"Bridge vs fresh HF max diff = {max_diff}" + + +def test_residual_branch_hooks_decompose_stream(granite_bridge, sample_tokens) -> None: + with torch.no_grad(): + _, cache = granite_bridge.run_with_cache(sample_tokens) + + for layer in range(granite_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_is_scaled_raw_output(granite_bridge, sample_tokens) -> None: + """The contribution must be raw * residual_multiplier; raw stays on the + architecture-shaped hooks.""" + with torch.no_grad(): + _, cache = granite_bridge.run_with_cache(sample_tokens) + + for layer in range(granite_bridge.cfg.n_layers): + torch.testing.assert_close( + cache[f"blocks.{layer}.hook_attn_out"], + cache[f"blocks.{layer}.attn.hook_out"] * RESIDUAL_MULTIPLIER, + ) + torch.testing.assert_close( + cache[f"blocks.{layer}.hook_mlp_out"], + cache[f"blocks.{layer}.mlp.hook_out"] * RESIDUAL_MULTIPLIER, + ) + + +def test_read_only_hooks_leave_forward_bit_exact(granite_bridge, sample_tokens) -> None: + """Cache-style hooks must not perturb the forward — the rewrite (with its + divide-multiply rounding) only happens when a hook changes the tensor.""" + grabbed = {} + + def grab(tensor: torch.Tensor, hook) -> torch.Tensor: + grabbed[hook.name] = tensor.detach().clone() + return tensor + + with torch.no_grad(): + baseline = granite_bridge(sample_tokens) + hooked = granite_bridge.run_with_hooks( + sample_tokens, + fwd_hooks=[ + ("blocks.0.hook_attn_out", grab), + ("blocks.0.hook_mlp_out", grab), + ], + ) + + assert torch.equal(baseline, hooked) + assert "blocks.0.hook_attn_out" in grabbed + + +def test_attn_out_ablation_collapses_residual_step(granite_bridge, sample_tokens) -> None: + 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 = granite_bridge(sample_tokens) + ablated = granite_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(granite_bridge, sample_tokens) -> None: + """Writing v must make the contribution v itself — the multiplier may not be + applied on top of the write (the issue's Granite failure mode).""" + torch.manual_seed(1) + replacement = torch.randn(1, sample_tokens.shape[1], granite_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(): + granite_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_in_place_mutation_is_not_dropped(granite_bridge, sample_tokens) -> None: + """A hook that zeroes the tensor in place (returning the same object) must + ablate the contribution — identity checks would silently drop it.""" + captured = {} + + def grab(key: str): + def hook_fn(tensor: torch.Tensor, hook) -> torch.Tensor: + captured[key] = tensor.detach().clone() + return tensor + + return hook_fn + + def zero_in_place(tensor: torch.Tensor, hook) -> torch.Tensor: + tensor.zero_() + return tensor + + with torch.no_grad(): + granite_bridge.run_with_hooks( + sample_tokens, + fwd_hooks=[ + ("blocks.0.hook_attn_out", zero_in_place), + ("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"]) + + +def test_backward_hook_receives_gradient(granite_bridge, sample_tokens) -> None: + """Backward hooks must sit on the compute path (the rewrite must happen when + bwd hooks are attached, or they observe a dead branch).""" + received = [] + + def bwd_hook(grad, hook): + received.append(grad.detach().clone()) + return grad + + granite_bridge.add_hook("blocks.0.hook_attn_out", bwd_hook, dir="bwd") + try: + logits = granite_bridge(sample_tokens) + logits.sum().backward() + finally: + granite_bridge.reset_hooks() + + assert received and received[0].abs().sum() > 0 + + +class TestGraniteMoeHookSemantics: + """Same contract for GraniteMoe: the MoE output is scaled before the add.""" + + @pytest.fixture(scope="class") + def moe_bridge(self, tmp_path_factory) -> TransformerBridge: + from transformers import AutoTokenizer, GraniteMoeConfig, GraniteMoeForCausalLM + + tok = AutoTokenizer.from_pretrained(TOKENIZER_SOURCE) + torch.manual_seed(0) + cfg = GraniteMoeConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=len(tok), + max_position_embeddings=512, + pad_token_id=tok.pad_token_id or 0, + residual_multiplier=RESIDUAL_MULTIPLIER, + num_local_experts=4, + num_experts_per_tok=2, + ) + path = tmp_path_factory.mktemp("granite_moe") / "tiny-granite-moe" + model = GraniteMoeForCausalLM(cfg).to(torch.float32) + model.save_pretrained(path) + tok.save_pretrained(path) + return TransformerBridge.boot_transformers(str(path), device="cpu", dtype=torch.float32) + + def test_residual_branch_hooks_decompose_stream(self, moe_bridge) -> None: + tokens = moe_bridge.to_tokens("The capital of France is Paris.") + with torch.no_grad(): + _, cache = moe_bridge.run_with_cache(tokens) + + for layer in range(moe_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"], + ) diff --git a/tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py b/tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py index f02de21826..d26ca6e40e 100644 --- a/tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py +++ b/tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py @@ -128,6 +128,35 @@ def test_inner_norm_is_gated(self, bridge: TransformerBridge) -> None: assert isinstance(mixer.inner_norm, GatedRMSNormBridge) +class TestGraniteMoeHybridContributionHooks: + """hook_attn_out must exist and fire on attention layers only — a HookPoint + that exists but never fires is a silent-no-op intervention trap (#1648).""" + + def test_mamba_layers_have_no_hook_attn_out(self, bridge: TransformerBridge) -> None: + hooks = bridge.hook_dict + for i in MAMBA_LAYERS: + assert f"blocks.{i}.hook_attn_out" not in hooks, ( + f"block {i} is a mamba layer; a dead hook_attn_out would " + f"silently no-op interventions" + ) + assert f"blocks.{ATTN_LAYER}.hook_attn_out" in hooks + + def test_attention_layer_hook_attn_out_fires(self, bridge: TransformerBridge, tokens) -> None: + fired = {} + + def grab(tensor, hook): + fired[hook.name] = tensor.detach().clone() + return tensor + + with torch.no_grad(): + bridge.run_with_hooks( + tokens, + fwd_hooks=[(f"blocks.{ATTN_LAYER}.hook_attn_out", grab)], + ) + + assert f"blocks.{ATTN_LAYER}.hook_attn_out" in fired + + # --------------------------------------------------------------------------- # Forward parity: bridge delegates fully, so logits match HF exactly # --------------------------------------------------------------------------- diff --git a/tests/integration/model_bridge/test_olmo_hybrid_adapter.py b/tests/integration/model_bridge/test_olmo_hybrid_adapter.py index e69d66de57..5c5dab595f 100644 --- a/tests/integration/model_bridge/test_olmo_hybrid_adapter.py +++ b/tests/integration/model_bridge/test_olmo_hybrid_adapter.py @@ -137,9 +137,7 @@ def test_full_attention_contributions_are_post_norm(self, olmo_bridge, sample_to cache[f"blocks.{layer}.hook_mlp_out"], ) - def test_hook_mlp_in_exposes_mid_residual_on_both_layer_types( - self, olmo_bridge, sample_tokens - ): + 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).""" diff --git a/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py b/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py index 2603d0f56a..41e617cd91 100644 --- a/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py @@ -127,6 +127,31 @@ def test_bridge_types(self, adapter: GraniteArchitectureAdapter) -> None: assert isinstance(mapping["ln_final"], RMSNormalizationBridge) assert isinstance(mapping["unembed"], UnembeddingBridge) + def test_scaled_residual_block(self, adapter: GraniteArchitectureAdapter) -> None: + """hook_attn_out / hook_mlp_out fire on the scaled contribution (#1648): + real HookPoints on the block, not aliases to raw module outputs.""" + from transformer_lens.model_bridge.generalized_components import ( + ScaledResidualBlockBridge, + ) + + block = adapter.component_mapping["blocks"] + assert isinstance(block, ScaledResidualBlockBridge) + assert "hook_attn_out" not in block.hook_aliases + assert "hook_mlp_out" not in block.hook_aliases + + def test_scale_comes_from_residual_multiplier(self) -> None: + cfg = _make_cfg() + cfg.residual_multiplier = 0.22 + adapter = GraniteArchitectureAdapter(cfg) + assert adapter.component_mapping["blocks"].residual_contribution_scale == 0.22 + + def test_zero_residual_multiplier_rejected(self) -> None: + """The write path divides by the scale, so zero must fail loudly.""" + cfg = _make_cfg() + cfg.residual_multiplier = 0.0 + with pytest.raises(ValueError, match="nonzero"): + GraniteArchitectureAdapter(cfg) + def test_top_level_hf_paths(self, adapter: GraniteArchitectureAdapter) -> None: mapping = adapter.component_mapping assert mapping["embed"].name == "model.embed_tokens" diff --git a/tests/unit/model_bridge/supported_architectures/test_granite_moe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_granite_moe_adapter.py index 0f96d26f83..945e63210c 100644 --- a/tests/unit/model_bridge/supported_architectures/test_granite_moe_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_granite_moe_adapter.py @@ -157,6 +157,18 @@ def test_top_level_hf_paths(self, adapter: GraniteMoeArchitectureAdapter) -> Non assert mapping["ln_final"].name == "model.norm" assert mapping["unembed"].name == "lm_head" + def test_scaled_residual_block(self, adapter: GraniteMoeArchitectureAdapter) -> None: + """hook_attn_out / hook_mlp_out fire on the scaled contribution (#1648).""" + from transformer_lens.model_bridge.generalized_components import ( + ScaledResidualBlockBridge, + ) + + block = adapter.component_mapping["blocks"] + assert isinstance(block, ScaledResidualBlockBridge) + assert block.scaled_mlp_submodule == "mlp" + assert "hook_attn_out" not in block.hook_aliases + assert "hook_mlp_out" not in block.hook_aliases + def test_block_submodule_keys(self, adapter: GraniteMoeArchitectureAdapter) -> None: blocks = adapter.component_mapping["blocks"] assert set(blocks.submodules.keys()) == {"ln1", "ln2", "attn", "mlp"} diff --git a/tests/unit/model_bridge/supported_architectures/test_granite_moe_hybrid_adapter.py b/tests/unit/model_bridge/supported_architectures/test_granite_moe_hybrid_adapter.py index d798435063..c25e8228e3 100644 --- a/tests/unit/model_bridge/supported_architectures/test_granite_moe_hybrid_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_granite_moe_hybrid_adapter.py @@ -143,6 +143,24 @@ def test_top_level_mapping(self, adapter: GraniteMoeHybridArchitectureAdapter) - assert mapping["ln_final"].name == "model.norm" assert mapping["unembed"].name == "lm_head" + def test_scaled_residual_block_mlp_wiring(self) -> None: + """With experts the MLP branch is moe + shared_mlp summed inline — no + single module produces the contribution, so hook_mlp_out stays absent. + Without experts it wires to shared_mlp (#1648).""" + from transformer_lens.model_bridge.generalized_components import ( + ScaledResidualBlockBridge, + ) + + with_experts = GraniteMoeHybridArchitectureAdapter(_make_cfg(num_experts=4)) + blocks = with_experts.component_mapping["blocks"] + assert isinstance(blocks, ScaledResidualBlockBridge) + assert blocks.scaled_mlp_submodule is None + assert not hasattr(blocks, "hook_mlp_out") + assert "hook_mlp_out" not in blocks.hook_aliases + + without_experts = GraniteMoeHybridArchitectureAdapter(_make_cfg(num_experts=0)) + assert without_experts.component_mapping["blocks"].scaled_mlp_submodule == "shared_mlp" + def test_block_submodule_mapping(self, adapter: GraniteMoeHybridArchitectureAdapter) -> None: blocks = adapter.component_mapping["blocks"] assert set(blocks.submodules.keys()) == { diff --git a/tests/unit/model_bridge/test_hook_alias_resolution.py b/tests/unit/model_bridge/test_hook_alias_resolution.py index 3dc30ef2bd..5d9fed376d 100644 --- a/tests/unit/model_bridge/test_hook_alias_resolution.py +++ b/tests/unit/model_bridge/test_hook_alias_resolution.py @@ -76,7 +76,6 @@ def _resolve(component: GeneralizedComponent, target: str) -> Any: "Gemma3ForConditionalGeneration": "audit H15 — multimodal vision encoder opaque", "Idefics3ForConditionalGeneration": "vision-encoder layer submodules unwired (same Siglip opacity as Llava/Gemma3 multimodal)", "OpenELMForCausalLM": "audit H23 — per-layer head counts break uniform q/k/v shape", - "GraniteMoeHybridForCausalLM": "new finding — MoE+shared-MLP block lacks proper submodule aliases", } diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index f1d758fe93..9a88bb7d4e 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -906,8 +906,8 @@ def set_compatibility_mode(component: Any) -> None: # Drop pre-ln capture handles from any prior call so they don't accumulate. if hasattr(self, "blocks"): for block in self.blocks: - if hasattr(block, "_teardown_pre_ln_capture"): - block._teardown_pre_ln_capture() + if hasattr(block, "_teardown_capture_hooks"): + block._teardown_capture_hooks() try: if not no_processing: self.process_weights( diff --git a/transformer_lens/model_bridge/generalized_components/__init__.py b/transformer_lens/model_bridge/generalized_components/__init__.py index 00225510b6..dd3ebb10b5 100644 --- a/transformer_lens/model_bridge/generalized_components/__init__.py +++ b/transformer_lens/model_bridge/generalized_components/__init__.py @@ -17,6 +17,7 @@ DelegatedAttentionBlockBridge, MLABlockBridge, ParallelBlockBridge, + ScaledResidualBlockBridge, ) from transformer_lens.model_bridge.generalized_components.bloom_attention import ( BloomAttentionBridge, @@ -152,6 +153,7 @@ "Lfm2ShortConvBridge", "MLABlockBridge", "ParallelBlockBridge", + "ScaledResidualBlockBridge", "BloomBlockBridge", "BloomAttentionBridge", "CodeGenAttentionBridge", diff --git a/transformer_lens/model_bridge/generalized_components/block.py b/transformer_lens/model_bridge/generalized_components/block.py index e983dcd6c1..721e14cd89 100644 --- a/transformer_lens/model_bridge/generalized_components/block.py +++ b/transformer_lens/model_bridge/generalized_components/block.py @@ -109,8 +109,8 @@ def __init__( ) self._original_block_forward: Optional[Callable[..., Any]] = None - self._pre_ln_capture_wired: bool = False - self._pre_ln_capture_handles: list[torch.utils.hooks.RemovableHandle] = [] + self._capture_hooks_wired: bool = False + self._capture_hook_handles: list[torch.utils.hooks.RemovableHandle] = [] # Fallback for _read_use_hook_mlp_in when block.config is None. self._use_hook_mlp_in: bool = False self.mlp_reads_resid_directly = mlp_reads_resid_directly @@ -118,7 +118,7 @@ def __init__( # blocks) when use_hook_mlp_in is set. See #1317. self.hook_mlp_in = HookPoint() - def _maybe_wire_pre_ln_capture(self) -> None: + def _maybe_wire_capture_hooks(self) -> None: """Install ln1/ln2 forward_pre_hooks that feed the bridge's pre-LN hooks (#1317). Hooks register on the NormalizationBridge instance, not on @@ -126,7 +126,7 @@ def _maybe_wire_pre_ln_capture(self) -> None: forward never calls the raw module, so a hook there would silently miss on most adapters. Idempotent. """ - if self._pre_ln_capture_wired: + if self._capture_hooks_wired: return from transformer_lens.model_bridge.generalized_components.attention import ( AttentionBridge, @@ -147,7 +147,7 @@ def _capture_pre_ln1(_module: torch.nn.Module, args: tuple) -> None: attn_ref._captured_pre_ln_residual = args[0] handle = ln1.register_forward_pre_hook(_capture_pre_ln1) - self._pre_ln_capture_handles.append(handle) + self._capture_hook_handles.append(handle) attn._ln1_module = ln1.original_component # hook_mlp_in must capture the MLP-branch entry point: ln2's input on @@ -173,16 +173,16 @@ def _capture_mlp_in(_module: torch.nn.Module, args: tuple) -> Any: return None handle = capture_target.register_forward_pre_hook(_capture_mlp_in) - self._pre_ln_capture_handles.append(handle) + self._capture_hook_handles.append(handle) - self._pre_ln_capture_wired = True + self._capture_hooks_wired = True - def _teardown_pre_ln_capture(self) -> None: - """Remove the ln1/ln2 forward_pre_hooks installed by _maybe_wire_pre_ln_capture.""" - for handle in self._pre_ln_capture_handles: + def _teardown_capture_hooks(self) -> None: + """Remove the ln1/ln2 forward_pre_hooks installed by _maybe_wire_capture_hooks.""" + for handle in self._capture_hook_handles: handle.remove() - self._pre_ln_capture_handles.clear() - self._pre_ln_capture_wired = False + self._capture_hook_handles.clear() + self._capture_hooks_wired = False def _read_use_hook_mlp_in(self) -> bool: """Prefer ``block.config.use_hook_mlp_in``; fall back to the block-local flag.""" @@ -209,7 +209,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: f"Original component not set for {self.name}. Call set_original_component() first." ) - self._maybe_wire_pre_ln_capture() + self._maybe_wire_capture_hooks() self._check_stop_at_layer(*args, **kwargs) args, kwargs = self._hook_input_hidden_states(args, kwargs) @@ -424,6 +424,144 @@ def __init__( self.hook_aliases.pop("hook_resid_mid", None) +class ScaledResidualBlockBridge(BlockBridge): + """Block whose sublayer outputs are scaled before the residual add. + + Granite-family HF blocks compute ``residual + sublayer_out * residual_multiplier`` + inline, so no submodule output equals the tensor added to the residual stream + and the legacy aliases cannot be fixed by re-pointing. ``hook_attn_out`` / + ``hook_mlp_out`` become real HookPoints on the block firing on the scaled + contribution: + + - no hooks attached: the forward is untouched (bit-exact); + - read-only hooks: they observe ``module_out * scale``, forward stays bit-exact; + - a hook that changes the tensor (returned new or mutated in place): the module + output is rewritten to ``hooked / scale`` so HF's multiply reconstructs the + written value as the contribution (~1-ulp rounding; exact for zero-ablation); + - any backward hooks: the rewrite always happens so the autograd graph routes + through the HookPoint. + """ + + def __init__( + self, + name: str, + config: Optional[Any] = None, + submodules: Optional[Dict[str, GeneralizedComponent]] = None, + hook_alias_overrides: Optional[Dict[str, str]] = None, + mlp_reads_resid_directly: bool = False, + residual_contribution_scale: float = 1.0, + scaled_attn_submodule: str = "attn", + scaled_mlp_submodule: Optional[str] = "mlp", + ): + """Initialize the scaled-residual block bridge. + + Args: + residual_contribution_scale: The HF block's residual multiplier. + scaled_attn_submodule: Submodule whose output feeds the attention-side + residual add. + scaled_mlp_submodule: Submodule whose output feeds the MLP-side residual + add, or None when no single submodule produces it (e.g. GraniteMoeHybrid + with experts sums ``block_sparse_moe + shared_mlp`` inline) — then + hook_mlp_out stays absent rather than firing with a partial tensor. + """ + super().__init__( + name, + config=config, + submodules=submodules, + hook_alias_overrides=hook_alias_overrides, + mlp_reads_resid_directly=mlp_reads_resid_directly, + ) + scale = float(residual_contribution_scale) + if scale == 0.0: + raise ValueError( + f"ScaledResidualBlockBridge at '{name}': residual_contribution_scale " + f"must be nonzero (the write path divides by it)." + ) + self.residual_contribution_scale = scale + self.scaled_attn_submodule = scaled_attn_submodule + self.scaled_mlp_submodule = scaled_mlp_submodule + if self.hook_aliases is BlockBridge.hook_aliases: + self.hook_aliases = dict(self.hook_aliases) + for alias in ("hook_attn_out", "hook_mlp_out"): + self.hook_aliases.pop(alias, None) + self.hook_attn_out = HookPoint() + if scaled_mlp_submodule is not None: + self.hook_mlp_out = HookPoint() + + def set_original_component(self, original_component: torch.nn.Module) -> None: + """Prune contribution HookPoints the bound layer cannot fire. + + Heterogeneous blocks (GraniteMoeHybrid mamba layers) lack the attn + submodule; a HookPoint that exists but never fires is a silent-no-op + intervention trap, so the name must be absent instead. + """ + super().set_original_component(original_component) + targets = [(self.scaled_attn_submodule, "hook_attn_out")] + if self.scaled_mlp_submodule is not None: + targets.append((self.scaled_mlp_submodule, "hook_mlp_out")) + for sub_name, hook_name in targets: + sub = self.submodules.get(sub_name) if self.submodules else None + remote = getattr(sub, "name", None) + first = remote.split(".", 1)[0] if isinstance(remote, str) else None + missing = sub is None or ( + first is not None and getattr(original_component, first, None) is None + ) + if missing: + if hasattr(self, hook_name): + delattr(self, hook_name) + # __setattr__ auto-registered the HookPoint; get_hooks() serves + # from this registry, so the module deletion alone is not enough. + self._hook_registry.pop(hook_name, None) + + def _maybe_wire_capture_hooks(self) -> None: + """Extend the base wiring with the scaled-contribution forward hooks. + + Shares the base flag and handle list so _teardown_capture_hooks also + removes these hooks and re-wiring stays idempotent. + """ + if self._capture_hooks_wired: + return + super()._maybe_wire_capture_hooks() + targets = [(self.scaled_attn_submodule, "hook_attn_out")] + if self.scaled_mlp_submodule is not None: + targets.append((self.scaled_mlp_submodule, "hook_mlp_out")) + for sub_name, hook_name in targets: + hook_point = getattr(self, hook_name, None) + sub = self.submodules.get(sub_name) if self.submodules else None + if ( + hook_point is None + or sub is None + or getattr(sub, "original_component", None) is None + ): + continue + handle = sub.register_forward_hook(self._make_scaled_contribution_hook(hook_point)) + self._capture_hook_handles.append(handle) + + def _make_scaled_contribution_hook( + self, hook_point: HookPoint + ) -> Callable[[torch.nn.Module, tuple, Any], Any]: + """Build a forward hook exposing ``output * scale`` through hook_point.""" + scale = self.residual_contribution_scale + + def _hook(_module: torch.nn.Module, _args: tuple, output: Any) -> Any: + if not hook_point.has_hooks(dir="both"): + return None + is_tuple = isinstance(output, tuple) + out = output[0] if is_tuple else output + if not isinstance(out, torch.Tensor): + return None + scaled = out * scale + hooked = hook_point(scaled) + # Compare against a freshly computed reference: an in-place mutation + # alters `scaled` itself, so an identity check would miss it. + if not hook_point.has_hooks(dir="bwd") and torch.equal(hooked, out * scale): + return None + new = hooked / scale + return ((new,) + output[1:]) if is_tuple else new + + return _hook + + class DelegatedAttentionBlockBridge(BlockBridge): """Block whose attention is delegated wholesale to HF (no split-qkv fork). diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index b8cbecd225..b6ba321c01 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -35,6 +35,7 @@ # Granite "position_embedding_type", "logits_scaling", + "residual_multiplier", # Falcon "parallel_attn", "multi_query", diff --git a/transformer_lens/model_bridge/supported_architectures/granite.py b/transformer_lens/model_bridge/supported_architectures/granite.py index 498451cfb4..47324fbe76 100644 --- a/transformer_lens/model_bridge/supported_architectures/granite.py +++ b/transformer_lens/model_bridge/supported_architectures/granite.py @@ -10,13 +10,13 @@ from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( - BlockBridge, EmbeddingBridge, GatedMLPBridge, LinearBridge, PositionEmbeddingsAttentionBridge, RMSNormalizationBridge, RotaryEmbeddingBridge, + ScaledResidualBlockBridge, UnembeddingBridge, ) @@ -78,7 +78,10 @@ def _build_component_mapping(self) -> dict: return { "embed": EmbeddingBridge(name="model.embed_tokens"), "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), - "blocks": BlockBridge( + # HF multiplies each sublayer output by residual_multiplier before the + # residual add, so hook_attn_out / hook_mlp_out must expose the scaled + # contribution, not the raw module output. + "blocks": ScaledResidualBlockBridge( name="model.layers", submodules={ "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), @@ -86,6 +89,7 @@ def _build_component_mapping(self) -> dict: "attn": self._build_attention_bridge(), "mlp": self._build_mlp_bridge(), }, + residual_contribution_scale=getattr(self.cfg, "residual_multiplier", 1.0), ), "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), diff --git a/transformer_lens/model_bridge/supported_architectures/granite_moe.py b/transformer_lens/model_bridge/supported_architectures/granite_moe.py index 5db23df313..5dadaeca30 100644 --- a/transformer_lens/model_bridge/supported_architectures/granite_moe.py +++ b/transformer_lens/model_bridge/supported_architectures/granite_moe.py @@ -1,11 +1,11 @@ """Granite MoE architecture adapter.""" from transformer_lens.model_bridge.generalized_components import ( - BlockBridge, EmbeddingBridge, MoEBridge, RMSNormalizationBridge, RotaryEmbeddingBridge, + ScaledResidualBlockBridge, UnembeddingBridge, ) from transformer_lens.model_bridge.supported_architectures.granite import ( @@ -26,7 +26,9 @@ def _build_component_mapping(self) -> dict: return { "embed": EmbeddingBridge(name="model.embed_tokens"), "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), - "blocks": BlockBridge( + # HF multiplies each sublayer output by residual_multiplier before the + # residual add; hook_attn_out / hook_mlp_out expose the scaled contribution. + "blocks": ScaledResidualBlockBridge( name="model.layers", submodules={ "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), @@ -37,6 +39,7 @@ def _build_component_mapping(self) -> dict: config=self.cfg, ), }, + residual_contribution_scale=getattr(self.cfg, "residual_multiplier", 1.0), ), "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), diff --git a/transformer_lens/model_bridge/supported_architectures/granite_moe_hybrid.py b/transformer_lens/model_bridge/supported_architectures/granite_moe_hybrid.py index 910ee4f4d0..01dc475481 100644 --- a/transformer_lens/model_bridge/supported_architectures/granite_moe_hybrid.py +++ b/transformer_lens/model_bridge/supported_architectures/granite_moe_hybrid.py @@ -15,7 +15,6 @@ from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( - BlockBridge, EmbeddingBridge, GatedRMSNormBridge, LinearBridge, @@ -23,6 +22,7 @@ MoEBridge, RMSNormalizationBridge, RotaryEmbeddingBridge, + ScaledResidualBlockBridge, SSM2MixerBridge, UnembeddingBridge, ) @@ -104,9 +104,20 @@ def _build_component_mapping(self) -> dict: config=self.cfg, ) + # HF multiplies each sublayer output by residual_multiplier before the + # residual add. hook_attn_out fires on attention layers (mamba layers have + # no attention-position hook). With experts, the MLP branch is + # block_sparse_moe + shared_mlp summed inline — no single module produces + # the contribution, so hook_mlp_out stays absent; without experts it fires + # on the scaled shared_mlp output. mapping: dict = { "embed": EmbeddingBridge(name="model.embed_tokens"), - "blocks": BlockBridge(name="model.layers", submodules=block_submodules), + "blocks": ScaledResidualBlockBridge( + name="model.layers", + submodules=block_submodules, + residual_contribution_scale=getattr(self.cfg, "residual_multiplier", 1.0), + scaled_mlp_submodule=None if (num_experts and num_experts > 0) else "shared_mlp", + ), "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), } diff --git a/transformer_lens/model_bridge/supported_architectures/llada.py b/transformer_lens/model_bridge/supported_architectures/llada.py index 2ef605679f..36eae41a65 100644 --- a/transformer_lens/model_bridge/supported_architectures/llada.py +++ b/transformer_lens/model_bridge/supported_architectures/llada.py @@ -190,15 +190,15 @@ def _route_pre_mlp_norm( return (self.hook_mlp_in(args[0]),) + args[1:] return None - def _maybe_wire_pre_ln_capture(self) -> None: + def _maybe_wire_capture_hooks(self) -> None: """Use deepcopy-safe bound hooks for LLaDA's pre-MLP residual hook.""" - if self._pre_ln_capture_wired: + if self._capture_hooks_wired: return if self.ln2.original_component is not None: - self._pre_ln_capture_handles.append( + self._capture_hook_handles.append( self.ln2.register_forward_pre_hook(self._route_pre_mlp_norm) ) - self._pre_ln_capture_wired = True + self._capture_hooks_wired = True def _wire_llada_container_hooks(self) -> None: if self._llada_container_hooks_wired: From 23538470560c0399e42b4b5ad53728368dc33b2f Mon Sep 17 00:00:00 2001 From: jlarson4 Date: Wed, 12 Aug 2026 13:13:41 -0500 Subject: [PATCH 4/4] comment cleanup --- docs/source/content/compatibility_mode.md | 5 +++++ transformer_lens/model_bridge/bridge.py | 5 +++-- .../model_bridge/generalized_components/block.py | 11 +++++++---- .../model_bridge/supported_architectures/granite.py | 4 +++- 4 files changed, 18 insertions(+), 7 deletions(-) diff --git a/docs/source/content/compatibility_mode.md b/docs/source/content/compatibility_mode.md index 083c497888..0b8dbf6c0f 100644 --- a/docs/source/content/compatibility_mode.md +++ b/docs/source/content/compatibility_mode.md @@ -73,6 +73,11 @@ One consequence for head-level direct logit attribution: per-head contributions `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. +The Granite family has the same raw-vs-contribution split for a different reason: +HF scales each sublayer output by `residual_multiplier` before the residual add, so +`hook_attn_out` / `hook_mlp_out` fire on the scaled contribution while +`attn.hook_out` / `mlp.hook_out` stay raw — `attn.hook_result` sums to the raw +output, off from the contribution by the multiplier. 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`. diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 9a88bb7d4e..2e89f497dc 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -903,7 +903,7 @@ def set_compatibility_mode(component: Any) -> None: apply_fn_to_all_components(self, set_compatibility_mode) self.clear_hook_registry() - # Drop pre-ln capture handles from any prior call so they don't accumulate. + # Drop block capture-hook handles from any prior call so they don't accumulate. if hasattr(self, "blocks"): for block in self.blocks: if hasattr(block, "_teardown_capture_hooks"): @@ -4149,7 +4149,8 @@ def set_use_attn_in(self, use_attn_in: bool): self._propagate_attention_flag("use_attn_in", use_attn_in) def set_use_hook_mlp_in(self, use_hook_mlp_in: bool) -> None: - """Toggle the pre-ln2 ``hook_mlp_in`` HookPoint, matching legacy semantics. + """Toggle the ``hook_mlp_in`` HookPoint (the MLP-branch entry: pre-ln2, or + the MLP input on post-norm blocks), matching legacy semantics. See :py:meth:`HookedTransformer.set_use_hook_mlp_in`. """ diff --git a/transformer_lens/model_bridge/generalized_components/block.py b/transformer_lens/model_bridge/generalized_components/block.py index 721e14cd89..79c7b73600 100644 --- a/transformer_lens/model_bridge/generalized_components/block.py +++ b/transformer_lens/model_bridge/generalized_components/block.py @@ -39,7 +39,8 @@ class BlockBridge(GeneralizedComponent): is_list_item: bool = True hook_out_is_single_residual_stream: bool = True # hook_mlp_in is a direct HookPoint on this class (not aliased) so it can - # fire pre-ln2; see __init__. The post-ln2 mlp input stays at block.mlp.hook_in. + # fire on the MLP-branch entry (pre-ln2, or the MLP input on post-norm + # blocks); see __init__. The normalized mlp input stays at block.mlp.hook_in. hook_aliases = { "hook_resid_pre": "hook_in", "hook_resid_mid": "ln2.hook_in", @@ -119,9 +120,11 @@ def __init__( self.hook_mlp_in = HookPoint() def _maybe_wire_capture_hooks(self) -> None: - """Install ln1/ln2 forward_pre_hooks that feed the bridge's pre-LN hooks (#1317). + """Install the block's capture hooks: the ln1 pre-hook feeding the + split-qkv fork and the MLP-branch-entry pre-hook feeding hook_mlp_in + (#1317). Subclasses extend this with their own captures. - Hooks register on the NormalizationBridge instance, not on + Hooks register on the bridge submodule instance, not on ``original_component`` — the manual (non-native-autograd) bridge forward never calls the raw module, so a hook there would silently miss on most adapters. Idempotent. @@ -178,7 +181,7 @@ def _capture_mlp_in(_module: torch.nn.Module, args: tuple) -> Any: self._capture_hooks_wired = True def _teardown_capture_hooks(self) -> None: - """Remove the ln1/ln2 forward_pre_hooks installed by _maybe_wire_capture_hooks.""" + """Remove the capture hooks installed by _maybe_wire_capture_hooks (and subclass extensions).""" for handle in self._capture_hook_handles: handle.remove() self._capture_hook_handles.clear() diff --git a/transformer_lens/model_bridge/supported_architectures/granite.py b/transformer_lens/model_bridge/supported_architectures/granite.py index 47324fbe76..79c4365c34 100644 --- a/transformer_lens/model_bridge/supported_architectures/granite.py +++ b/transformer_lens/model_bridge/supported_architectures/granite.py @@ -26,7 +26,9 @@ class GraniteArchitectureAdapter(ArchitectureAdapter): Granite is a Llama-like architecture with RMSNorm, rotary position embeddings (RoPE), GQA, and a gated MLP (SiLU activation). Granite-specific scaling - multipliers are handled by the HF model's native forward pass. + multipliers are applied by the HF model's native forward pass; + ScaledResidualBlockBridge accounts for residual_multiplier so + hook_attn_out / hook_mlp_out expose the scaled residual contributions. Optional Parameters (may not exist in state_dict): -------------------------------------------------