fix: register unmapped leaf parameters so module conversions reach them

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

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(

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