From 813ec4d516941d76adf319c87996b0f512ed43e4 Mon Sep 17 00:00:00 2001 From: kigland Date: Fri, 21 Aug 2026 13:34:03 +0800 Subject: [PATCH 1/2] preserve generalized component input dtype --- .../generalized_components/test_base.py | 26 +++++++++++++++++++ .../generalized_components/base.py | 18 ------------- 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/tests/unit/model_bridge/generalized_components/test_base.py b/tests/unit/model_bridge/generalized_components/test_base.py index df3c9a152..ffff1d061 100644 --- a/tests/unit/model_bridge/generalized_components/test_base.py +++ b/tests/unit/model_bridge/generalized_components/test_base.py @@ -377,5 +377,31 @@ 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) + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/transformer_lens/model_bridge/generalized_components/base.py b/transformer_lens/model_bridge/generalized_components/base.py index af8087282..3d092489a 100644 --- a/transformer_lens/model_bridge/generalized_components/base.py +++ b/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) From 22dad242d0edcdb7ea44ec48e5e7cd781ca16e78 Mon Sep 17 00:00:00 2001 From: kigland Date: Sat, 22 Aug 2026 02:03:11 +0800 Subject: [PATCH 2/2] preserve dtype in specialized bridge forwards --- .../generalized_components/test_base.py | 55 +++++++++++++++++++ .../test_moe_bridge_tuple_output.py | 17 ++++++ .../test_olmo_adapter.py | 34 ++++++++++++ .../generalized_components/attention.py | 28 ---------- .../generalized_components/moe.py | 17 ------ .../position_embeddings_attention.py | 14 ----- .../generalized_components/unembedding.py | 11 ---- 7 files changed, 106 insertions(+), 70 deletions(-) diff --git a/tests/unit/model_bridge/generalized_components/test_base.py b/tests/unit/model_bridge/generalized_components/test_base.py index ffff1d061..23754c5dd 100644 --- a/tests/unit/model_bridge/generalized_components/test_base.py +++ b/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): @@ -403,5 +409,54 @@ def forward(self, hidden_states): 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__]) diff --git a/tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py b/tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py index 3b644497b..c321c9ec4 100644 --- a/tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py +++ b/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"): diff --git a/tests/unit/model_bridge/supported_architectures/test_olmo_adapter.py b/tests/unit/model_bridge/supported_architectures/test_olmo_adapter.py index ab0f5d04d..201d2381d 100644 --- a/tests/unit/model_bridge/supported_architectures/test_olmo_adapter.py +++ b/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.""" diff --git a/transformer_lens/model_bridge/generalized_components/attention.py b/transformer_lens/model_bridge/generalized_components/attention.py index 73770f31c..2fbd820ce 100644 --- a/transformer_lens/model_bridge/generalized_components/attention.py +++ b/transformer_lens/model_bridge/generalized_components/attention.py @@ -797,42 +797,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. diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py index 334b7dc51..19ba06ae6 100644 --- a/transformer_lens/model_bridge/generalized_components/moe.py +++ b/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): diff --git a/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py b/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py index f5dbbdaf2..a26b62516 100644 --- a/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py +++ b/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) diff --git a/transformer_lens/model_bridge/generalized_components/unembedding.py b/transformer_lens/model_bridge/generalized_components/unembedding.py index ac7ab7e74..7c43e1bd9 100644 --- a/transformer_lens/model_bridge/generalized_components/unembedding.py +++ b/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)