diff --git a/tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py b/tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py index b31720e02..efe527ca0 100644 --- a/tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py +++ b/tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py @@ -2,7 +2,8 @@ Blocks without attention (gated-delta-net, Mamba) expose their sequence-mixer output instead of ``hook_attn_out``; blocks without an MLP (Mamba, NemotronH) -write only through their mixer. Every decomposition must still sum to the +write only through their mixer. Parallel hybrids (Falcon-H1) add a Mamba output +beside attention in every block. Every decomposition must still sum to the residual stream it decomposes. """ from unittest.mock import MagicMock @@ -169,6 +170,51 @@ def _granite_dense(): return _boot(AutoModelForCausalLM.from_config(c).eval(), "GraniteMoeHybridForCausalLM") +def _falcon_h1(): + from transformers import AutoModelForCausalLM + from transformers.models.falcon_h1 import FalconH1Config + + torch.manual_seed(0) + c = FalconH1Config( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=64, + max_position_embeddings=64, + mamba_d_ssm=128, + mamba_n_heads=4, + mamba_d_head=32, + mamba_n_groups=1, + mamba_d_state=16, + mamba_d_conv=4, + mamba_chunk_size=8, + attention_bias=True, + mlp_bias=True, + mamba_proj_bias=True, + projectors_bias=True, + # Released checkpoints scale every branch; 1.0 would hide a missing multiplier. + embedding_multiplier=2.0, + attention_in_multiplier=1.1, + attention_out_multiplier=0.7, + key_multiplier=0.8, + ssm_in_multiplier=1.4, + ssm_out_multiplier=1.9, + ssm_multipliers=[0.5, 1.5, 0.8, 1.2, 0.9], + mlp_multipliers=[1.3, 0.6], + pad_token_id=0, + ) + c.architectures = ["FalconH1ForCausalLM"] + model = AutoModelForCausalLM.from_config(c).eval() + with torch.no_grad(): + for param in model.parameters(): + if param.ndim == 1: + param.normal_(std=0.3) # nonzero biases and norm weights + return _boot(model, "FalconH1ForCausalLM") + + def _zamba2(): """Zamba2 with hybrid layers: the shared transformer feeds the Mamba-2 input of layers 1 and 3 but does not write to the residual stream.""" @@ -250,6 +296,7 @@ def _lfm2(): _FACTORIES = { + "falcon_h1": _falcon_h1, "granite_dense": _granite_dense, "lfm2": _lfm2, "mamba2": _mamba2, @@ -307,6 +354,7 @@ def test_labels_name_the_mixer(self, bridge_and_cache): name, bridge, cache = bridge_and_cache _, labels = cache.decompose_resid(return_labels=True) expected = { + "falcon_h1": ["0_attn_out", "0_mamba_out", "0_mlp_out", "1_attn_out", "1_mamba_out"], "qwen3_5": ["0_linear_attn_out", "0_mlp_out", "3_attn_out", "3_mlp_out"], "granite_dense": ["0_mixer_out", "0_mlp_out", "1_attn_out", "1_mlp_out"], "lfm2": ["0_conv_out", "0_mlp_out", "1_attn_out", "1_mlp_out", "2_conv_out"], @@ -346,14 +394,29 @@ def test_mlp_type_mixers_fill_the_mlp_slot(self, bridge_and_cache): def test_mlp_input_sums_to_resid_mid(self, bridge_and_cache): name, bridge, cache = bridge_and_cache - if name not in ("qwen3_5", "recurrent_gemma"): - pytest.skip("only Qwen3.5 and RecurrentGemma blocks have a separate MLP input") - for layer in (0, 2, 3): + if name not in ("qwen3_5", "falcon_h1", "recurrent_gemma"): + pytest.skip("Qwen3.5, RecurrentGemma and Falcon-H1 blocks have a separate MLP input") + for layer in (0, 3) if name == "qwen3_5" else (0, 2): comps = cache.decompose_resid(layer=layer, mlp_input=True) torch.testing.assert_close( comps.sum(0), cache[f"blocks.{layer}.hook_resid_mid"], atol=ATOL, rtol=0 ) + def test_parallel_branches_sum_to_their_residual_add(self, bridge_and_cache): + name, bridge, cache = bridge_and_cache + if name != "falcon_h1": + pytest.skip("only Falcon-H1 runs a mixer beside attention") + ssm_out_multiplier = bridge.cfg.ssm_out_multiplier + for layer in range(bridge.cfg.n_layers): + added = ( + cache[f"blocks.{layer}.hook_resid_mid"] - cache[f"blocks.{layer}.hook_resid_pre"] + ) + branches = ( + cache[f"blocks.{layer}.hook_attn_out"] + + cache[f"blocks.{layer}.mamba.hook_out"] * ssm_out_multiplier + ) + torch.testing.assert_close(branches, added, atol=ATOL, rtol=0) + def test_logit_attrs_runs(self, bridge_and_cache): _, bridge, cache = bridge_and_cache attrs = cache.logit_attrs(cache.decompose_resid(), tokens=torch.tensor([3, 3])) @@ -379,7 +442,9 @@ def test_head_results_sum_to_attn_out(self, bridge_and_cache): heads = cache[f"blocks.{layer}.attn.hook_result"].sum(-2) b_O = bridge.blocks[layer].attn.b_O if b_O is not None: - heads = heads + b_O + heads = heads + b_O * getattr( + bridge.blocks[layer], "residual_contribution_scale", 1.0 + ) torch.testing.assert_close( heads, cache[f"blocks.{layer}.hook_attn_out"], atol=ATOL, rtol=0 ) diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index 2fa617dc6..7b244ac61 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -402,20 +402,40 @@ def _has_attn(self, layer: int) -> bool: ) def _mixer_out_names(self, layer: int) -> List[str]: - """Variant slots (``linear_attn``, ``mixer``, …) of a block without attention. + """Variant slots (``linear_attn``, ``mixer``, …) that write outside ``hook_attn_out``. Hybrid and SSM blocks write their sequence-mixer output in place of - ``hook_attn_out``; empty for attention blocks. + ``hook_attn_out``. Parallel hybrids (Falcon-H1) run a mixer beside attention + and add both outputs to the residual stream; that mixer is listed too. """ from transformer_lens.model_bridge.generalized_components.block import ( VARIANT_SUBMODULE_NAMES, ) - if self._has_attn(layer): - return [] modules = self._block_modules(layer) or {} + if self._has_attn(layer): + if "attn" not in modules: + return [] + return [name for name in VARIANT_SUBMODULE_NAMES if name != "attn" and name in modules] return [name for name in VARIANT_SUBMODULE_NAMES if name in modules] + @staticmethod + def _mlp_contribution_scale(block: Any) -> float: + """Multiplier HF applies to the MLP output before the residual add.""" + return getattr( + block, "mlp_contribution_scale", getattr(block, "residual_contribution_scale", 1.0) + ) + + @classmethod + def _neuron_contribution_scale(cls, block: Any) -> float: + """Multiplier from ``post @ W_out + b_out`` to the residual contribution. + + Adds any multiplier inside the MLP after ``W_out`` (Falcon-H1's down multiplier), + which ``hook_mlp_out`` already includes but the neuron expansion does not. + """ + mlp = block._modules.get("mlp") if isinstance(block._modules, dict) else None + return cls._mlp_contribution_scale(block) * getattr(mlp, "_neuron_contribution_scale", 1.0) + def _has_mlp(self, layer: int) -> bool: """Whether block ``layer`` writes an MLP output separate from its mixer.""" from transformer_lens.model_bridge.generalized_components.block import ( @@ -470,24 +490,28 @@ def _mixer_stack( """Residual contributions of a non-attention block's mixer, e.g. ``0_linear_attn_out``. Scaled by the block's ``residual_contribution_scale`` (Granite's ``residual_multiplier`` - is applied after ``hook_out``), matching the head and neuron paths. ``subtract_bias`` - drops the mixer's output bias first, for decompositions that carry it in their own - ``bias`` row. + is applied after ``hook_out``), matching the head and neuron paths, or by the mixer's + own entry in ``mixer_contribution_scales`` (Falcon-H1's ``ssm_out_multiplier``). + ``subtract_bias`` drops the mixer's output bias first, for decompositions that carry it + in their own ``bias`` row. """ components: List[torch.Tensor] = [] labels: List[str] = [] names = self._mixer_out_names(layer) block = self.model.blocks[layer] - scale = getattr(block, "residual_contribution_scale", 1.0) - # Same bias accumulated_bias counts for this block, so the two cancel exactly. + default_scale = getattr(block, "residual_contribution_scale", 1.0) + mixer_scales = getattr(block, "mixer_contribution_scales", {}) + # Same bias accumulated_bias counts for this block, so the two cancel exactly. On a + # parallel hybrid that bias is the attention's, so the mixer keeps its own. bias = None - if subtract_bias and names: + if subtract_bias and names and not self._has_attn(layer): get_bias = getattr(self.model, "_get_block_variant_bias", None) bias = get_bias(block) if callable(get_bias) else None for i, name in enumerate(names): out = pos_slice.apply(self[f"blocks.{layer}.{name}.hook_out"], dim=-2) if i == 0 and isinstance(bias, torch.Tensor): out = out - bias.to(out.device, out.dtype) + scale = mixer_scales.get(name, default_scale) if scale != 1.0: out = out * scale components.append(out) @@ -845,7 +869,8 @@ def decompose_resid( attention components. Defaults to "all". On SSM / hybrid models, blocks without attention contribute their sequence-mixer output (e.g. ``0_linear_attn_out``, ``0_mixer_out``) in the attention slot, and blocks without an MLP contribute no - ``mlp_out``. + ``mlp_out``. Parallel hybrids (Falcon-H1) add their mixer output + (``0_mamba_out``) after ``0_attn_out``. apply_ln: Whether to apply LayerNorm to the stack. pos_slice: @@ -891,6 +916,8 @@ def add_attn_slot(l: int) -> None: if self._has_attn(l): components_list.append(self[("attn_out", l)]) labels.append(f"{l}_attn_out") + # Parallel hybrids add a mixer output beside attention. + add_mixer(l) elif self._mixer_role(l) == "attn": add_mixer(l) @@ -1230,7 +1257,11 @@ def _post_norm_bias_correction( mapped = self._apply_post_norm_map(raw, scale, weight, centers) if norm_bias is not None: mapped = mapped + norm_bias.to(device=scale.device) - delta = (mapped - raw) * getattr(block, "residual_contribution_scale", 1.0) + delta = (mapped - raw) * ( + getattr(block, "residual_contribution_scale", 1.0) + if kind == "attn" + else self._neuron_contribution_scale(block) + ) correction = delta if correction is None else correction + delta return correction @@ -1408,7 +1439,7 @@ def get_neuron_results( neuron_acts = self[("post", layer, "mlp")] block = self.model.blocks[layer] W_out = block.mlp.W_out - contribution_scale = getattr(block, "residual_contribution_scale", 1.0) + contribution_scale = self._neuron_contribution_scale(block) if pos_slice is not None: # Note - order is important, as Slice.apply *may* collapse a dimension, so this ensures # that position dimension is -2 when we apply position slice @@ -1521,7 +1552,7 @@ def _stack_neuron_results_apply_ln_projected( continue block = self.model.blocks[l] W_out_l = block.mlp.W_out # [d_mlp, d_model] - contribution_scale = getattr(block, "residual_contribution_scale", 1.0) + contribution_scale = self._neuron_contribution_scale(block) post = self._post_norm_map(l, "ln2_post", pos_slice) if post is not None: post_scale, post_weight, post_centers, _ = post @@ -1851,7 +1882,8 @@ def get_full_resid_decomposition( layers only on SSM / hybrid models) 2. Sequence-mixer results of blocks without attention heads (SSM / hybrid models only; one row per mixer, e.g. ``0_linear_attn_out``, output bias - excluded since it is in the ``bias`` row). MLP-type passthrough mixers + excluded since it is in the ``bias`` row), and of mixers that run beside + attention (Falcon-H1's ``0_mamba_out``, bias included). MLP-type passthrough mixers (NemotronH) appear here only when ``expand_neurons=True``; otherwise they are part of the MLP rows below. 3. Neuron / MLP results (only if ``cfg.attn_only=False`` and diff --git a/transformer_lens/model_bridge/generalized_components/block.py b/transformer_lens/model_bridge/generalized_components/block.py index a383bd661..9fa8a5d36 100644 --- a/transformer_lens/model_bridge/generalized_components/block.py +++ b/transformer_lens/model_bridge/generalized_components/block.py @@ -520,10 +520,17 @@ def __init__( residual_contribution_scale: float = 1.0, scaled_attn_submodule: str = "attn", scaled_mlp_submodule: Optional[str] = "mlp", + mlp_contribution_scale: Optional[float] = None, + mixer_contribution_scales: Optional[Dict[str, float]] = None, ): """scaled_mlp_submodule=None when no single submodule feeds the MLP-side add (GraniteMoeHybrid sums moe + shared_mlp inline) — hook_mlp_out then - stays absent rather than firing with a partial tensor.""" + stays absent rather than firing with a partial tensor. + + ``residual_contribution_scale`` scales the attention side, and the MLP side + unless ``mlp_contribution_scale`` is given. ``mixer_contribution_scales`` maps + a sequence mixer that runs beside attention (Falcon-H1's ``mamba``) to the + multiplier HF applies after its ``hook_out``.""" super().__init__( name, config=config, @@ -538,6 +545,16 @@ def __init__( f"must be nonzero (the write path divides by it)." ) self.residual_contribution_scale = scale + mlp_scale = scale if mlp_contribution_scale is None else float(mlp_contribution_scale) + if mlp_scale == 0.0: + raise ValueError( + f"ScaledResidualBlockBridge at '{name}': mlp_contribution_scale " + f"must be nonzero (the write path divides by it)." + ) + self.mlp_contribution_scale = mlp_scale + self.mixer_contribution_scales = { + k: float(v) for k, v in (mixer_contribution_scales or {}).items() + } self.scaled_attn_submodule = scaled_attn_submodule self.scaled_mlp_submodule = scaled_mlp_submodule attn = self.submodules.get(scaled_attn_submodule) if self.submodules else None @@ -585,10 +602,10 @@ def _maybe_wire_capture_hooks(self) -> None: if self._capture_hooks_wired: return super()._maybe_wire_capture_hooks() - targets = [(self.scaled_attn_submodule, "hook_attn_out")] + targets = [(self.scaled_attn_submodule, "hook_attn_out", self.residual_contribution_scale)] if self.scaled_mlp_submodule is not None: - targets.append((self.scaled_mlp_submodule, "hook_mlp_out")) - for sub_name, hook_name in targets: + targets.append((self.scaled_mlp_submodule, "hook_mlp_out", self.mlp_contribution_scale)) + for sub_name, hook_name, scale in targets: hook_point = getattr(self, hook_name, None) sub = self.submodules.get(sub_name) if self.submodules else None if ( @@ -597,14 +614,15 @@ def _maybe_wire_capture_hooks(self) -> None: or getattr(sub, "original_component", None) is None ): continue - handle = sub.register_forward_hook(self._make_scaled_contribution_hook(hook_point)) + handle = sub.register_forward_hook( + self._make_scaled_contribution_hook(hook_point, scale) + ) self._capture_hook_handles.append(handle) def _make_scaled_contribution_hook( - self, hook_point: HookPoint + self, hook_point: HookPoint, scale: float ) -> 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"): diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 1cf0b3ef8..713562deb 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -73,6 +73,9 @@ "mamba_d_conv", "mamba_chunk_size", "lm_head_multiplier", + "attention_out_multiplier", + "ssm_out_multiplier", + "mlp_multipliers", # Multimodal "vision_config", # Cohere diff --git a/transformer_lens/model_bridge/supported_architectures/falcon_h1.py b/transformer_lens/model_bridge/supported_architectures/falcon_h1.py index b633cd42f..f985adb79 100644 --- a/transformer_lens/model_bridge/supported_architectures/falcon_h1.py +++ b/transformer_lens/model_bridge/supported_architectures/falcon_h1.py @@ -19,12 +19,14 @@ needed for forward parity. Folding would only matter for compatibility mode. Key adapter decisions: -- ``BlockBridge`` is the block container: the block has two norms +- ``ScaledResidualBlockBridge`` is the block container: the block has two norms (``input_layernorm`` + ``pre_ff_layernorm``) and a real post-attention residual, so transformer-shaped hooks (``hook_resid_mid`` via ``ln2``) are meaningful — unlike single-norm SSM blocks that need ``SSMBlockBridge``. - ``BlockBridge.forward`` delegates the whole block to HF, so the parallel - attn+mamba combine happens natively. + The forward delegates the whole block to HF, so the parallel attn+mamba + combine happens natively. ``hook_attn_out`` exposes the attention output times + ``attention_out_multiplier``, and the residual decompositions add the Mamba + output times ``ssm_out_multiplier``. - ``SSM2MixerBridge`` wraps the Mamba-2 mixer (passthrough forward). Its inner ``in_proj`` / ``conv1d`` / ``out_proj`` are mapped for hookability; the gated ``inner_norm`` is optional (absent when ``mamba_rms_norm=false``). @@ -45,7 +47,6 @@ from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( - BlockBridge, DepthwiseConv1DBridge, EmbeddingBridge, GatedRMSNormBridge, @@ -53,6 +54,7 @@ PositionEmbeddingsAttentionBridge, RMSNormalizationBridge, RotaryEmbeddingBridge, + ScaledResidualBlockBridge, SSM2MixerBridge, UnembeddingBridge, ) @@ -119,10 +121,24 @@ def __init__(self, cfg: Any) -> None: # forward parity. self.weight_processing_conversions = {} + mlp = self._gated_mlp(name="feed_forward") + # feed_forward multiplies down_proj's output by mlp_multipliers[1]; hook_mlp_out + # includes it, the neuron expansion (post @ W_out) does not. + mlp_multipliers = getattr(self.cfg, "mlp_multipliers", None) or (1.0, 1.0) + mlp._neuron_contribution_scale = float(mlp_multipliers[1]) + self.component_mapping = { - "embed": EmbeddingBridge(name="model.embed_tokens"), + # HF multiplies the embeddings by embedding_multiplier before block 0. + "embed": EmbeddingBridge( + name="model.embed_tokens", + output_contribution_scale=getattr(self.cfg, "embedding_multiplier", 1.0), + ), "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), - "blocks": BlockBridge( + # HF scales the attention and Mamba outputs inline before their shared residual + # add, so hook_attn_out exposes the scaled attention contribution and the + # decompositions scale the Mamba output by ssm_out_multiplier. The MLP's + # multipliers are inside feed_forward. + "blocks": ScaledResidualBlockBridge( name="model.layers", submodules={ "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), @@ -152,7 +168,12 @@ def __init__(self, cfg: Any) -> None: "out_proj": LinearBridge(name="out_proj"), }, ), - "mlp": self._gated_mlp(name="feed_forward"), + "mlp": mlp, + }, + residual_contribution_scale=getattr(self.cfg, "attention_out_multiplier", 1.0), + mlp_contribution_scale=1.0, + mixer_contribution_scales={ + "mamba": getattr(self.cfg, "ssm_out_multiplier", 1.0), }, ), "ln_final": RMSNormalizationBridge(name="model.final_layernorm", config=self.cfg), diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index c13997eec..3befbb191 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -1799,7 +1799,10 @@ def accumulated_bias( if include_mlp_biases and "mlp" in block._modules: b_out = getattr(block.mlp, "b_out", None) if b_out is not None: - accumulated = accumulated + b_out.to(accumulated.device) * scale + mlp_scale = getattr(block, "mlp_contribution_scale", scale) * getattr( + block.mlp, "_neuron_contribution_scale", 1.0 + ) + accumulated = accumulated + b_out.to(accumulated.device) * mlp_scale if mlp_input: assert layer < self.cfg.n_layers, "Cannot include attn_bias from beyond the final layer" block = self.blocks[layer]