[Bug Report] .half() and parent-module conversions skip some Gemma 4 E2B/E4B and Gemma 3n weights #1935

Describe the bug

On Gemma 4 models with per-layer embeddings (E2B/E4B) and on Gemma 3n, .half(), .bfloat16(), .float(), .double() and .type() leave weights that the text forward uses in the old dtype. Dtype and device conversions of a parent module that holds the bridge (parent.half(), parent.to(torch.float16), parent.to("mps")) miss the same weights. bridge.to() works, because it also calls original_model.to().

These conversions go through nn.Module._apply, which only reaches the bridge's registered tree. _register_unowned_container_state (transformer_lens/model_bridge/component_setup.py:120-178, from #1671) registers unmapped buffers of every module, but unmapped parameters only of modules that have children (:137-142). The adapters leave these leaf modules unmapped: embed_tokens_per_layer, per_layer_model_projection and per_layer_projection_norm on both families, and the Linear layers in altup_projections and altup_unembed_projections on Gemma 3n. Their weights are also missing from parent.parameters(), so parent.requires_grad_(False) leaves them trainable and an optimizer built from the parent skips them.

Code example

import torch
from transformer_lens.model_bridge import TransformerBridge

bridge = TransformerBridge.boot_transformers("tiny-random/gemma-4-e", device="cpu")
bridge.half()
bridge(torch.tensor([[1, 2, 3, 4, 5]]))
# RuntimeError: expected m1 and m2 to have the same dtype, but got: c10::Half != float

On CPU, onnx-internal-testing/tiny-random-Gemma3nForConditionalGeneration raises the same error. On MPS, the .half() above aborts the process (MPSNDArrayMatrixMultiplication datatype assertion). Gemma 3n's .half() and Gemma 4's .float() on a bf16 bridge do not fail there: they run with those weights unconverted, and their logits differ from the matching .to() by about 1e-3 on these tiny models. With either tiny model, parent.to("mps") leaves the missed weights on CPU, and the forward raises RuntimeError: Placeholder storage has not been allocated on MPS device!.

System Info

• Source checkout installed with uv, dev at a3dff503.
• macOS 27.2, Apple M5; Python 3.12.14; PyTorch 2.11.0; transformers 5.13.0; CPU and MPS.

Additional context

Proposed fix: drop the children check, so unmapped parameters get the same owner views as unmapped buffers. That also reaches modules the text path does not use: audio_tower and embed_audio on Gemma 4 E2B/E4B; audio_tower, vision_tower, embed_audio and embed_vision on Gemma 3n; vision_tower and multi_modal_projector on T5Gemma2; and the VQ image tokenizer vqmodel on Emu3. On tiny configs built from cached remote code, three more adapters have the same gap: Raven's transformer.adapter, Ouro's model.early_exit_gate and BD3LM's backbone.output_layer.adaLN_modulation when adaln=True. The fix makes all three reachable, and on Raven it also fixes the forward after .half() on CPU, which raises the same error on dev.

Direct bridge.state_dict() keys do not change. The owner views add no hooks, and test_parent_and_direct_traversal_have_identical_state (#1671) would no longer catch an adapter that leaves a leaf unmapped.

Mapping the modules in the adapters, as #1671 did for BERT, AST and ViT, adds hooks and keeps that check. But the Gemma 4 mapping needs a condition (31B, 26B-A4B and 12B have no per-layer embeddings), it renames the mapped keys in bridge.state_dict(), and the towers' leaf parameters stay unreachable unless they are mapped too. Raven and Ouro leave theirs unmapped on purpose.

Either way, parent.state_dict() gains the newly reachable parameter keys (_container_state_owners.* with the proposed fix), so parent checkpoints saved before the fix report missing keys with strict=True.

Should the fix go in _register_unowned_container_state, or would you rather map the modules in the adapters?

Checklist

• I have checked that there is no similar issue in the repo.