refactor(model): build MiMo from a builder-backed model config

MimoBridge extends the Qwen2 builder kwargs with its attention settings and
MTP depth (mtp_num_layers, mtp_loss_scaling_factor) instead of overriding
provider_bridge(). It inherits the Qwen2 conversion flag, so MiMo checkpoints
convert through the ModelBuilder as well.

Signed-off-by: Kamran Jafari <kjafarisadeg@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: kamran-nvidia

src/megatron/bridge/models/mimo/mimo_bridge.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Mapping
+from typing import Any, Mapping
 
 import torch
 from megatron.core.models.gpt.gpt_model import GPTModel
+from transformers import PretrainedConfig
 
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
 from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge, WeightConversionTask
@@ -31,20 +32,18 @@
 class MimoBridge(Qwen2Bridge):
     """Megatron Bridge for MiMo Causal LM."""
 
-    def provider_bridge(self, hf_pretrained):
-        provider = super().provider_bridge(hf_pretrained)
-        hf_config = hf_pretrained.config
+    def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Convert a Hugging Face MiMo config to builder config kwargs."""
+        config_kwargs = super().hf_config_to_model_config_kwargs(hf_config)
 
         # MiMo follows Qwen2 attention behavior and adds MTP on top.
-        provider.qk_layernorm = False
-        provider.add_qkv_bias = True
+        config_kwargs.update(qk_layernorm=False, add_qkv_bias=True)
 
         num_mtp_layers = getattr(hf_config, "num_nextn_predict_layers", 0)
         if num_mtp_layers > 0:
-            provider.mtp_num_layers = num_mtp_layers
-            provider.mtp_loss_scaling_factor = 0.1
+            config_kwargs.update(mtp_num_layers=num_mtp_layers, mtp_loss_scaling_factor=0.1)
 
-        return provider
+        return config_kwargs
 
     def mapping_registry(self) -> MegatronMappingRegistry:
         mapping_list = list(super().mapping_registry().mappings)

tests/unit_tests/models/mimo/test_mimo_bridge.py

@@ -126,3 +126,37 @@ def test_mtp_input_proj_swap_on_hf_export(self):
 
         expected = torch.cat((weight[:, 4:], weight[:, :4]), dim=1)
         assert torch.equal(modified["model.mtp_layers.0.input_proj.weight"], expected)
+
+    def test_provider_bridge_is_inherited_compatibility_only(self):
+        """MiMo builds both paths from its builder-config kwargs."""
+        assert "provider_bridge" not in MimoBridge.__dict__
+
+    def test_conversion_uses_builder_config(self):
+        assert MimoBridge.USE_MODEL_CONFIG_FOR_CONVERSION is True
+
+    def test_model_config_maps_mtp_config(self, mock_pretrained_mimo):
+        from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
+
+        model_config = MimoBridge().hf_config_to_model_config(mock_pretrained_mimo.config)
+
+        assert isinstance(model_config, BridgeGPTModelConfig)
+        assert model_config.qk_layernorm is False
+        assert model_config.add_qkv_bias is True
+        assert model_config.mtp_num_layers == mock_pretrained_mimo.config.num_nextn_predict_layers
+        assert model_config.mtp_loss_scaling_factor == 0.1
+        assert model_config.position_embedding_type == "rope"
+        assert model_config.normalization == "RMSNorm"
+
+    def test_model_config_matches_provider_runtime_config(self, mock_pretrained_mimo):
+        from dataclasses import fields
+
+        bridge = MimoBridge()
+        model_config = bridge.hf_config_to_model_config(mock_pretrained_mimo.config)
+        with pytest.warns(FutureWarning, match=r"deprecated.*get_model_config.*get_model"):
+            provider = bridge.provider_bridge(mock_pretrained_mimo)
+
+        provider_fields = {field.name for field in fields(provider)}
+        model_config_fields = {field.name for field in fields(model_config)}
+        model_config_fields.update(field.name for field in fields(model_config.transformer))
+        for field_name in sorted((provider_fields & model_config_fields) - {"transformer_layer_spec"}):
+            assert getattr(model_config, field_name) == getattr(provider, field_name), field_name