Skip to content

Commit 3d8a518

Browse files
kamran-nvidiaclaude
andcommitted
refactor(model): opt the migrated families into builder-backed model configs
After #6305, MODEL_CONFIG_CLASS defaults to None and each family opts in. Set MODEL_CONFIG_CLASS = BridgeGPTModelConfig on the bridges this PR migrates, matching Qwen3 and Llama, and assert it in their tests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Kamran Jafari <kjafarisadeg@nvidia.com>
1 parent 465dd2c commit 3d8a518

4 files changed

Lines changed: 7 additions & 0 deletions

File tree

‎src/megatron/bridge/models/hy_v3/hy_v3_bridge.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
GatedMLPMapping,
2828
QKVMapping,
2929
)
30+
from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
3031
from megatron.bridge.models.gpt_provider import GPTModelProvider
3132

3233

@@ -61,6 +62,7 @@ class HYV3Bridge(MegatronModelBridge):
6162
>>> model_config = bridge.get_model_config()
6263
"""
6364

65+
MODEL_CONFIG_CLASS = BridgeGPTModelConfig
6466
USE_MODEL_CONFIG_FOR_CONVERSION = True
6567

6668
def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:

‎src/megatron/bridge/models/minimax_m2/minimax_m2_bridge.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
MegatronParamMapping,
3030
QKVMapping,
3131
)
32+
from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
3233
from megatron.bridge.models.minimax_m2.minimax_m2_provider import minimax_m2_layer_spec
3334

3435

@@ -132,6 +133,7 @@ class MiniMaxM2Bridge(MegatronModelBridge):
132133
>>> model_config = bridge.get_model_config()
133134
"""
134135

136+
MODEL_CONFIG_CLASS = BridgeGPTModelConfig
135137
USE_MODEL_CONFIG_FOR_CONVERSION = True
136138

137139
def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:

‎tests/unit_tests/models/hy_v3/test_hy_v3_bridge.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ def test_provider_bridge_is_inherited_compatibility_only(self):
5858
assert "provider_bridge" not in HYV3Bridge.__dict__
5959

6060
def test_conversion_uses_builder_config(self):
61+
assert HYV3Bridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
6162
assert HYV3Bridge.USE_MODEL_CONFIG_FOR_CONVERSION is True
6263

6364
def test_hf_config_to_model_config_uses_direct_mapping(self):

‎tests/unit_tests/models/minimax_m2/test_minimax_m2_bridges.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from transformers import GenerationConfig
2424

2525
from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
26+
from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
2627
from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
2728
from megatron.bridge.models.minimax_m2.minimax_m2_bridge import (
2829
MiniMaxM2Bridge,
@@ -535,6 +536,7 @@ def test_provider_bridge_is_inherited_compatibility_only(self):
535536
assert "provider_bridge" not in MiniMaxM2Bridge.__dict__
536537

537538
def test_conversion_uses_builder_config(self):
539+
assert MiniMaxM2Bridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig
538540
assert MiniMaxM2Bridge.USE_MODEL_CONFIG_FOR_CONVERSION is True
539541

540542
def test_hf_config_to_model_config_uses_direct_mapping(self):

0 commit comments

Comments
 (0)