From c08dcf4161f151c90482f7316fe866bec118312b Mon Sep 17 00:00:00 2001 From: Kamran Jafari Date: Fri, 9 Oct 2026 12:27:14 -0700 Subject: [PATCH 1/3] refactor(model): build Qwen2 from a builder-backed model config Move Qwen2/Qwen2.5 onto the ModelConfig/ModelBuilder path, following Muse Glimmer and Qwen3: - Qwen2Bridge maps Hugging Face configs in hf_config_to_model_config_kwargs(); the deprecated provider path reuses the same kwargs, so provider_bridge() is no longer overridden. - USE_MODEL_CONFIG_FOR_CONVERSION = True: checkpoint import, the conversion scripts and the compare tool build Qwen2 through its ModelBuilder. - All 30 Qwen2/2.5 H100 recipes use get_model_config(). - Add Qwen2 bridge unit tests (none existed) and a strict builder-config recipe test. Signed-off-by: Kamran Jafari Co-Authored-By: Claude Opus 5.5 --- .../bridge/models/qwen/qwen2_bridge.py | 39 +++-- .../bridge/recipes/qwen/h100/qwen2.py | 60 +++---- .../models/qwen/test_qwen2_bridge.py | 158 ++++++++++++++++++ .../recipes/qwen/test_qwen2_recipes.py | 50 +++++- 4 files changed, 263 insertions(+), 44 deletions(-) create mode 100644 tests/unit_tests/models/qwen/test_qwen2_bridge.py diff --git a/src/megatron/bridge/models/qwen/qwen2_bridge.py b/src/megatron/bridge/models/qwen/qwen2_bridge.py index f62382e007..1d070917e2 100644 --- a/src/megatron/bridge/models/qwen/qwen2_bridge.py +++ b/src/megatron/bridge/models/qwen/qwen2_bridge.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any + import torch from megatron.core.models.gpt.gpt_model import GPTModel -from transformers import Qwen2ForCausalLM +from transformers import PretrainedConfig, Qwen2ForCausalLM from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge @@ -35,26 +37,37 @@ class Qwen2Bridge(MegatronModelBridge): configuration translation. Qwen2 inherits CONFIG_MAPPING from MegatronModelBridge base class. - Model-specific settings are applied in provider_bridge. + Model-specific settings are applied in hf_config_to_model_config_kwargs. Example: >>> from megatron.bridge import AutoBridge >>> bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen2-7B") - >>> provider = bridge.to_megatron_provider() + >>> model_config = bridge.get_model_config() """ - def provider_bridge(self, hf_pretrained): - """Convert HuggingFace Qwen2 config to GPTModelProvider.""" - provider = super().provider_bridge(hf_pretrained) + USE_MODEL_CONFIG_FOR_CONVERSION = True - provider.normalization = "RMSNorm" - provider.gated_linear_unit = True - provider.add_bias_linear = False - provider.add_qkv_bias = True - provider.hidden_dropout = 0.0 - provider.autocast_dtype = torch.bfloat16 + def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]: + """Convert a Hugging Face Qwen2 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, + add_bias_linear=False, + add_qkv_bias=True, + hidden_dropout=0.0, + autocast_dtype=torch.bfloat16, + masked_softmax_fusion=True, + rope_scaling=False, + rope_scaling_factor=1.0, + ) + # The shared mapping already selects "yarn" for YaRN-scaled checkpoints. + 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.""" + return self.hf_config_to_model_config_kwargs(hf_config) def mapping_registry(self) -> MegatronMappingRegistry: # Return MegatronMappingRegistry containing parameter mappings from Megatron to HF format diff --git a/src/megatron/bridge/recipes/qwen/h100/qwen2.py b/src/megatron/bridge/recipes/qwen/h100/qwen2.py index c3943780f5..7c2619363a 100644 --- a/src/megatron/bridge/recipes/qwen/h100/qwen2.py +++ b/src/megatron/bridge/recipes/qwen/h100/qwen2.py @@ -37,7 +37,7 @@ def qwen2_500m_pretrain_1gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config (--tensor-model-parallel-size, --pipeline-model-parallel-size, etc.) - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-0.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-0.5B").get_model_config() # Tokenizer # Qwen2 uses NullTokenizer by default for pretraining @@ -128,7 +128,7 @@ def qwen2_1p5b_pretrain_1gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config (--tensor-model-parallel-size, --pipeline-model-parallel-size, etc.) - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-1.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-1.5B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_type = "NullTokenizer" @@ -217,7 +217,7 @@ def qwen2_7b_pretrain_2gpu_h100_bf16_config() -> ConfigContainer: """ cfg = _pretrain_common() - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-7B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-7B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_type = "NullTokenizer" @@ -307,7 +307,7 @@ def qwen2_72b_pretrain_32gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-72B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-72B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_type = "NullTokenizer" @@ -402,7 +402,7 @@ def qwen25_500m_pretrain_1gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config (--tensor-model-parallel-size, --pipeline-model-parallel-size, etc.) - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-0.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-0.5B").get_model_config() # Tokenizer (--tokenizer-model) cfg.tokenizer.tokenizer_type = "NullTokenizer" @@ -487,7 +487,7 @@ def qwen25_1p5b_pretrain_1gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-1.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-1.5B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_type = "NullTokenizer" @@ -572,7 +572,7 @@ def qwen25_7b_pretrain_2gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-7B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-7B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_type = "NullTokenizer" @@ -657,7 +657,7 @@ def qwen25_14b_pretrain_4gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-14B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-14B").get_model_config() # Tokenizer (--tokenizer-model) cfg.tokenizer.tokenizer_type = "NullTokenizer" @@ -747,7 +747,7 @@ def qwen25_32b_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-32B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-32B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_type = "NullTokenizer" @@ -832,7 +832,7 @@ def qwen25_72b_pretrain_32gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-72B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-72B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_type = "NullTokenizer" @@ -927,7 +927,7 @@ def qwen2_500m_sft_1gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-0.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-0.5B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2-0.5B" @@ -1029,7 +1029,7 @@ def qwen2_1p5b_sft_1gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-1.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-1.5B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2-1.5B" @@ -1131,7 +1131,7 @@ def qwen2_7b_sft_2gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-7B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-7B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2-7B" @@ -1233,7 +1233,7 @@ def qwen2_72b_sft_32gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-72B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-72B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2-72B" @@ -1340,7 +1340,7 @@ def qwen25_500m_sft_1gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-0.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-0.5B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-0.5B" @@ -1442,7 +1442,7 @@ def qwen25_1p5b_sft_1gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-1.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-1.5B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-1.5B" @@ -1544,7 +1544,7 @@ def qwen25_7b_sft_2gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-7B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-7B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-7B" @@ -1646,7 +1646,7 @@ def qwen25_14b_sft_4gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-14B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-14B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-14B" @@ -1748,7 +1748,7 @@ def qwen25_32b_sft_16gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-32B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-32B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-32B" @@ -1850,7 +1850,7 @@ def qwen25_72b_sft_32gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-72B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-72B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-72B" @@ -1960,7 +1960,7 @@ def qwen2_500m_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> C cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-0.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-0.5B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2-0.5B" @@ -2067,7 +2067,7 @@ def qwen2_1p5b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> C cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-1.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-1.5B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2-1.5B" @@ -2174,7 +2174,7 @@ def qwen2_7b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Con cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-7B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-7B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2-7B" @@ -2281,7 +2281,7 @@ def qwen2_72b_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Co cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-72B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2-72B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2-72B" @@ -2393,7 +2393,7 @@ def qwen25_500m_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-0.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-0.5B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-0.5B" @@ -2500,7 +2500,7 @@ def qwen25_1p5b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-1.5B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-1.5B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-1.5B" @@ -2607,7 +2607,7 @@ def qwen25_7b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Co cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-7B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-7B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-7B" @@ -2714,7 +2714,7 @@ def qwen25_14b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> C cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-14B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-14B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-14B" @@ -2821,7 +2821,7 @@ def qwen25_32b_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> C cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-32B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-32B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-32B" @@ -2928,7 +2928,7 @@ def qwen25_72b_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> C cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-72B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-72B").get_model_config() # Tokenizer cfg.tokenizer.tokenizer_model = "Qwen/Qwen2.5-72B" diff --git a/tests/unit_tests/models/qwen/test_qwen2_bridge.py b/tests/unit_tests/models/qwen/test_qwen2_bridge.py new file mode 100644 index 0000000000..f164ac9724 --- /dev/null +++ b/tests/unit_tests/models/qwen/test_qwen2_bridge.py @@ -0,0 +1,158 @@ +# 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.nn.functional as F +from transformers import Qwen2Config + +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.transformers_compat import rope_theta_from_hf +from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig +from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.qwen.qwen2_bridge import Qwen2Bridge + + +def _qwen2_config(**overrides) -> Qwen2Config: + kwargs = dict( + architectures=["Qwen2ForCausalLM"], + hidden_size=64, + intermediate_size=128, + max_position_embeddings=32768, + num_attention_heads=8, + num_hidden_layers=2, + num_key_value_heads=4, + rms_norm_eps=1e-6, + rope_theta=1000000.0, + tie_word_embeddings=False, + torch_dtype="bfloat16", + vocab_size=128, + ) + kwargs.update(overrides) + return Qwen2Config(**kwargs) + + +class TestQwen2Bridge: + def test_bridge_registration(self): + assert issubclass(Qwen2Bridge, MegatronModelBridge) + + def test_provider_bridge_is_inherited_compatibility_only(self): + """Qwen2 should not maintain a separate provider construction path.""" + assert "provider_bridge" not in Qwen2Bridge.__dict__ + + def test_conversion_uses_builder_config(self): + """Checkpoint conversion constructs Qwen2 through its ModelBuilder.""" + assert Qwen2Bridge.USE_MODEL_CONFIG_FOR_CONVERSION is True + + def test_hf_config_to_model_config_uses_direct_mapping(self): + """The builder config path must not route through the legacy provider.""" + config = _qwen2_config(tie_word_embeddings=True) + bridge = Qwen2Bridge() + + 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 path used"), + ), + ): + result = bridge.hf_config_to_model_config(config) + + assert isinstance(result, BridgeGPTModelConfig) + assert result.num_layers == config.num_hidden_layers + assert result.hidden_size == config.hidden_size + assert result.ffn_hidden_size == config.intermediate_size + assert result.num_attention_heads == config.num_attention_heads + assert result.num_query_groups == config.num_key_value_heads + assert result.seq_length == config.max_position_embeddings + assert result.rotary_base == rope_theta_from_hf(config) + assert result.vocab_size == config.vocab_size + assert result.share_embeddings_and_output_weights is True + assert result.layernorm_epsilon == config.rms_norm_eps + assert result.activation_func is F.silu + assert result.normalization == "RMSNorm" + assert result.gated_linear_unit is True + assert result.add_bias_linear is False + assert result.add_qkv_bias is True + assert result.qk_layernorm is False + assert result.hidden_dropout == 0.0 + assert result.position_embedding_type == "rope" + + @pytest.mark.parametrize( + ("tie_word_embeddings", "torch_dtype", "rope_scaling"), + [ + (True, "bfloat16", None), + (False, "float16", None), + ( + False, + "bfloat16", + { + "rope_type": "yarn", + "factor": 4.0, + "original_max_position_embeddings": 32768, + "beta_fast": 32.0, + "beta_slow": 1.0, + }, + ), + ], + ids=["tied-bf16", "untied-fp16", "yarn"], + ) + def test_model_config_matches_provider_runtime_config(self, tie_word_embeddings, torch_dtype, rope_scaling): + """Builder and provider paths agree on every comparable runtime field.""" + config = _qwen2_config( + tie_word_embeddings=tie_word_embeddings, torch_dtype=torch_dtype, rope_scaling=rope_scaling + ) + mock_pretrained = Mock(spec=PreTrainedCausalLM) + mock_pretrained.config = config + bridge = Qwen2Bridge() + + mapped_kwargs = bridge.hf_config_to_model_config_kwargs(config) + 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)) + comparable_fields = sorted((provider_fields & model_config_fields) - {"transformer_layer_spec"}) + + assert set(mapped_kwargs) <= set(comparable_fields) + assert len(comparable_fields) > len(mapped_kwargs) + for field_name in comparable_fields: + assert getattr(model_config, field_name) == getattr(provider, field_name), field_name + + # The provider stores its model-construction callable directly; the new + # config intentionally delegates layer-spec selection to its builder. + assert model_config.transformer_layer_spec is None + assert callable(provider.transformer_layer_spec) + + def test_model_config_maps_yarn_to_transformer_config(self): + """YaRN checkpoints keep YaRN and place its parameters where Megatron Core reads them.""" + config = _qwen2_config( + max_position_embeddings=131072, + rope_scaling={"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 32768}, + ) + + result = Qwen2Bridge().hf_config_to_model_config(config) + + assert result.position_embedding_type == "yarn" + assert result.transformer.yarn_rotary_scaling_factor == 4.0 + assert result.transformer.yarn_original_max_position_embeddings == 32768 + + hf_config = Qwen2Bridge.megatron_to_hf_config(result) + assert hf_config["rope_scaling"]["rope_type"] == "yarn" + assert hf_config["rope_scaling"]["factor"] == 4.0 diff --git a/tests/unit_tests/recipes/qwen/test_qwen2_recipes.py b/tests/unit_tests/recipes/qwen/test_qwen2_recipes.py index e3457bbd17..79226eb4e1 100644 --- a/tests/unit_tests/recipes/qwen/test_qwen2_recipes.py +++ b/tests/unit_tests/recipes/qwen/test_qwen2_recipes.py @@ -16,7 +16,10 @@ from typing import Callable import pytest +from transformers import Qwen2Config +from megatron.bridge import AutoBridge +from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global @@ -79,14 +82,43 @@ class _FakeBridge: def __init__(self): pass - def to_megatron_provider(self, load_weights: bool = False): + def get_model_config(self): return _FakeModelCfg() + def to_megatron_provider(self, load_weights: bool = False): + raise AssertionError("Qwen2 recipes must use get_model_config(), not the legacy provider API") + @staticmethod def from_hf_pretrained(hf_path: str, **kwargs): return _FakeBridge() +class _BuilderOnlyBridge: + """Return a real strict ModelConfig while rejecting the provider API.""" + + @staticmethod + def from_hf_pretrained(hf_path: str, **kwargs) -> "_BuilderOnlyBridge": + return _BuilderOnlyBridge() + + def get_model_config(self) -> BridgeGPTModelConfig: + config = Qwen2Config( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=80, + num_attention_heads=8, + num_key_value_heads=4, + max_position_embeddings=32768, + vocab_size=151936, + ) + config.architectures = ["Qwen2ForCausalLM"] + model_config = AutoBridge.from_hf_config(config).get_model_config() + assert isinstance(model_config, BridgeGPTModelConfig) + return model_config + + def to_megatron_provider(self, load_weights: bool = False): + raise AssertionError("Qwen2 recipes must not call the legacy provider API") + + def _assert_basic_config(cfg): from megatron.bridge.training.config import ConfigContainer @@ -450,3 +482,19 @@ def test_qwen25_72b_lora_defaults(monkeypatch: pytest.MonkeyPatch): # For LoRA, 72B should use TP=8, PP=1 assert cfg.model.tensor_model_parallel_size == 8 assert cfg.model.pipeline_model_parallel_size == 1 + + +@pytest.mark.parametrize("recipe_func", _QWEN2_RECIPE_FUNCS, ids=lambda func: func.__name__) +def test_each_qwen2_recipe_uses_strict_builder_config(recipe_func: Callable, monkeypatch: pytest.MonkeyPatch) -> None: + """Every Qwen2/2.5 recipe should configure a ModelConfig without provider fallback.""" + patch_recipe_module_global(monkeypatch, recipe_func, "AutoBridge", _BuilderOnlyBridge) + + if "peft" in recipe_func.__name__.lower(): + cfg = recipe_func(peft_scheme="lora") + else: + cfg = recipe_func() + + assert isinstance(cfg.model, BridgeGPTModelConfig) + assert cfg.model.normalization == "RMSNorm" + assert cfg.model.add_qkv_bias is True + assert cfg.model.position_embedding_type == "rope" From 098c1bf0f56e0eabaa6b256414b0bd2a2b802862 Mon Sep 17 00:00:00 2001 From: Kamran Jafari Date: Fri, 9 Oct 2026 12:27:15 -0700 Subject: [PATCH 2/3] 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 Co-Authored-By: Claude Opus 5.5 --- .../bridge/models/mimo/mimo_bridge.py | 17 +++++----- .../models/mimo/test_mimo_bridge.py | 34 +++++++++++++++++++ 2 files changed, 42 insertions(+), 9 deletions(-) diff --git a/src/megatron/bridge/models/mimo/mimo_bridge.py b/src/megatron/bridge/models/mimo/mimo_bridge.py index 9ec9f63317..8cb10f2de4 100644 --- a/src/megatron/bridge/models/mimo/mimo_bridge.py +++ b/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) diff --git a/tests/unit_tests/models/mimo/test_mimo_bridge.py b/tests/unit_tests/models/mimo/test_mimo_bridge.py index 95ccea8088..7692e133dd 100644 --- a/tests/unit_tests/models/mimo/test_mimo_bridge.py +++ b/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 From 96e9e4d79f08aa3ae11363250c44ab439cdb5214 Mon Sep 17 00:00:00 2001 From: Kamran Jafari Date: Wed, 14 Oct 2026 07:21:08 -0700 Subject: [PATCH 3/3] 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 Signed-off-by: Kamran Jafari --- src/megatron/bridge/models/qwen/qwen2_bridge.py | 2 ++ tests/unit_tests/models/mimo/test_mimo_bridge.py | 2 ++ tests/unit_tests/models/qwen/test_qwen2_bridge.py | 1 + tests/unit_tests/models/test_gpt_model_config.py | 6 ++++-- 4 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/models/qwen/qwen2_bridge.py b/src/megatron/bridge/models/qwen/qwen2_bridge.py index 1d070917e2..7bdf7cf533 100644 --- a/src/megatron/bridge/models/qwen/qwen2_bridge.py +++ b/src/megatron/bridge/models/qwen/qwen2_bridge.py @@ -25,6 +25,7 @@ GatedMLPMapping, QKVMapping, ) +from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig @MegatronModelBridge.register_bridge(source=Qwen2ForCausalLM, target=GPTModel, model_type="qwen2") @@ -45,6 +46,7 @@ class Qwen2Bridge(MegatronModelBridge): >>> model_config = bridge.get_model_config() """ + MODEL_CONFIG_CLASS = BridgeGPTModelConfig USE_MODEL_CONFIG_FOR_CONVERSION = True def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]: diff --git a/tests/unit_tests/models/mimo/test_mimo_bridge.py b/tests/unit_tests/models/mimo/test_mimo_bridge.py index 7692e133dd..d50627bec3 100644 --- a/tests/unit_tests/models/mimo/test_mimo_bridge.py +++ b/tests/unit_tests/models/mimo/test_mimo_bridge.py @@ -19,6 +19,7 @@ from transformers import GenerationConfig from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge, WeightConversionTask +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 from megatron.bridge.models.mimo.mimo_bridge import MimoBridge @@ -132,6 +133,7 @@ def test_provider_bridge_is_inherited_compatibility_only(self): assert "provider_bridge" not in MimoBridge.__dict__ def test_conversion_uses_builder_config(self): + assert MimoBridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig assert MimoBridge.USE_MODEL_CONFIG_FOR_CONVERSION is True def test_model_config_maps_mtp_config(self, mock_pretrained_mimo): diff --git a/tests/unit_tests/models/qwen/test_qwen2_bridge.py b/tests/unit_tests/models/qwen/test_qwen2_bridge.py index f164ac9724..fad46eca6a 100644 --- a/tests/unit_tests/models/qwen/test_qwen2_bridge.py +++ b/tests/unit_tests/models/qwen/test_qwen2_bridge.py @@ -55,6 +55,7 @@ def test_provider_bridge_is_inherited_compatibility_only(self): def test_conversion_uses_builder_config(self): """Checkpoint conversion constructs Qwen2 through its ModelBuilder.""" + assert Qwen2Bridge.MODEL_CONFIG_CLASS is BridgeGPTModelConfig assert Qwen2Bridge.USE_MODEL_CONFIG_FOR_CONVERSION is True def test_hf_config_to_model_config_uses_direct_mapping(self): diff --git a/tests/unit_tests/models/test_gpt_model_config.py b/tests/unit_tests/models/test_gpt_model_config.py index 406c1a8619..33de06377b 100644 --- a/tests/unit_tests/models/test_gpt_model_config.py +++ b/tests/unit_tests/models/test_gpt_model_config.py @@ -217,10 +217,12 @@ def test_unmigrated_bridge_keeps_provider_without_warning() -> None: vocab_size=256, ) - assert Qwen2Bridge.MODEL_CONFIG_CLASS is None + class UnmigratedQwen2Bridge(Qwen2Bridge): + MODEL_CONFIG_CLASS = None + with warnings.catch_warnings(): warnings.simplefilter("error", FutureWarning) - provider = Qwen2Bridge().provider_bridge(Mock(config=hf_config)) + provider = UnmigratedQwen2Bridge().provider_bridge(Mock(config=hf_config)) assert provider.position_embedding_type == "rope" assert provider.add_qkv_bias is True