refactor(model): build Hy V3 and MiniMax-M2 from builder-backed model configs (#6392)

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/hy_v3/hy_v3_bridge.py

@@ -13,10 +13,12 @@
 # limitations under the License.
 
 from functools import partial
+from typing import Any
 
 import torch
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
 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
@@ -25,8 +27,8 @@
     GatedMLPMapping,
     QKVMapping,
 )
+from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
 from megatron.bridge.models.gpt_provider import GPTModelProvider
-from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
 
 
 try:
@@ -57,49 +59,62 @@ class HYV3Bridge(MegatronModelBridge):
     Example:
         >>> from megatron.bridge import AutoBridge
         >>> bridge = AutoBridge.from_hf_pretrained("tencent/Hy3-preview-Base")
-        >>> provider = bridge.to_megatron_provider()
+        >>> model_config = bridge.get_model_config()
     """
 
-    def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider:
-        """Convert HuggingFace Hy V3 config to GPTModelProvider."""
-        provider = super().provider_bridge(hf_pretrained)
-        hf_config = hf_pretrained.config
+    MODEL_CONFIG_CLASS = BridgeGPTModelConfig
+    USE_MODEL_CONFIG_FOR_CONVERSION = True
 
-        provider.transformer_layer_spec = partial(get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE)
+    def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Convert a Hugging Face Hy V3 config to builder config kwargs.
 
-        provider.normalization = "RMSNorm"
-        provider.gated_linear_unit = True
-        provider.add_bias_linear = False
-        provider.add_qkv_bias = False  # Hy V3 does NOT have QKV bias
-        provider.hidden_dropout = 0.0
-        provider.qk_layernorm = True  # Hy V3 uses QK layernorm
-        provider.attention_softmax_in_fp32 = False
-        provider.autocast_dtype = torch.bfloat16
-        provider.fp16 = False
-        provider.bf16 = True
-        provider.params_dtype = torch.bfloat16
-        provider.moe_grouped_gemm = True
-        provider.moe_token_dispatcher_type = "alltoall"
-        provider.moe_router_load_balancing_type = "none"
-        provider.moe_router_pre_softmax = False
-        provider.moe_router_score_function = "sigmoid"
-        provider.moe_router_enable_expert_bias = True
-        provider.moe_router_bias_update_rate = 0
-        provider.moe_router_dtype = "fp32"
-        provider.moe_permute_fusion = True
-        provider.moe_shared_expert_overlap = False
-        provider.moe_aux_loss_coeff = 0.0
-        provider.moe_router_topk_scaling_factor = float(hf_config.router_scaling_factor)
-
-        # Shared experts hidden dim = moe_intermediate_size * num_shared_experts.
-        provider.moe_shared_expert_intermediate_size = hf_config.moe_intermediate_size * hf_config.num_shared_experts
-
-        # Dense-first layer pattern.
-        provider.moe_layer_freq = [0] * hf_config.first_k_dense_replace + [1] * (
-            hf_config.num_hidden_layers - hf_config.first_k_dense_replace
+        The builder's default MoE layer spec already uses the GPT decoder block
+        spec, which honors the dense-first ``moe_layer_freq``.
+        """
+        config_kwargs = super().hf_config_to_model_config_kwargs(hf_config)
+        config_kwargs.update(
+            normalization="RMSNorm",
+            gated_linear_unit=True,
+            add_bias_linear=False,
+            add_qkv_bias=False,  # Hy V3 does NOT have QKV bias
+            hidden_dropout=0.0,
+            qk_layernorm=True,  # Hy V3 uses QK layernorm
+            attention_softmax_in_fp32=False,
+            autocast_dtype=torch.bfloat16,
+            fp16=False,
+            bf16=True,
+            params_dtype=torch.bfloat16,
+            masked_softmax_fusion=True,
+            rope_scaling=False,
+            rope_scaling_factor=1.0,
+            moe_grouped_gemm=True,
+            moe_token_dispatcher_type="alltoall",
+            moe_router_load_balancing_type="none",
+            moe_router_pre_softmax=False,
+            moe_router_score_function="sigmoid",
+            moe_router_enable_expert_bias=True,
+            moe_router_bias_update_rate=0,
+            moe_router_dtype="fp32",
+            moe_permute_fusion=True,
+            moe_shared_expert_overlap=False,
+            moe_aux_loss_coeff=0.0,
+            moe_router_topk_scaling_factor=float(hf_config.router_scaling_factor),
+            # Shared experts hidden dim = moe_intermediate_size * num_shared_experts.
+            moe_shared_expert_intermediate_size=hf_config.moe_intermediate_size * hf_config.num_shared_experts,
+            # Dense-first layer pattern.
+            moe_layer_freq=[0] * hf_config.first_k_dense_replace
+            + [1] * (hf_config.num_hidden_layers - hf_config.first_k_dense_replace),
         )
+        config_kwargs.setdefault("position_embedding_type", "rope")
+        return config_kwargs
 
-        return provider
+    def hf_config_to_provider_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Adapt the canonical builder mapping to the deprecated provider path."""
+        provider_kwargs = self.hf_config_to_model_config_kwargs(hf_config)
+        # The provider's default layer spec builds homogeneous layers; the decoder
+        # block spec is needed for the dense-first MoE layer pattern.
+        provider_kwargs["transformer_layer_spec"] = partial(get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE)
+        return provider_kwargs
 
     @classmethod
     def megatron_to_hf_config(cls, provider: GPTModelProvider) -> dict:

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
@@ -28,6 +29,7 @@
     MegatronParamMapping,
     QKVMapping,
 )
+from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
 from megatron.bridge.models.minimax_m2.minimax_m2_provider import minimax_m2_layer_spec
 
 
@@ -128,46 +130,53 @@ 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
+    MODEL_CONFIG_CLASS = BridgeGPTModelConfig
+    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/hy_v3/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unit_tests/models/hy_v3/test_hy_v3_bridge.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import fields
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
+from transformers import HYV3Config
+
+from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
+from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
+from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
+from megatron.bridge.models.hy_v3.hy_v3_bridge import HYV3Bridge
+
+
+def _hy_v3_config(**overrides) -> HYV3Config:
+    kwargs = dict(
+        hidden_size=64,
+        intermediate_size=128,
+        moe_intermediate_size=32,
+        num_hidden_layers=4,
+        num_attention_heads=8,
+        num_key_value_heads=4,
+        head_dim=8,
+        num_experts=8,
+        num_experts_per_tok=2,
+        num_shared_experts=1,
+        first_k_dense_replace=1,
+        router_scaling_factor=2.5,
+        max_position_embeddings=4096,
+        vocab_size=128,
+        torch_dtype="bfloat16",
+    )
+    kwargs.update(overrides)
+    config = HYV3Config(**kwargs)
+    config.architectures = ["HYV3ForCausalLM"]
+    return config
+
+
+class TestHYV3Bridge:
+    def test_bridge_registration(self):
+        assert issubclass(HYV3Bridge, MegatronModelBridge)
+
+    def test_provider_bridge_is_inherited_compatibility_only(self):
+        assert "provider_bridge" not in HYV3Bridge.__dict__
+
+    def test_conversion_uses_builder_config(self):
+        assert HYV3Bridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
+        assert HYV3Bridge.USE_MODEL_CONFIG_FOR_CONVERSION is True
+
+    def test_hf_config_to_model_config_uses_direct_mapping(self):
+        config = _hy_v3_config()
+        bridge = HYV3Bridge()
+        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)
+        # The builder's default MoE spec is the decoder block spec, so none is configured here.
+        assert result.transformer_layer_spec is None
+        assert result.moe_layer_freq == [0, 1, 1, 1]
+        assert result.moe_shared_expert_intermediate_size == config.moe_intermediate_size * config.num_shared_experts
+        assert result.moe_router_topk_scaling_factor == 2.5
+        assert result.moe_router_score_function == "sigmoid"
+        assert result.moe_router_enable_expert_bias is True
+        assert result.qk_layernorm is True
+        assert result.params_dtype == torch.bfloat16
+        assert result.bf16 is True
+        assert result.position_embedding_type == "rope"
+
+    @pytest.mark.parametrize("torch_dtype", ["bfloat16", "float16"])
+    def test_model_config_matches_provider_runtime_config(self, torch_dtype):
+        config = _hy_v3_config(torch_dtype=torch_dtype)
+        mock_pretrained = Mock(spec=PreTrainedCausalLM)
+        mock_pretrained.config = config
+        bridge = HYV3Bridge()
+
+        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
+        # Hy V3 always builds in bf16, whatever the source checkpoint dtype.
+        assert model_config.params_dtype == provider.params_dtype == torch.bfloat16
+        # The provider keeps its explicit decoder block spec for the dense-first layer pattern.
+        assert provider.transformer_layer_spec.func is get_gpt_decoder_block_spec
+
+    def test_megatron_to_hf_config_from_model_config(self):
+        model_config = HYV3Bridge().hf_config_to_model_config(_hy_v3_config())
+
+        hf_config = HYV3Bridge.megatron_to_hf_config(model_config)
+
+        assert hf_config["first_k_dense_replace"] == 1
+        assert hf_config["num_shared_experts"] == 1
+        assert hf_config["router_scaling_factor"] == 2.5
+        assert hf_config["moe_intermediate_size"] == 32