diff --git a/tests/integration/model_bridge/test_gemma3n_bridge.py b/tests/integration/model_bridge/test_gemma3n_bridge.py index ae3318386..92b592ecb 100644 --- a/tests/integration/model_bridge/test_gemma3n_bridge.py +++ b/tests/integration/model_bridge/test_gemma3n_bridge.py @@ -99,3 +99,17 @@ def test_contribution_hooks_are_post_norm_outputs(bridge): for key, expected in reference.items(): torch.testing.assert_close(cache[key], expected, atol=1e-5, rtol=1e-5) + + +def test_half_converts_unmapped_parameters(): + bridge = TransformerBridge.boot_transformers(MODEL_NAME, device="cpu", dtype=torch.float32) + bridge.half() + + unconverted = [ + name + for name, parameter in bridge.original_model.named_parameters() + if parameter.dtype != torch.float16 + ] + assert unconverted == [] + with torch.no_grad(): + assert bridge(IDS).dtype == torch.float16 diff --git a/tests/integration/model_bridge/test_parent_module_traversal.py b/tests/integration/model_bridge/test_parent_module_traversal.py index 3d80d25d1..3d6ebb3a4 100644 --- a/tests/integration/model_bridge/test_parent_module_traversal.py +++ b/tests/integration/model_bridge/test_parent_module_traversal.py @@ -17,6 +17,12 @@ BertForPreTraining, BloomConfig, BloomForCausalLM, + Gemma4AudioConfig, + Gemma4Config, + Gemma4ForCausalLM, + Gemma4ForConditionalGeneration, + Gemma4TextConfig, + Gemma4VisionConfig, GPT2Config, GPT2LMHeadModel, GPTNeoXConfig, @@ -36,6 +42,7 @@ ViTModel, ) +from tests.integration.model_bridge.helpers import make_tiny_pair from transformer_lens.model_bridge.sources import build_bridge_from_module @@ -100,6 +107,42 @@ def _hubert_config() -> HubertConfig: ) +def _gemma4_text_config() -> Gemma4TextConfig: + return Gemma4TextConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + max_position_embeddings=16, + vocab_size_per_layer_input=32, + hidden_size_per_layer_input=4, + ) + + +def _gemma4_config() -> Gemma4Config: + return Gemma4Config( + text_config=_gemma4_text_config(), + vision_config=Gemma4VisionConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=4, + ), + audio_config=Gemma4AudioConfig( + hidden_size=16, + num_attention_heads=4, + num_hidden_layers=1, + output_proj_dims=16, + subsampling_conv_channels=[8, 8], + ), + ) + + ARCHITECTURE_CASES = ( ArchitectureCase( "gpt2-joint-qkv", @@ -240,6 +283,18 @@ def _hubert_config() -> HubertConfig: ), "ASTForAudioClassification", ), + ArchitectureCase( + "gemma4-per-layer-embeddings", + Gemma4ForCausalLM, + _gemma4_text_config, + "Gemma4ForCausalLM", + ), + ArchitectureCase( + "gemma4-audio-vision", + Gemma4ForConditionalGeneration, + _gemma4_config, + "Gemma4ForConditionalGeneration", + ), ) ARCHITECTURE_CASE_BY_NAME = {case.name: case for case in ARCHITECTURE_CASES} @@ -304,6 +359,28 @@ def test_parent_dtype_conversion_updates_container_owned_state() -> None: } +@pytest.mark.parametrize("through_parent", (False, True), ids=("bridge-half", "parent-to")) +def test_dtype_conversion_updates_unmapped_leaf_parameters(through_parent: bool) -> None: + bridge, reference = make_tiny_pair( + _gemma4_text_config(), "Gemma4ForCausalLM", loader=Gemma4ForCausalLM + ) + if through_parent: + parent = nn.Module() + parent.add_module("bridge", bridge) + parent.to(torch.float16) + else: + bridge.half() + reference.half() + tokens = torch.tensor([[2, 5, 7, 9, 11]]) + + with torch.no_grad(): + logits = bridge(tokens) + expected = reference(input_ids=tokens).logits + + assert logits.dtype == torch.float16 + assert torch.equal(logits, expected) + + def test_parent_assign_load_updates_container_owned_state() -> None: bart_config = _bart_config() bridge = build_bridge_from_module( diff --git a/transformer_lens/model_bridge/component_setup.py b/transformer_lens/model_bridge/component_setup.py index 9064e8e09..db6b90d31 100644 --- a/transformer_lens/model_bridge/component_setup.py +++ b/transformer_lens/model_bridge/component_setup.py @@ -135,11 +135,10 @@ def _register_unowned_container_state(bridge_module: nn.Module, original_model: missing_buffers: list[tuple[str, nn.Module, str, Any]] = [] for module_path, module in original_model.named_modules(): - if any(child is not None for child in module._modules.values()): - for parameter_name, parameter in module._parameters.items(): - if parameter is not None and id(parameter) not in registered_parameter_ids: - missing_parameters.append((module_path, module, parameter_name, parameter)) - registered_parameter_ids.add(id(parameter)) + for parameter_name, parameter in module._parameters.items(): + if parameter is not None and id(parameter) not in registered_parameter_ids: + missing_parameters.append((module_path, module, parameter_name, parameter)) + registered_parameter_ids.add(id(parameter)) for buffer_name, buffer in module._buffers.items(): if buffer is not None and id(buffer) not in registered_buffer_ids: missing_buffers.append((module_path, module, buffer_name, buffer))