diff --git a/docs/source/content/compatibility_mode.md b/docs/source/content/compatibility_mode.md index 083c49788..0b8dbf6c0 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/tests/integration/model_bridge/test_granite_hook_semantics.py b/tests/integration/model_bridge/test_granite_hook_semantics.py new file mode 100644 index 000000000..ea6a3f54b --- /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 f02de2182..d26ca6e40 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/unit/model_bridge/supported_architectures/test_granite_adapter.py b/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py index 2603d0f56..41e617cd9 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 0f96d26f8..945e63210 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 d79843506..c25e8228e 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 3dc30ef2b..5d9fed376 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 f1d758fe9..2e89f497d 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -903,11 +903,11 @@ 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_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( @@ -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/__init__.py b/transformer_lens/model_bridge/generalized_components/__init__.py index 00225510b..dd3ebb10b 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 e983dcd6c..79c7b7360 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", @@ -109,8 +110,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,15 +119,17 @@ 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: - """Install ln1/ln2 forward_pre_hooks that feed the bridge's pre-LN hooks (#1317). + def _maybe_wire_capture_hooks(self) -> None: + """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. """ - if self._pre_ln_capture_wired: + if self._capture_hooks_wired: return from transformer_lens.model_bridge.generalized_components.attention import ( AttentionBridge, @@ -147,7 +150,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 +176,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 capture hooks installed by _maybe_wire_capture_hooks (and subclass extensions).""" + 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 +212,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 +427,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 b8cbecd22..b6ba321c0 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 498451cfb..79c4365c3 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, ) @@ -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): ------------------------------------------------- @@ -78,7 +80,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 +91,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 5db23df31..5dadaeca3 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 910ee4f4d..01dc47548 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 2ef605679..36eae41a6 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: