fix: apply dtype and device changes to gated MLPs in compatibility mode

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())

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

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)

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])