From d5152daa179e9a5a249fc05cb8defc1fa4e7cc90 Mon Sep 17 00:00:00 2001 From: tivris Date: Wed, 21 Oct 2026 17:46:15 +0300 Subject: [PATCH] fix: apply dtype and device changes to gated MLPs in compatibility mode --- .../generalized_components/test_mlp_bridge.py | 94 +++++++++++++------ tests/unit/test_weight_processing.py | 33 +++++++ .../generalized_components/gated_mlp.py | 33 +++---- .../joint_gate_up_mlp.py | 9 +- 4 files changed, 113 insertions(+), 56 deletions(-) diff --git a/tests/unit/model_bridge/generalized_components/test_mlp_bridge.py b/tests/unit/model_bridge/generalized_components/test_mlp_bridge.py index 85169ab56b..920773c679 100644 --- a/tests/unit/model_bridge/generalized_components/test_mlp_bridge.py +++ b/tests/unit/model_bridge/generalized_components/test_mlp_bridge.py @@ -12,6 +12,7 @@ import copy from types import SimpleNamespace +import pytest import torch import torch.nn as nn @@ -144,32 +145,43 @@ def test_residual_inside_mlp_out_hook_keeps_projection_output(self) -> None: torch.testing.assert_close(seen[0], out - residual) +def _processed_gated_mlp() -> tuple[GatedMLPBridge, torch.Tensor]: + """A GatedMLPBridge on its processed path over biased nn.Linear projections, and an input.""" + d_model, d_mlp = 8, 16 + bridge = GatedMLPBridge( + name="mlp", + config=SimpleNamespace(hidden_act="silu"), + submodules={ + "gate": LinearBridge(name="gate_proj"), + "in": LinearBridge(name="up_proj"), + "out": LinearBridge(name="down_proj"), + }, + ) + shapes = {"gate": (d_mlp, d_model), "in": (d_mlp, d_model), "out": (d_model, d_mlp)} + weights = {} + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + for key, (d_out, d_in) in shapes.items(): + bridge.submodules[key].set_original_component(nn.Linear(d_in, d_out)) + bridge.add_module(key, bridge.submodules[key]) + weights[f"{key}.weight"] = torch.randn(d_out, d_in) + weights[f"{key}.bias"] = torch.randn(d_out) + x = torch.randn(2, 3, d_model) + bridge.set_processed_weights(weights) + return bridge, x + + +def _gated_reference(bridge: GatedMLPBridge, x: torch.Tensor) -> torch.Tensor: + """down(silu(gate(x)) * up(x)) through the wrapped nn.Linear projections.""" + gate, up, down = (bridge.submodules[key].original_component for key in ("gate", "in", "out")) + return down(torch.nn.functional.silu(gate(x)) * up(x)) + + class TestGatedMLPProcessedPathHooks: def test_processed_path_fires_out_hook_out(self) -> None: """The functional processed-weights path bypasses the wrapped projections; out.hook_out must still fire with the down-projection output.""" - d_model, d_mlp = 8, 16 - bridge = GatedMLPBridge( - name="mlp", - config=SimpleNamespace(hidden_act="silu"), - submodules={ - "gate": LinearBridge(name="gate_proj"), - "in": LinearBridge(name="up_proj"), - "out": LinearBridge(name="down_proj"), - }, - ) - for key, sub in bridge.submodules.items(): - bridge.add_module(key, sub) - with torch.random.fork_rng(devices=[]): - torch.manual_seed(0) - bridge._use_processed_weights = True - bridge._processed_W_gate = torch.randn(d_mlp, d_model) - bridge._processed_b_gate = None - bridge._processed_W_in = torch.randn(d_mlp, d_model) - bridge._processed_b_in = None - bridge._processed_W_out = torch.randn(d_model, d_mlp) - bridge._processed_b_out = None - x = torch.randn(2, 3, d_model) + bridge, x = _processed_gated_mlp() out_bridge = bridge.submodules["out"] seen: list[torch.Tensor] = [] @@ -181,10 +193,38 @@ def test_processed_path_fires_out_hook_out(self) -> None: out_bridge.hook_out.remove_hooks() assert len(seen) == 1 - expected = torch.nn.functional.linear( - torch.nn.functional.silu(torch.nn.functional.linear(x, bridge._processed_W_gate)) - * torch.nn.functional.linear(x, bridge._processed_W_in), - bridge._processed_W_out, - ) + with torch.no_grad(): + expected = _gated_reference(bridge, x) torch.testing.assert_close(seen[0], expected) torch.testing.assert_close(out, expected) + + +class TestGatedMLPProcessedPathWeights: + @pytest.mark.parametrize("dtype", [torch.float16, torch.float64]) + def test_follows_dtype_change(self, dtype: torch.dtype) -> None: + """After .to(dtype), the processed path computes with the converted projections.""" + bridge, x = _processed_gated_mlp() + bridge.to(dtype) + x = x.to(dtype) + with torch.no_grad(): + out = bridge(x) + expected = _gated_reference(bridge, x) + assert out.dtype == dtype + torch.testing.assert_close(out, expected) + + def test_deepcopy_reads_its_own_projections(self) -> None: + """In-place edits to a deep copy's projections reach the copy's processed path.""" + bridge, x = _processed_gated_mlp() + clone = copy.deepcopy(bridge) + with torch.no_grad(): + clone.out.original_component.weight.zero_() + clone.out.original_component.bias.zero_() + assert torch.count_nonzero(clone(x)) == 0 + + def test_projections_get_no_gradient(self) -> None: + """The processed path reads the projections detached from autograd.""" + bridge, x = _processed_gated_mlp() + x.requires_grad_() + bridge(x).sum().backward() + assert x.grad is not None + assert all(param.grad is None for param in bridge.parameters()) diff --git a/tests/unit/test_weight_processing.py b/tests/unit/test_weight_processing.py index 47b2ed7aea..901115da06 100644 --- a/tests/unit/test_weight_processing.py +++ b/tests/unit/test_weight_processing.py @@ -1146,3 +1146,36 @@ def test_process_weights_created_biases_keep_model_dtype(architecture, dtype): assert not [name for name, param in hf_model.named_parameters() if param.dtype != dtype] bridge(torch.tensor([[1, 2, 3, 4]])) + + +GATED_MLP_CONFIGS = { + "Phi3ForCausalLM": dict( + model_type="phi3", + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + pad_token_id=0, + ), + "Qwen2ForCausalLM": BIAS_CREATING_CONFIGS["Qwen2ForCausalLM"], + "StableLmForCausalLM": BIAS_CREATING_CONFIGS["StableLmForCausalLM"], +} + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("architecture", sorted(GATED_MLP_CONFIGS)) +def test_compat_gated_mlp_follows_dtype_change(architecture, dtype): + """Gated MLPs keep their processed forward when a compat bridge changes dtype.""" + hf_config = AutoConfig.for_model(vocab_size=64, **GATED_MLP_CONFIGS[architecture]) + hf_config.architectures = [architecture] + hf_model = AutoModelForCausalLM.from_config(hf_config) + bridge = build_bridge_from_module(hf_model, architecture, hf_config=hf_config) + bridge.enable_compatibility_mode() + + bridge.to(dtype) + logits, cache = bridge.run_with_cache(torch.tensor([[1, 2, 3, 4]])) + + assert logits.dtype == dtype + # gate.hook_in fires only off the processed path. + assert "blocks.0.mlp.gate.hook_in" not in cache diff --git a/transformer_lens/model_bridge/generalized_components/gated_mlp.py b/transformer_lens/model_bridge/generalized_components/gated_mlp.py index 59dbaed4fe..fc3a9c391a 100644 --- a/transformer_lens/model_bridge/generalized_components/gated_mlp.py +++ b/transformer_lens/model_bridge/generalized_components/gated_mlp.py @@ -106,22 +106,15 @@ def forward(self, *args, **kwargs) -> torch.Tensor: Output hidden states """ if hasattr(self, "_use_processed_weights") and self._use_processed_weights: - assert hasattr(self, "_processed_W_gate") and hasattr(self, "_processed_W_in"), ( - "Processed weights flag is set but weights are missing. " - "This indicates a bug in set_processed_weights()." - ) - assert self._processed_W_in is not None - assert self._processed_W_out is not None + W_gate, b_gate = self._processed_projection("gate") + W_in, b_in = self._processed_projection("in") + W_out, b_out = self._processed_projection("out") hidden_states = args[0] hidden_states = self.hook_in(hidden_states) - gate_output = torch.nn.functional.linear( - hidden_states, self._processed_W_gate, self._processed_b_gate - ) + gate_output = torch.nn.functional.linear(hidden_states, W_gate, b_gate) if hasattr(self, "gate") and hasattr(self.gate, "hook_out"): gate_output = self.gate.hook_out(gate_output) - linear_output = torch.nn.functional.linear( - hidden_states, self._processed_W_in, self._processed_b_in - ) + linear_output = torch.nn.functional.linear(hidden_states, W_in, b_in) in_module = getattr(self, "in", None) if in_module is not None and hasattr(in_module, "hook_out"): linear_output = in_module.hook_out(linear_output) # type: ignore[misc] @@ -130,9 +123,7 @@ def forward(self, *args, **kwargs) -> torch.Tensor: hidden = activated * linear_output if hasattr(self, "out") and hasattr(self.out, "hook_in"): hidden = self.out.hook_in(hidden) - output = torch.nn.functional.linear( - hidden, self._processed_W_out, self._processed_b_out - ) + output = torch.nn.functional.linear(hidden, W_out, b_out) # The functional path bypasses the wrapped `out` projection, so fire # its hook_out here — it is the down-projection output. if hasattr(self, "out") and hasattr(self.out, "hook_out"): @@ -150,6 +141,12 @@ def forward(self, *args, **kwargs) -> torch.Tensor: output = self.hook_out(output) return output + def _processed_projection(self, name: str) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Current weight and bias of a projection, detached from autograd.""" + projection = getattr(self, name) + bias = projection.bias + return projection.weight.detach(), None if bias is None else bias.detach() + def set_processed_weights( self, weights: Mapping[str, torch.Tensor | None], verbose: bool = False ) -> None: @@ -191,12 +188,6 @@ def set_processed_weights( print(f" Setting W_out with shape: {W_out.shape}") self._use_processed_weights = True - self._processed_W_gate = W_gate - self._processed_b_gate = b_gate - self._processed_W_in = W_in - self._processed_b_in = b_in - self._processed_W_out = W_out - self._processed_b_out = b_out # Distribute to submodules if they support it gate_module = getattr(self, "gate", None) diff --git a/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py b/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py index e0e998a0b2..9f0b28b4d7 100644 --- a/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py +++ b/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py @@ -175,14 +175,7 @@ def _resolve_activation_fn(self) -> Callable: def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: """Reconstructed gated MLP forward with individual hook access.""" - # Delegate to GatedMLPBridge's processed-weights path only when ALL - # processed weights exist; its fallback bypasses intermediate hooks. - if ( - hasattr(self, "_use_processed_weights") - and self._use_processed_weights - and hasattr(self, "_processed_W_gate") - and hasattr(self, "_processed_W_in") - ): + if hasattr(self, "_use_processed_weights") and self._use_processed_weights: return super().forward(*args, **kwargs) hidden_states = self.hook_in(args[0])