Preserve input dtype in GeneralizedComponent

tests/unit/model_bridge/generalized_components/test_base.py

@@ -5,9 +5,15 @@
 import torch.nn as nn
 
 from transformer_lens.hook_points import HookPoint
+from transformer_lens.model_bridge.generalized_components.attention import (
+    AttentionBridge,
+)
 from transformer_lens.model_bridge.generalized_components.base import (
     GeneralizedComponent,
 )
+from transformer_lens.model_bridge.generalized_components.unembedding import (
+    UnembeddingBridge,
+)
 
 
 class MockOriginalComponent(nn.Module):
@@ -377,5 +383,80 @@ def test_complex_object_attributes(self):
         assert component.complex_attr["nested"]["deep"] == [1, 2, 3]
 
 
+@pytest.mark.parametrize("use_keyword", [False, True])
+def test_forward_preserves_input_dtype(use_keyword: bool):
+    """Parameter storage dtype must not determine the component compute dtype."""
+
+    class MixedPrecisionComponent(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.auxiliary = nn.Parameter(torch.ones((), dtype=torch.float32))
+            self.received_dtype = None
+
+        def forward(self, hidden_states):
+            self.received_dtype = hidden_states.dtype
+            return hidden_states.clone()
+
+    original = MixedPrecisionComponent()
+    component = MockGeneralizedComponent("mixed_precision")
+    component.set_original_component(original)
+    inputs = torch.ones(2, 3, dtype=torch.bfloat16)
+
+    output = component(hidden_states=inputs) if use_keyword else component(inputs)
+
+    assert original.received_dtype == torch.bfloat16
+    assert output.dtype == torch.bfloat16
+    torch.testing.assert_close(output, inputs)
+
+
+@pytest.mark.parametrize("input_name", ["positional", "hidden_states", "query_input"])
+def test_attention_forward_preserves_input_dtype(input_name: str):
+    class MixedPrecisionAttention(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.auxiliary = nn.Parameter(torch.ones((), dtype=torch.float32))
+            self.received_dtype = None
+
+        def forward(self, hidden_states=None, query_input=None):
+            value = query_input if query_input is not None else hidden_states
+            self.received_dtype = value.dtype
+            return value
+
+    original = MixedPrecisionAttention()
+    bridge = AttentionBridge(name="attention", config=None)
+    bridge.set_original_component(original)
+    inputs = torch.ones(2, 3, dtype=torch.bfloat16)
+
+    if input_name == "positional":
+        output = bridge(inputs)
+    else:
+        output = bridge(**{input_name: inputs})
+
+    assert original.received_dtype == torch.bfloat16
+    assert output.dtype == torch.bfloat16
+
+
+def test_unembedding_forward_preserves_input_dtype():
+    class MixedPrecisionUnembedding(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.auxiliary = nn.Parameter(torch.ones((), dtype=torch.float32))
+            self.received_dtype = None
+
+        def forward(self, hidden_states):
+            self.received_dtype = hidden_states.dtype
+            return hidden_states
+
+    original = MixedPrecisionUnembedding()
+    bridge = UnembeddingBridge(name="unembed")
+    bridge.set_original_component(original)
+    inputs = torch.ones(2, 3, dtype=torch.bfloat16)
+
+    output = bridge(inputs)
+
+    assert original.received_dtype == torch.bfloat16
+    assert output.dtype == torch.bfloat16
+
+
 if __name__ == "__main__":
     pytest.main([__file__])

tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py

@@ -39,6 +39,23 @@ def _bridge_with_stub(fake_forward) -> MoEBridge:
 
 
 class TestMoEBridgeTupleOutput:
+    @pytest.mark.parametrize("use_keyword", [False, True])
+    def test_preserves_input_dtype(self, use_keyword: bool) -> None:
+        received_dtype = None
+
+        def fake_forward(hidden_states):
+            nonlocal received_dtype
+            received_dtype = hidden_states.dtype
+            return hidden_states
+
+        bridge = _bridge_with_stub(fake_forward)
+        hidden_states = torch.ones(1, 3, 4, dtype=torch.bfloat16)
+
+        output = bridge(hidden_states=hidden_states) if use_keyword else bridge(hidden_states)
+
+        assert received_dtype == torch.bfloat16
+        assert output.dtype == torch.bfloat16
+
     def test_empty_tuple_raises_clear_type_error(self) -> None:
         bridge = _bridge_with_stub(lambda *a, **kw: ())
         with pytest.raises(TypeError, match="torch.Tensor"):

tests/unit/model_bridge/supported_architectures/test_olmo_adapter.py

@@ -100,6 +100,14 @@ def __init__(self, d_model: int, n_heads: int, n_kv_heads: int) -> None:
         self.o_proj = nn.Linear(n_heads * head_dim, d_model, bias=False)
 
 
+class _RecordingLinear(nn.Linear):
+    """Record the bridge input dtype while handling conversion internally."""
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        self.received_dtype = input.dtype
+        return super().forward(input.to(self.weight.dtype))
+
+
 def _wire_attention_bridge(
     adapter: OlmoArchitectureAdapter,
     cfg: TransformerBridgeConfig,
@@ -297,6 +305,32 @@ def _hook(tensor: torch.Tensor, hook: Any) -> None:
         assert seen["v"] == torch.Size([batch, seq_len, cfg.n_key_value_heads, cfg.d_head])
         assert seen["z"] == torch.Size([batch, seq_len, cfg.n_heads, cfg.d_head])
 
+    def test_forward_preserves_input_dtype_at_projection_boundary(
+        self, adapter: OlmoArchitectureAdapter, cfg: TransformerBridgeConfig
+    ) -> None:
+        attn_bridge = _wire_attention_bridge(adapter, cfg)
+        recordings = []
+        for name in ("q", "k", "v"):
+            projection = getattr(attn_bridge, name)
+            original = projection.original_component
+            assert isinstance(original, nn.Linear)
+            recording = _RecordingLinear(original.in_features, original.out_features, bias=False)
+            recording.load_state_dict(original.state_dict())
+            projection.set_original_component(recording)
+            recordings.append(recording)
+
+        hidden_states = torch.randn(1, 3, cfg.d_model, dtype=torch.bfloat16)
+        position_embeddings = identity_rope(3, cfg.d_head)
+
+        attn_bridge(
+            hidden_states=hidden_states,
+            position_embeddings=position_embeddings,
+            attention_mask=None,
+        )
+
+        for projection in recordings:
+            assert projection.received_dtype == torch.bfloat16
+
 
 class TestOlmoClipQkv:
     """The reconstructed forward must clamp Q/K/V when config.clip_qkv is set."""

transformer_lens/model_bridge/generalized_components/attention.py

@@ -813,42 +813,14 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
             raise RuntimeError(
                 f"Original component not set for {self.name}. Call set_original_component() first."
             )
-        # Skip non-fp params: quantized weights (bnb uint8/int8, GPTQ/AWQ int32,
-        # HQQ, torchao) are stored in integer dtypes and dequantized internally
-        # during matmul. The compute dtype must come from a fp parameter; casting
-        # fp inputs to an integer storage dtype destroys precision.
-        target_dtype = None
-        for p in self.original_component.parameters():
-            if not p.dtype.is_floating_point:
-                continue
-            target_dtype = p.dtype
-            break
         if "query_input" in kwargs:
             hooked = self.hook_in(kwargs["query_input"])
-            if (
-                target_dtype is not None
-                and isinstance(hooked, torch.Tensor)
-                and hooked.is_floating_point()
-            ):
-                hooked = hooked.to(dtype=target_dtype)
             kwargs["query_input"] = hooked
         elif "hidden_states" in kwargs:
             hooked = self.hook_in(kwargs["hidden_states"])
-            if (
-                target_dtype is not None
-                and isinstance(hooked, torch.Tensor)
-                and hooked.is_floating_point()
-            ):
-                hooked = hooked.to(dtype=target_dtype)
             kwargs["hidden_states"] = hooked
         elif len(args) > 0 and isinstance(args[0], torch.Tensor):
             hooked = self.hook_in(args[0])
-            if (
-                target_dtype is not None
-                and isinstance(hooked, torch.Tensor)
-                and hooked.is_floating_point()
-            ):
-                hooked = hooked.to(dtype=target_dtype)
             args = (hooked,) + args[1:]
         # try/finally so the captured tensor (and its autograd graph) is
         # released even if original_component raises.

transformer_lens/model_bridge/generalized_components/base.py

@@ -320,16 +320,6 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
             raise RuntimeError(
                 f"Original component not set for {self.name}. Call set_original_component() first."
             )
-        # Skip non-fp params: quantized weights (bnb uint8/int8, GPTQ/AWQ int32,
-        # HQQ, torchao) are stored in integer dtypes and dequantized internally
-        # during matmul. The compute dtype must come from a fp parameter; casting
-        # fp inputs to an integer storage dtype destroys precision.
-        target_dtype = None
-        for p in original_component.parameters():
-            if not p.dtype.is_floating_point:
-                continue
-            target_dtype = p.dtype
-            break
         input_arg_names = [
             "input",
             "hidden_states",
@@ -342,19 +332,11 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
         for name in input_arg_names:
             if name in kwargs:
                 hooked = self.hook_in(kwargs[name])
-                if (
-                    target_dtype is not None
-                    and isinstance(hooked, torch.Tensor)
-                    and hooked.is_floating_point()
-                ):
-                    hooked = hooked.to(dtype=target_dtype)
                 kwargs[name] = hooked
                 input_found = True
                 break
         if not input_found and len(args) > 0 and isinstance(args[0], torch.Tensor):
             hooked_input = self.hook_in(args[0])
-            if target_dtype is not None and hooked_input.is_floating_point():
-                hooked_input = hooked_input.to(dtype=target_dtype)
             args = (hooked_input,) + args[1:]
             input_found = True
         output = original_component(*args, **kwargs)

transformer_lens/model_bridge/generalized_components/moe.py

@@ -281,28 +281,11 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
             raise RuntimeError(
                 f"Original component not set for {self.name}. Call set_original_component() first."
             )
-        target_dtype = None
-        try:
-            target_dtype = next(self.original_component.parameters()).dtype
-        except StopIteration:
-            pass
         if len(args) > 0:
             hooked = self.hook_in(args[0])
-            if (
-                target_dtype is not None
-                and isinstance(hooked, torch.Tensor)
-                and hooked.is_floating_point()
-            ):
-                hooked = hooked.to(dtype=target_dtype)
             args = (hooked,) + args[1:]
         elif "hidden_states" in kwargs:
             hooked = self.hook_in(kwargs["hidden_states"])
-            if (
-                target_dtype is not None
-                and isinstance(hooked, torch.Tensor)
-                and hooked.is_floating_point()
-            ):
-                hooked = hooked.to(dtype=target_dtype)
             kwargs = {**kwargs, "hidden_states": hooked}
         output = self.original_component(*args, **kwargs)
         if isinstance(output, tuple):

transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py

@@ -348,20 +348,6 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
         # Apply input hook
         hidden_states = self.hook_in(hidden_states)
 
-        # Match dtype of HF module. Skip non-fp params: quantized weights (bnb
-        # uint8/int8, GPTQ/AWQ int32, HQQ, torchao) are stored in integer dtypes
-        # and dequantized internally during matmul. The compute dtype must come
-        # from a fp parameter; casting fp inputs to an integer storage dtype
-        # destroys precision.
-        target_dtype = None
-        for p in hf_attn.parameters():
-            if not p.dtype.is_floating_point:
-                continue
-            target_dtype = p.dtype
-            break
-        if target_dtype is not None and hidden_states.is_floating_point():
-            hidden_states = hidden_states.to(dtype=target_dtype)
-
         input_shape = hidden_states.shape[:-1]
         head_dim = hf_attn.head_dim
         hidden_shape = (*input_shape, -1, head_dim)

transformer_lens/model_bridge/generalized_components/unembedding.py

@@ -86,18 +86,7 @@ def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor:
             raise RuntimeError(
                 f"Original component not set for {self.name}. Call set_original_component() first."
             )
-        target_dtype = None
-        try:
-            target_dtype = next(self.original_component.parameters()).dtype
-        except StopIteration:
-            pass
         hidden_states = self.hook_in(hidden_states)
-        if (
-            target_dtype is not None
-            and isinstance(hidden_states, torch.Tensor)
-            and hidden_states.is_floating_point()
-        ):
-            hidden_states = hidden_states.to(dtype=target_dtype)
         output = self.original_component(hidden_states, **kwargs)
 
         output = self.hook_out(output)