fix(model): handle GEMMA4_CONVERSION_MODE=text for gemma4 MoE checkpoints

Signed-off-by: Josephasafg <ajgard7@gmail.com>

Author: Josephasafg

src/megatron/bridge/models/gemma/gemma4_bridge.py

@@ -504,71 +504,75 @@ def _hf_layer_prefix(self) -> str:
         """Text-only CausalLM: weights at ``model.*``; override in VL subclass."""
         return "model."
 
-    def _moe_mapping_registry(self) -> MegatronMappingRegistry:
+    def _moe_mapping_registry(self, megatron_prefix: str = "") -> MegatronMappingRegistry:
         """Parameter mappings for the MoE variant."""
+        mp = megatron_prefix
+        hp = self._hf_layer_prefix()
         param_mappings = {
-            "embedding.word_embeddings.weight": "model.embed_tokens.weight",
-            "decoder.final_layernorm.weight": "model.norm.weight",
-            "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight",
-            "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_norm.weight",
-            "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight",
-            "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight",
-            "decoder.layers.*.self_attention.linear_proj.post_layernorm.weight": (
-                "model.layers.*.post_attention_layernorm.weight"
+            f"{mp}embedding.word_embeddings.weight": f"{hp}embed_tokens.weight",
+            f"{mp}decoder.final_layernorm.weight": f"{hp}norm.weight",
+            f"{mp}decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": (
+                f"{hp}layers.*.input_layernorm.weight"
+            ),
+            f"{mp}decoder.layers.*.self_attention.q_layernorm.weight": f"{hp}layers.*.self_attn.q_norm.weight",
+            f"{mp}decoder.layers.*.self_attention.k_layernorm.weight": f"{hp}layers.*.self_attn.k_norm.weight",
+            f"{mp}decoder.layers.*.self_attention.linear_proj.weight": f"{hp}layers.*.self_attn.o_proj.weight",
+            f"{mp}decoder.layers.*.self_attention.linear_proj.post_layernorm.weight": (
+                f"{hp}layers.*.post_attention_layernorm.weight"
             ),
-            "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.pre_feedforward_layernorm_2.weight",
-            "decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight",
-            "decoder.layers.*.mlp.post_shared_expert_layernorm.weight": (
-                "model.layers.*.post_feedforward_layernorm_1.weight"
+            f"{mp}decoder.layers.*.pre_mlp_layernorm.weight": f"{hp}layers.*.pre_feedforward_layernorm_2.weight",
+            f"{mp}decoder.layers.*.mlp.shared_experts.linear_fc2.weight": f"{hp}layers.*.mlp.down_proj.weight",
+            f"{mp}decoder.layers.*.mlp.post_shared_expert_layernorm.weight": (
+                f"{hp}layers.*.post_feedforward_layernorm_1.weight"
             ),
-            "decoder.layers.*.mlp.router.weight": "model.layers.*.router.proj.weight",
+            f"{mp}decoder.layers.*.mlp.router.weight": f"{hp}layers.*.router.proj.weight",
         }
 
         mapping_list = [AutoMapping(megatron_param=m, hf_param=h) for m, h in param_mappings.items()]
         mapping_list.extend(
             [
                 _Gemma4QKVMapping(
-                    megatron_param="decoder.layers.*.self_attention.linear_qkv.weight",
-                    q="model.layers.*.self_attn.q_proj.weight",
-                    k="model.layers.*.self_attn.k_proj.weight",
-                    v="model.layers.*.self_attn.v_proj.weight",
+                    megatron_param=f"{mp}decoder.layers.*.self_attention.linear_qkv.weight",
+                    q=f"{hp}layers.*.self_attn.q_proj.weight",
+                    k=f"{hp}layers.*.self_attn.k_proj.weight",
+                    v=f"{hp}layers.*.self_attn.v_proj.weight",
                 ),
                 GatedMLPMapping(
-                    megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight",
-                    gate="model.layers.*.mlp.gate_proj.weight",
-                    up="model.layers.*.mlp.up_proj.weight",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.shared_experts.linear_fc1.weight",
+                    gate=f"{hp}layers.*.mlp.gate_proj.weight",
+                    up=f"{hp}layers.*.mlp.up_proj.weight",
                 ),
                 FusedGatedExpertMapping(
-                    megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*",
-                    hf_param="model.layers.*.experts.gate_up_proj",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.experts.linear_fc1.weight*",
+                    hf_param=f"{hp}layers.*.experts.gate_up_proj",
                 ),
                 FusedExpertMapping(
-                    megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*",
-                    hf_param="model.layers.*.experts.down_proj",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.experts.linear_fc2.weight*",
+                    hf_param=f"{hp}layers.*.experts.down_proj",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.layer_scalar",
-                    hf_param="model.layers.*.layer_scalar",
+                    megatron_param=f"{mp}decoder.layers.*.layer_scalar",
+                    hf_param=f"{hp}layers.*.layer_scalar",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.mlp.router.per_expert_scale",
-                    hf_param="model.layers.*.router.per_expert_scale",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.router.per_expert_scale",
+                    hf_param=f"{hp}layers.*.router.per_expert_scale",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.mlp.router.scale",
-                    hf_param="model.layers.*.router.scale",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.router.scale",
+                    hf_param=f"{hp}layers.*.router.scale",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.pre_shared_expert_layernorm.weight",
-                    hf_param="model.layers.*.pre_feedforward_layernorm.weight",
+                    megatron_param=f"{mp}decoder.layers.*.pre_shared_expert_layernorm.weight",
+                    hf_param=f"{hp}layers.*.pre_feedforward_layernorm.weight",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.mlp.post_moe_layernorm.weight",
-                    hf_param="model.layers.*.post_feedforward_layernorm_2.weight",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.post_moe_layernorm.weight",
+                    hf_param=f"{hp}layers.*.post_feedforward_layernorm_2.weight",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.post_ffn_layernorm.weight",
-                    hf_param="model.layers.*.post_feedforward_layernorm.weight",
+                    megatron_param=f"{mp}decoder.layers.*.post_ffn_layernorm.weight",
+                    hf_param=f"{hp}layers.*.post_feedforward_layernorm.weight",
                 ),
             ]
         )

src/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py

@@ -21,10 +21,14 @@
 
   AutoBridge.from_hf_pretrained("google/gemma-4-E4B-it")
     └─ Gemma4VLBridge (registered for Gemma4ForConditionalGeneration)
-         ├─ provider_bridge()  text mode → Gemma4DenseProvider (pretraining)
-         │                     auto/vl   → Gemma4DenseVLProvider (full VL)
-         └─ mapping_registry()  Dense → _dense_vl_mapping_registry()
-                                 MoE   → _moe_vl_mapping_registry()
+         ├─ provider_bridge()  Dense: text mode → Gemma4DenseProvider (pretraining)
+         │                            auto/vl   → Gemma4DenseVLProvider (full VL)
+         │                     MoE:   text mode → Gemma4ModelProvider (text-only training)
+         │                            auto/vl   → Gemma4VLModelProvider (full VL)
+         └─ mapping_registry()  Dense: text → _dense_mapping_registry("")
+                                        vl   → _dense_vl_mapping_registry()
+                                 MoE:   text → _moe_mapping_registry()
+                                        vl   → _moe_vl_mapping_registry()
 """
 
 import os
@@ -91,6 +95,8 @@ def provider_bridge(
             return self._build_dense_vl_provider(hf_config, text_config, vision_config)
 
         self._is_dense = False
+        if self._conversion_mode() == "text":
+            return self._build_moe_provider(text_config)
 
         provider_kwargs = self.hf_config_to_provider_kwargs(text_config)
         provider = Gemma4VLModelProvider(**provider_kwargs)
@@ -226,6 +232,8 @@ def mapping_registry(self) -> MegatronMappingRegistry:
             if self._conversion_mode() == "text":
                 return self._dense_mapping_registry(megatron_prefix="")
             return self._dense_vl_mapping_registry()
+        if self._conversion_mode() == "text":
+            return self._moe_mapping_registry()
         return self._moe_vl_mapping_registry()
 
     def _dense_vl_mapping_registry(self) -> MegatronMappingRegistry:

tests/unit_tests/models/gemma_vl/test_gemma4_vl_bridge.py

@@ -735,6 +735,26 @@ def test_vision_config_set(self, bridge, mock_hf_pretrained_moe):
         assert p.vision_config is mock_hf_pretrained_moe.config.vision_config
         assert p.text_config is mock_hf_pretrained_moe.config.text_config
 
+    def test_text_mode_returns_text_moe_provider(self, bridge, mock_hf_pretrained_moe, monkeypatch):
+        monkeypatch.setenv("GEMMA4_CONVERSION_MODE", "text")
+        p = bridge.provider_bridge(mock_hf_pretrained_moe)
+        assert type(p) is Gemma4ModelProvider
+
+    def test_text_mode_moe_provider_is_bf16(self, bridge, mock_hf_pretrained_moe, monkeypatch):
+        monkeypatch.setenv("GEMMA4_CONVERSION_MODE", "text")
+        p = bridge.provider_bridge(mock_hf_pretrained_moe)
+        assert p.bf16 is True
+        assert p.params_dtype == torch.bfloat16
+
+    def test_text_mode_moe_provider_config(self, bridge, mock_hf_pretrained_moe, monkeypatch):
+        monkeypatch.setenv("GEMMA4_CONVERSION_MODE", "text")
+        p = bridge.provider_bridge(mock_hf_pretrained_moe)
+        assert p.num_moe_experts == 128
+        assert p.moe_router_topk == 8
+        assert p.moe_ffn_hidden_size == 704
+        assert p.global_head_dim == 512
+        assert p.num_global_key_value_heads == 2
+
 
 class TestGemma4VLBridgeProviderBridgeDense:
     def test_accepts_dense_with_per_layer_inputs(self, bridge, mock_hf_pretrained_dense):
@@ -947,6 +967,22 @@ def test_uses_language_model_prefix_for_vl(self, bridge, mock_hf_config_moe):
         lm_keys = [n for n in names if "layers" in n and "vision" not in n and "audio" not in n]
         assert any("language_model" in n for n in lm_keys)
 
+    def test_moe_text_mode_registry_targets_text_model_from_vl_checkpoint(
+        self, bridge, mock_hf_config_moe, monkeypatch
+    ):
+        """MoE text mode: unprefixed Megatron params sourced from model.language_model.* HF weights."""
+        monkeypatch.setenv("GEMMA4_CONVERSION_MODE", "text")
+        bridge.hf_config = mock_hf_config_moe
+        registry = bridge.mapping_registry()
+
+        megatron_params = [str(m.megatron_param) for m in registry.mappings if hasattr(m, "megatron_param")]
+        assert "embedding.word_embeddings.weight" in megatron_params
+        assert all(not n.startswith("language_model.") for n in megatron_params)
+        assert not any("vision_tower" in n or "audio_tower" in n for n in megatron_params)
+
+        hf_targets = self._collect_hf_targets(registry)
+        assert all(n.startswith("model.language_model.") for n in hf_targets)
+
     def test_dense_vl_audio_tower_replicated_mappings(self, bridge, mock_hf_config_dense, monkeypatch):
         monkeypatch.setenv("GEMMA4_CONVERSION_MODE", "vl")
         bridge.hf_config = mock_hf_config_dense