diff --git a/tests/integration/model_bridge/test_granite_embedding_decomposition.py b/tests/integration/model_bridge/test_granite_embedding_decomposition.py new file mode 100644 index 000000000..ee3c64039 --- /dev/null +++ b/tests/integration/model_bridge/test_granite_embedding_decomposition.py @@ -0,0 +1,229 @@ +"""Offline regressions for Granite's post-embedding runtime multiplier.""" + +import copy + +import pytest +import torch +from transformers import GraniteConfig, GraniteForCausalLM + +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_from_module, +) +from transformer_lens.tools.analysis import direct_logit_attribution + + +@pytest.fixture(params=[1.0, 0.5, 2.0, 12.0]) +def granite(request): + torch.manual_seed(0) + config = GraniteConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=32, + embedding_multiplier=request.param, + residual_multiplier=0.22, + logits_scaling=1.0, + attention_multiplier=0.5, + attn_implementation="eager", + ) + model = GraniteForCausalLM(config).eval() + reference = copy.deepcopy(model) + bridge = build_bridge_from_module(model, "GraniteForCausalLM", hf_config=config, device="cpu") + yield bridge, reference, request.param + bridge.close() + + +@pytest.mark.parametrize("compatibility", [False, True]) +@pytest.mark.parametrize("remove_batch", [False, True]) +@pytest.mark.parametrize("expand_neurons", [False, True]) +def test_embedding_contribution_reconstructs_residual( + granite, compatibility, remove_batch, expand_neurons +): + bridge, _, multiplier = granite + if compatibility: + bridge.enable_compatibility_mode() + assert bridge.cfg.embedding_multiplier == multiplier + tokens = torch.tensor([[3, 4, 5]]) + with torch.no_grad(): + _, cache = bridge.run_with_cache(tokens, remove_batch_dim=remove_batch) + stack, labels = cache.decompose_resid(return_labels=True) + torch.testing.assert_close(stack[labels.index("embed")], cache["hook_embed"]) + torch.testing.assert_close(stack.sum(0), cache["blocks.1.hook_out"], rtol=1e-5, atol=1e-7) + torch.testing.assert_close( + cache.decompose_resid(layer=0).sum(0), + cache["blocks.0.hook_in"], + rtol=1e-5, + atol=1e-7, + ) + torch.testing.assert_close( + cache.decompose_resid(pos_slice=-1).sum(0), + cache["blocks.1.hook_out"][..., -1, :], + rtol=1e-5, + atol=1e-7, + ) + full = cache.get_full_resid_decomposition(expand_neurons=expand_neurons) + torch.testing.assert_close(full.sum(0), cache["blocks.1.hook_out"], rtol=1e-5, atol=1e-7) + + +def test_component_dla_reconstructs_logit_difference(granite): + bridge, _, _ = granite + bridge.enable_compatibility_mode() + tokens = torch.tensor([[3, 4, 5]]) + with torch.no_grad(): + logits, cache = bridge.run_with_cache(tokens) + result = direct_logit_attribution( + bridge, answer_tokens=7, incorrect_tokens=9, unit="component", cache=cache + ) + torch.testing.assert_close( + result.attribution.sum(0), + logits[:, -1, 7] - logits[:, -1, 9], + rtol=1e-5, + atol=1e-7, + ) + + +def test_embedding_hook_exposes_residual_contribution_without_changing_forward(granite): + bridge, reference, multiplier = granite + tokens = torch.tensor([[3, 4, 5]]) + with torch.no_grad(): + logits, cache = bridge.run_with_cache(tokens) + torch.testing.assert_close(logits, reference(tokens).logits, rtol=0, atol=0) + torch.testing.assert_close( + cache["hook_embed"], reference.model.embed_tokens(tokens) * multiplier + ) + torch.testing.assert_close(cache["hook_embed"], cache["blocks.0.hook_in"]) + bridge.add_hook("embed.hook_out", lambda value, hook: value * 3) + try: + _, edited = bridge.run_with_cache(tokens) + finally: + bridge.reset_hooks() + torch.testing.assert_close(edited["hook_embed"], cache["hook_embed"] * 3) + torch.testing.assert_close(edited["blocks.0.hook_in"], edited["hook_embed"]) + + +@pytest.mark.parametrize("in_place", [False, True]) +def test_embedding_ablation_lands_in_residual_stream(granite, in_place): + bridge, _, _ = granite + + def ablate(value, hook): + return value.zero_() if in_place else torch.zeros_like(value) + + bridge.add_hook("embed.hook_out", ablate) + try: + with torch.no_grad(): + _, cache = bridge.run_with_cache(torch.tensor([[3, 4, 5]])) + torch.testing.assert_close(cache["blocks.0.hook_in"], torch.zeros_like(cache["hook_embed"])) + finally: + bridge.reset_hooks() + + +def test_embedding_backward_hook_receives_residual_gradient(granite): + bridge, _, _ = granite + gradients = {} + + def record(gradient, hook): + gradients[hook.name] = gradient.detach().clone() + + bridge.add_hook("embed.hook_out", record, dir="bwd") + bridge.add_hook("blocks.0.hook_in", record, dir="bwd") + try: + bridge(torch.tensor([[3, 4, 5]])).square().sum().backward() + assert torch.isfinite(gradients["embed.hook_out"]).all() + assert gradients["embed.hook_out"].abs().sum() > 0 + torch.testing.assert_close(gradients["embed.hook_out"], gradients["blocks.0.hook_in"]) + finally: + bridge.reset_hooks() + + +def test_live_and_reconstructed_head_results_have_same_scale(granite): + bridge, _, _ = granite + tokens = torch.tensor([[3, 4, 5]]) + with torch.no_grad(): + _, reconstructed = bridge.run_with_cache(tokens) + reconstructed.compute_head_results() + bridge.set_use_attn_result(True) + _, live = bridge.run_with_cache(tokens) + for layer in range(bridge.cfg.n_layers): + key = f"blocks.{layer}.attn.hook_result" + torch.testing.assert_close(live[key], reconstructed[key], rtol=1e-5, atol=1e-7) + torch.testing.assert_close( + live[key].sum(-2), live[f"blocks.{layer}.hook_attn_out"], rtol=1e-5, atol=1e-7 + ) + torch.testing.assert_close( + live.get_full_resid_decomposition(expand_neurons=True).sum(0), + live["blocks.1.hook_out"], + rtol=1e-5, + atol=1e-7, + ) + bridge.add_hook("blocks.0.attn.hook_result", lambda value, hook: torch.zeros_like(value)) + try: + with torch.no_grad(): + _, edited = bridge.run_with_cache(tokens) + torch.testing.assert_close( + edited["blocks.0.hook_attn_out"], torch.zeros_like(edited["blocks.0.hook_attn_out"]) + ) + finally: + bridge.reset_hooks() + + +def test_neuron_projection_and_scaled_bias_reconstruct_residual(granite): + _, reference, _ = granite + model = copy.deepcopy(reference) + # Bind the bias-bearing model before Bridge snapshots its parameter aliases. + for layer in model.model.layers: + layer.self_attn.o_proj.bias = torch.nn.Parameter(torch.randn(16) * 0.01) + layer.mlp.down_proj.bias = torch.nn.Parameter(torch.randn(16) * 0.01) + bridge = build_bridge_from_module( + model, "GraniteForCausalLM", hf_config=model.config, device="cpu" + ) + try: + with torch.no_grad(): + _, cache = bridge.run_with_cache(torch.tensor([[3, 4, 5]])) + full = cache.get_full_resid_decomposition(expand_neurons=True) + torch.testing.assert_close(full.sum(0), cache["blocks.1.hook_out"], rtol=1e-5, atol=1e-7) + directions = torch.randn(16, 2) + dense = cache.stack_neuron_results(2, apply_ln=True) @ directions + projected = cache.stack_neuron_results(2, apply_ln=True, project_output_onto=directions) + torch.testing.assert_close(projected, dense, rtol=1e-5, atol=1e-7) + finally: + bridge.close() + + +@pytest.mark.parametrize("multiplier", [0.5, 2.0]) +def test_moe_embedding_contribution(multiplier): + from transformers import GraniteMoeConfig, GraniteMoeForCausalLM + + torch.manual_seed(0) + config = GraniteMoeConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=32, + num_local_experts=2, + num_experts_per_tok=1, + embedding_multiplier=multiplier, + residual_multiplier=0.22, + logits_scaling=1.0, + attn_implementation="eager", + ) + model = GraniteMoeForCausalLM(config).eval() + bridge = build_bridge_from_module( + model, "GraniteMoeForCausalLM", hf_config=config, device="cpu" + ) + try: + with torch.no_grad(): + _, cache = bridge.run_with_cache(torch.tensor([[3, 4, 5]])) + torch.testing.assert_close( + cache.decompose_resid().sum(0), + cache["blocks.0.hook_out"], + rtol=1e-5, + atol=1e-7, + ) + finally: + bridge.close() 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 a80ef2e01..faa2371a2 100644 --- a/tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py +++ b/tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py @@ -20,6 +20,7 @@ """ import contextlib +import copy import pytest import torch @@ -44,6 +45,34 @@ ATTN_LAYER = 1 +def test_hybrid_embedding_hook_exposes_scaled_contribution(hf_model, tokens): + cfg = GraniteMoeHybridConfig.from_dict(hf_model.config.to_dict()) + cfg.embedding_multiplier = 12.0 + cfg.residual_multiplier = 0.22 + cfg._attn_implementation = "eager" + model = AutoModelForCausalLM.from_config(cfg).to(torch.float32).eval() + reference = copy.deepcopy(model) + bridge_config = build_bridge_config_from_hf( + cfg, "GraniteMoeHybridForCausalLM", "granitemoehybrid-tiny", torch.float32 + ) + adapter = GraniteMoeHybridArchitectureAdapter(bridge_config) + bridge = TransformerBridge(model=model, adapter=adapter, tokenizer=_MockTokenizer()) + try: + with torch.no_grad(): + baseline = bridge(tokens, use_cache=False) + logits, cache = bridge.run_with_cache( + tokens, use_cache=False, names_filter=["hook_embed", "blocks.0.hook_in"] + ) + expected = reference(tokens, use_cache=False).logits + embed = reference.model.embed_tokens(tokens) * cfg.embedding_multiplier + torch.testing.assert_close(logits, baseline, rtol=0, atol=0) + torch.testing.assert_close(logits, expected, rtol=1e-5, atol=1e-7) + torch.testing.assert_close(cache["hook_embed"], embed) + torch.testing.assert_close(cache["hook_embed"], cache["blocks.0.hook_in"]) + finally: + bridge.close() + + class _MockTokenizer: """Stand-in to satisfy TransformerBridge(tokenizer=...) without a Hub call.""" diff --git a/tests/unit/model_bridge/generalized_components/test_scaled_embedding.py b/tests/unit/model_bridge/generalized_components/test_scaled_embedding.py new file mode 100644 index 000000000..6d6ccc29b --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_scaled_embedding.py @@ -0,0 +1,45 @@ +"""Embedding contribution hooks preserve the outer model's runtime scale.""" + +import pytest +import torch + +from transformer_lens.model_bridge.generalized_components.embedding import ( + EmbeddingBridge, +) + + +@pytest.mark.parametrize("scale", [0.0, 0.22, 1.0, 12.0]) +def test_scaled_embedding_observation_preserves_raw_output(scale): + embedding = torch.nn.Embedding(8, 4) + bridge = EmbeddingBridge(name="embed", output_contribution_scale=scale) + bridge.set_original_component(embedding) + tokens = torch.tensor([[1, 2]]) + expected = embedding(tokens) + torch.testing.assert_close(bridge(tokens), expected, rtol=0, atol=0) + observed = [] + bridge.hook_out.add_hook(lambda value, hook: observed.append(value.detach().clone())) + torch.testing.assert_close(bridge(tokens), expected, rtol=0, atol=0) + torch.testing.assert_close(observed[0], expected * scale, rtol=0, atol=0) + + +@pytest.mark.parametrize("scale", [0.22, 1.0, 12.0]) +@pytest.mark.parametrize("in_place", [False, True]) +def test_scaled_embedding_write_is_not_scaled_twice(scale, in_place): + bridge = EmbeddingBridge(name="embed", output_contribution_scale=scale) + bridge.set_original_component(torch.nn.Embedding(8, 4)) + replacement = torch.randn(1, 2, 4) + + def write(value, hook): + return value.copy_(replacement) if in_place else replacement + + bridge.hook_out.add_hook(write) + actual = bridge(torch.tensor([[1, 2]])) * scale + torch.testing.assert_close(actual, replacement) + + +def test_zero_scale_rejects_unrealizable_edit(): + bridge = EmbeddingBridge(name="embed", output_contribution_scale=0.0) + bridge.set_original_component(torch.nn.Embedding(8, 4)) + bridge.hook_out.add_hook(lambda value, hook: torch.ones_like(value)) + with pytest.raises(ValueError, match="zero runtime scale"): + bridge(torch.tensor([[1, 2]])) diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index ad097942a..15529ffd4 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -726,7 +726,8 @@ def decompose_resid( A slice object to apply to the pos dimension. Defaults to None, do nothing. incl_embeds: - Whether to include embed & pos_embed + Whether to include embed & pos_embed, whose hooks expose their + contributions to the residual stream. return_labels: Whether to return a list of labels for the residual stream components. Useful for labelling graphs. @@ -820,7 +821,9 @@ def _ensure_head_results(self) -> bool: # Preserve broadcast promotion for caches collected under autocast. dtype = torch.promote_types(z.dtype, weights.dtype) z, weights = z.to(dtype), weights.to(dtype) - self.cache_dict[result_key] = torch.einsum("...hd,hdm->...hm", z, weights) + result = torch.einsum("...hd,hdm->...hm", z, weights) + scale = getattr(self.model.blocks[layer].attn, "_result_contribution_scale", 1.0) + self.cache_dict[result_key] = result if scale == 1.0 else result * scale return all_cached def ssm_layers(self, mixer_type: Optional[Union[type, Tuple[type, ...]]] = None) -> List[int]: @@ -1158,6 +1161,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) 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 @@ -1165,6 +1169,8 @@ def get_neuron_results( if neuron_slice is not None: neuron_acts = neuron_slice.apply(neuron_acts, dim=-1) W_out = neuron_slice.apply(W_out, dim=0) + if contribution_scale != 1.0: + neuron_acts = neuron_acts * contribution_scale if project_output_onto is None: return neuron_acts[..., None] * W_out # W_out: [d_mlp, d_model]; project: [d_model] or [d_model, n_outs] @@ -1239,6 +1245,7 @@ def _stack_neuron_results_apply_ln_projected( for l in range(layer): 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) W_out_l_sliced = neuron_slice.apply(W_out_l, dim=0) W_proj_l = W_out_l_sliced @ project_2d # [d_mlp, n_outs] if apply_centering: @@ -1247,6 +1254,8 @@ def _stack_neuron_results_apply_ln_projected( lin_form_l = W_proj_l - W_means_l[:, None] * sum_p[None, :] else: lin_form_l = W_proj_l + if contribution_scale != 1.0: + lin_form_l = lin_form_l * contribution_scale a_l = self[("post", l, "mlp")] a_l = pos_slice.apply(a_l, dim=-2) a_l = neuron_slice.apply(a_l, dim=-1) @@ -1530,7 +1539,7 @@ def get_full_resid_decomposition( 2. Neuron / MLP results (only if ``cfg.attn_only=False`` and ``layer > 0``; ``L * d_mlp`` rows when ``expand_neurons=True``, else ``L`` rows) - 3. ``embed`` (1 row, if the model has token embeddings) + 3. ``embed`` (1 row, including any post-embedding runtime multiplier) 4. ``pos_embed`` (1 row, if the model has positional embeddings) 5. ``bias`` (1 row, the accumulated layer biases) @@ -1626,7 +1635,7 @@ def _ln_then_project(stack: torch.Tensor) -> torch.Tensor: components.append(mlp_stack) if self.has_embed: - embed = pos_slice.apply(self["embed"], -2)[None] + embed = pos_slice.apply(self["hook_embed"], -2)[None] if ln_folded: embed = _ln_then_project(embed) elif project_2d is not None: diff --git a/transformer_lens/model_bridge/generalized_components/attention.py b/transformer_lens/model_bridge/generalized_components/attention.py index 40bb4786c..4ef6058b0 100644 --- a/transformer_lens/model_bridge/generalized_components/attention.py +++ b/transformer_lens/model_bridge/generalized_components/attention.py @@ -810,7 +810,12 @@ def _compute_per_head_result( d_head=d_head, ) per_head = torch.einsum("bshd,hdm->bshm", z_4d, w_per_head) - per_head = self.hook_result(per_head) + scale = getattr(self, "_result_contribution_scale", 1.0) + if scale == 1.0: + # Keep callable hook replacements on the original unscaled path. + per_head = self.hook_result(per_head) + else: + per_head = self._hook_scaled_output(per_head, self.hook_result, scale) summed = per_head.sum(dim=-2) if bias is not None: summed = summed + bias diff --git a/transformer_lens/model_bridge/generalized_components/base.py b/transformer_lens/model_bridge/generalized_components/base.py index 80d96b80c..044c69b60 100644 --- a/transformer_lens/model_bridge/generalized_components/base.py +++ b/transformer_lens/model_bridge/generalized_components/base.py @@ -135,6 +135,22 @@ def __init__( self.hook_aliases = self.__class__.hook_aliases.copy() self.hook_aliases.update(hook_alias_overrides) + def _hook_scaled_output( + self, output: torch.Tensor, hook_point: HookPoint, scale: float + ) -> torch.Tensor: + """Expose a residual contribution while leaving an outer HF scale intact.""" + if scale == 1.0: + return hook_point(output) + if not hook_point.has_hooks(dir="both"): + return output + hooked = hook_point(output * scale) + # Recompute the reference because an in-place hook can mutate its input. + if not hook_point.has_hooks(dir="bwd") and torch.equal(hooked, output * scale): + return output + if scale == 0.0: + raise ValueError("Cannot edit or backward-hook a contribution with zero runtime scale") + return hooked / scale + def mark_pos_dim(self, pos_dim: int) -> None: """Declare the position axis of every HookPoint this component owns. diff --git a/transformer_lens/model_bridge/generalized_components/block.py b/transformer_lens/model_bridge/generalized_components/block.py index f3bba6f66..3bb3e98f5 100644 --- a/transformer_lens/model_bridge/generalized_components/block.py +++ b/transformer_lens/model_bridge/generalized_components/block.py @@ -535,6 +535,9 @@ def __init__( self.residual_contribution_scale = scale 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 + if attn is not None: + attn._result_contribution_scale = scale if self.hook_aliases is BlockBridge.hook_aliases: self.hook_aliases = dict(self.hook_aliases) for alias in ("hook_attn_out", "hook_mlp_out"): @@ -605,13 +608,9 @@ def _hook(_module: torch.nn.Module, _args: tuple, output: Any) -> Any: 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): + new = self._hook_scaled_output(out, hook_point, scale) + if new is out: return None - new = hooked / scale return ((new,) + output[1:]) if is_tuple else new return _hook diff --git a/transformer_lens/model_bridge/generalized_components/embedding.py b/transformer_lens/model_bridge/generalized_components/embedding.py index c69d9144e..102c435c7 100644 --- a/transformer_lens/model_bridge/generalized_components/embedding.py +++ b/transformer_lens/model_bridge/generalized_components/embedding.py @@ -25,6 +25,7 @@ def __init__( name: str, config: Optional[Any] = None, submodules: Optional[Dict[str, GeneralizedComponent]] = {}, + output_contribution_scale: float = 1.0, ): """Initialize the embedding bridge. @@ -32,8 +33,11 @@ def __init__( name: The name of this component config: Optional configuration (unused for EmbeddingBridge) submodules: Dictionary of GeneralizedComponent submodules to register + output_contribution_scale: Outer HF runtime scale exposed by hook_out. + Observation preserves the raw forward; edits are unscaled before returning. """ super().__init__(name, config, submodules=submodules) + self._output_contribution_scale = float(output_contribution_scale) @property def W_E(self) -> torch.Tensor: @@ -90,5 +94,5 @@ def forward( output = output[0] if target_dtype is not None and output.dtype != target_dtype: output = output.to(dtype=target_dtype) - output = self.hook_out(output) + output = self._hook_scaled_output(output, self.hook_out, self._output_contribution_scale) return output diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index bebebe656..1ca8cb440 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -46,6 +46,7 @@ "position_embedding_type", "logits_scaling", "residual_multiplier", + "embedding_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 eef2b5710..1aada0b28 100644 --- a/transformer_lens/model_bridge/supported_architectures/granite.py +++ b/transformer_lens/model_bridge/supported_architectures/granite.py @@ -30,6 +30,8 @@ class GraniteArchitectureAdapter(ArchitectureAdapter): 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. + EmbeddingBridge similarly exposes embedding_multiplier at hook_embed, + while per-head hook_result includes the residual multiplier. Optional Parameters (may not exist in state_dict): ------------------------------------------------- @@ -60,6 +62,13 @@ def _build_attention_bridge(self, optional: bool = False) -> AttentionBridge: """Build the standard Granite attention bridge (GraniteMoeHybrid passes optional).""" return self._qkvo_attention_bridge(optional=optional) + def _build_embedding_bridge(self) -> EmbeddingBridge: + """Expose the embedding's scaled residual contribution at hook_out.""" + return EmbeddingBridge( + name="model.embed_tokens", + output_contribution_scale=getattr(self.cfg, "embedding_multiplier", 1.0), + ) + def _build_mlp_bridge(self) -> GatedMLPBridge: """Build the dense gated MLP bridge.""" return self._gated_mlp() @@ -84,7 +93,7 @@ def _build_moe_bridge(self) -> MoEBridge: def _build_component_mapping(self) -> dict: """Build the full component mapping for dense Granite.""" return { - "embed": EmbeddingBridge(name="model.embed_tokens"), + "embed": self._build_embedding_bridge(), "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), # HF multiplies each sublayer output by residual_multiplier before the # residual add, so hook_attn_out / hook_mlp_out must expose the scaled diff --git a/transformer_lens/model_bridge/supported_architectures/granite_moe.py b/transformer_lens/model_bridge/supported_architectures/granite_moe.py index c71a12cbb..194dfdb1c 100644 --- a/transformer_lens/model_bridge/supported_architectures/granite_moe.py +++ b/transformer_lens/model_bridge/supported_architectures/granite_moe.py @@ -1,7 +1,6 @@ """Granite MoE architecture adapter.""" from transformer_lens.model_bridge.generalized_components import ( - EmbeddingBridge, RMSNormalizationBridge, RotaryEmbeddingBridge, ScaledResidualBlockBridge, @@ -23,7 +22,7 @@ class GraniteMoeArchitectureAdapter(GraniteArchitectureAdapter): def _build_component_mapping(self) -> dict: """Build component mapping with MoE instead of dense MLP.""" return { - "embed": EmbeddingBridge(name="model.embed_tokens"), + "embed": self._build_embedding_bridge(), "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), # HF multiplies each sublayer output by residual_multiplier before the # residual add; hook_attn_out / hook_mlp_out expose the scaled contribution. 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 2f94390a5..506557446 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 ( - EmbeddingBridge, GatedRMSNormBridge, JointGateUpMLPBridge, LinearBridge, @@ -110,7 +109,7 @@ def _build_component_mapping(self) -> dict: # 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"), + "embed": self._build_embedding_bridge(), "blocks": ScaledResidualBlockBridge( name="model.layers", submodules=block_submodules, diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index f4c0c43d1..45ec8609c 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -1705,19 +1705,21 @@ def accumulated_bias( accumulated = torch.zeros(self.cfg.d_model, device=self.cfg.device) for i in range(layer): block = self.blocks[i] + scale = getattr(block, "residual_contribution_scale", 1.0) b_O = self._get_block_variant_bias(block) if b_O is not None: - accumulated = accumulated + b_O.to(accumulated.device) + accumulated = accumulated + b_O.to(accumulated.device) * scale 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) + accumulated = accumulated + b_out.to(accumulated.device) * scale if mlp_input: assert layer < self.cfg.n_layers, "Cannot include attn_bias from beyond the final layer" block = self.blocks[layer] + scale = getattr(block, "residual_contribution_scale", 1.0) b_O = self._get_block_variant_bias(block) if b_O is not None: - accumulated = accumulated + b_O.to(accumulated.device) + accumulated = accumulated + b_O.to(accumulated.device) * scale return accumulated def all_composition_scores(self, mode: str) -> CompositionScores: