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

Move MiniMax-M2 onto the ModelConfig/ModelBuilder path, following Muse
Glimmer. MiniMaxM2Bridge maps Hugging Face configs in
hf_config_to_model_config_kwargs(), including the full-dimension QK-norm
layer spec, and the deprecated provider path reuses the same kwargs.
USE_MODEL_CONFIG_FOR_CONVERSION = True.

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/minimax_m2/minimax_m2_bridge.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 
 from collections.abc import Mapping
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 import torch
 import torch.nn as nn
 from megatron.core.models.gpt.gpt_model import GPTModel
+from transformers import PretrainedConfig
 
 from megatron.bridge.models.conversion import quantization_utils
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
@@ -128,46 +129,52 @@ class MiniMaxM2Bridge(MegatronModelBridge):
     Example:
         >>> from megatron.bridge import AutoBridge
         >>> bridge = AutoBridge.from_hf_pretrained("MiniMaxAI/MiniMax-M2")
-        >>> provider = bridge.to_megatron_provider()
+        >>> model_config = bridge.get_model_config()
     """
 
-    def provider_bridge(self, hf_pretrained):
-        """Convert HuggingFace MiniMax-M2 config to GPTModelProvider."""
-        provider = super().provider_bridge(hf_pretrained)
-
-        hf_config = hf_pretrained.config
-
-        provider.normalization = "RMSNorm"
-        provider.gated_linear_unit = True
-        provider.position_embedding_type = "rope"
-        provider.add_bias_linear = False
-        provider.add_qkv_bias = False
-        provider.hidden_dropout = 0.0
-        provider.autocast_dtype = torch.bfloat16
+    USE_MODEL_CONFIG_FOR_CONVERSION = True
+
+    def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Convert a Hugging Face MiniMax-M2 config to builder config kwargs."""
+        config_kwargs = super().hf_config_to_model_config_kwargs(hf_config)
+        config_kwargs.update(
+            normalization="RMSNorm",
+            gated_linear_unit=True,
+            position_embedding_type="rope",
+            add_bias_linear=False,
+            add_qkv_bias=False,
+            hidden_dropout=0.0,
+            autocast_dtype=torch.bfloat16,
+            masked_softmax_fusion=True,
+            rope_scaling=False,
+            rope_scaling_factor=1.0,
+            # Full-dimension QK norm via custom layer spec (see minimax_m2_provider.py).
+            # qk_layernorm=True so mcore creates QK norms; the spec overrides the default
+            # TENorm with FullDimQNorm/FullDimKNorm for full-dimension normalization.
+            qk_layernorm=True,
+            transformer_layer_spec=minimax_m2_layer_spec,
+            # MoE settings — sigmoid routing with expert bias (same pattern as DeepSeek V3)
+            moe_grouped_gemm=True,
+            moe_router_pre_softmax=False,
+            moe_router_load_balancing_type="aux_loss",
+            moe_aux_loss_coeff=getattr(hf_config, "router_aux_loss_coef", 1e-3),
+            moe_token_dispatcher_type="alltoall",
+            moe_permute_fusion=True,
+            moe_router_score_function="sigmoid",
+            moe_router_enable_expert_bias=True,
+        )
 
         # MiniMax-M2 uses rotary_dim instead of partial_rotary_factor
         rotary_dim = getattr(hf_config, "rotary_dim", None)
         head_dim = getattr(hf_config, "head_dim", None)
         if rotary_dim is not None and head_dim is not None:
-            provider.rotary_percent = rotary_dim / head_dim
-
-        # Full-dimension QK norm via custom layer spec (see minimax_m2_provider.py).
-        # qk_layernorm=True so mcore creates QK norms; the spec overrides the default
-        # TENorm with FullDimQNorm/FullDimKNorm for full-dimension normalization.
-        provider.qk_layernorm = True
-        provider.transformer_layer_spec = minimax_m2_layer_spec
-
-        # MoE settings — sigmoid routing with expert bias (same pattern as DeepSeek V3)
-        provider.moe_grouped_gemm = True
-        provider.moe_router_pre_softmax = False
-        provider.moe_router_load_balancing_type = "aux_loss"
-        provider.moe_aux_loss_coeff = getattr(hf_config, "router_aux_loss_coef", 1e-3)
-        provider.moe_token_dispatcher_type = "alltoall"
-        provider.moe_permute_fusion = True
-        provider.moe_router_score_function = "sigmoid"
-        provider.moe_router_enable_expert_bias = True
-
-        return provider
+            config_kwargs["rotary_percent"] = rotary_dim / head_dim
+
+        return config_kwargs
+
+    def hf_config_to_provider_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Adapt the canonical builder mapping to the deprecated provider path."""
+        return self.hf_config_to_model_config_kwargs(hf_config)
 
     def maybe_modify_loaded_hf_weight(
         self, hf_param: str | dict[str, str], hf_state_dict: Mapping[str, torch.Tensor]

tests/unit_tests/models/minimax_m2/test_minimax_m2_bridges.py

@@ -33,6 +33,7 @@
     FullDimKNorm,
     FullDimQNorm,
     _FullDimRMSNorm,
+    minimax_m2_layer_spec,
 )
 
 
@@ -503,3 +504,89 @@ def test_megatron_to_hf_deduplicates_replicated_kv_shards(self):
             result["model.layers.*.self_attn.k_norm.weight"],
             torch.arange(8, dtype=torch.bfloat16),
         )
+
+
+class TestMiniMaxM2BuilderConfig:
+    """Builder-backed ModelConfig path for MiniMax-M2."""
+
+    @staticmethod
+    def _config(**attrs):
+        from transformers import MiniMaxM2Config
+
+        config = MiniMaxM2Config(
+            hidden_size=64,
+            intermediate_size=32,
+            num_hidden_layers=2,
+            num_attention_heads=8,
+            num_key_value_heads=4,
+            head_dim=8,
+            num_local_experts=8,
+            num_experts_per_tok=2,
+            max_position_embeddings=4096,
+            vocab_size=128,
+            torch_dtype="bfloat16",
+        )
+        config.architectures = ["MiniMaxM2ForCausalLM"]
+        for key, value in attrs.items():
+            setattr(config, key, value)
+        return config
+
+    def test_provider_bridge_is_inherited_compatibility_only(self):
+        assert "provider_bridge" not in MiniMaxM2Bridge.__dict__
+
+    def test_conversion_uses_builder_config(self):
+        assert MiniMaxM2Bridge.USE_MODEL_CONFIG_FOR_CONVERSION is True
+
+    def test_hf_config_to_model_config_uses_direct_mapping(self):
+        from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
+
+        config = self._config(rotary_dim=4)
+        bridge = MiniMaxM2Bridge()
+        with (
+            patch.object(bridge, "provider_bridge", side_effect=AssertionError("provider path used")),
+            patch.object(bridge, "hf_config_to_provider_kwargs", side_effect=AssertionError("provider kwargs used")),
+        ):
+            result = bridge.hf_config_to_model_config(config)
+
+        assert isinstance(result, BridgeGPTModelConfig)
+        assert result.transformer_layer_spec is minimax_m2_layer_spec
+        assert result.qk_layernorm is True
+        assert result.rotary_percent == 0.5
+        assert result.position_embedding_type == "rope"
+        assert result.moe_router_score_function == "sigmoid"
+        assert result.moe_router_enable_expert_bias is True
+        assert result.moe_aux_loss_coeff == config.router_aux_loss_coef
+
+    @pytest.mark.parametrize("attrs", [{}, {"rotary_dim": 4}], ids=["full-rotary", "rotary-dim"])
+    def test_model_config_matches_provider_runtime_config(self, attrs):
+        from dataclasses import fields
+
+        config = self._config(**attrs)
+        mock_pretrained = Mock(spec=PreTrainedCausalLM)
+        mock_pretrained.config = config
+        bridge = MiniMaxM2Bridge()
+
+        model_config = bridge.hf_config_to_model_config(config)
+        with pytest.warns(FutureWarning, match=r"deprecated.*get_model_config.*get_model"):
+            provider = bridge.provider_bridge(mock_pretrained)
+
+        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
+        # Both paths build the full-dimension QK norm through the same layer spec.
+        assert model_config.transformer_layer_spec is minimax_m2_layer_spec
+        assert provider.transformer_layer_spec is minimax_m2_layer_spec
+
+    def test_layer_spec_matches_between_provider_and_builder_config(self):
+        config = self._config()
+        mock_pretrained = Mock(spec=PreTrainedCausalLM)
+        mock_pretrained.config = config
+        bridge = MiniMaxM2Bridge()
+        model_config = bridge.hf_config_to_model_config(config)
+        with pytest.warns(FutureWarning):
+            provider = bridge.provider_bridge(mock_pretrained)
+
+        # functools.partial (used for the MoE layer) compares by identity, so compare structure.
+        assert repr(minimax_m2_layer_spec(model_config)) == repr(minimax_m2_layer_spec(provider))