refactor(model): add Qwen3 builder-backed model config with YaRN support (#6365)

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

Author: kamran-nvidia

docs/bridge-guide.md

@@ -106,12 +106,13 @@ The provider pattern is especially useful when you need to:
 - Configure advanced features like MoE, activation recomputation, or mixed precision
 - Set up distributed training parameters
 
-### Using Builder-backed Llama Configs
+### Using Builder-backed Model Configs
 
-Llama also supports the builder-backed configuration path. This keeps model
-configuration as serializable data and leaves construction to Megatron Core's
-`GPTModelBuilder`. The provider API remains available for compatibility while
-other model families migrate. Calling the legacy provider API for a
+Model families are migrating to the builder-backed configuration path. It keeps
+model configuration as serializable data and leaves construction to the
+family's Megatron Core `ModelBuilder`, such as `GPTModelBuilder` or
+`HybridModelBuilder`. The provider API remains available for compatibility
+while other model families migrate. Calling the legacy provider API for a
 builder-backed family emits a deprecation warning.
 
 ```python
@@ -134,9 +135,10 @@ Use `load_weights=False` for random initialization. A bridge created with
 `from_hf_config()` has no weights, so it requires `load_weights=False` or an
 explicit `hf_path`.
 
-Llama training recipes store the result of `get_model_config()` in
-`ConfigContainer.model`. The training setup recognizes `ModelConfig` and calls
-its `ModelBuilder` directly; it does not create a legacy model provider.
+Training recipes for builder-backed families store the result of
+`get_model_config()` in `ConfigContainer.model`. The training setup recognizes
+`ModelConfig` and calls its `ModelBuilder` directly; it does not create a
+legacy model provider.
 
 ## Check Supported Models
 

src/megatron/bridge/models/gpt_provider.py

@@ -148,15 +148,7 @@ class GPTModelProvider(TransformerConfig, ModelProviderMixin[MCoreGPTModel]):
     rotary_scaling_factor: Optional[float] = None
     seq_len_interpolation_factor: Optional[float] = None
 
-    # YARN (Yet Another RoPE extensioN) position embedding parameters
-    # Used when position_embedding_type == "yarn"
-    yarn_rotary_scaling_factor: Optional[float] = None
-    yarn_original_max_position_embeddings: Optional[int] = None
-    yarn_beta_fast: Optional[float] = None
-    yarn_beta_slow: Optional[float] = None
-    yarn_mscale: Optional[float] = None
-    yarn_mscale_all_dim: Optional[float] = None
-    yarn_correction_range_round_to_int: Optional[bool] = None
+    # YaRN parameters (yarn_*) are declared on the base TransformerConfig.
 
     seq_length: int = 1024
     attention_softmax_in_fp32: bool = False
@@ -405,37 +397,3 @@ def mtp_block_spec(config: "GPTModelProvider", vp_stage: Optional[int] = None) -
         return get_gpt_mtp_block_spec(config, spec, use_transformer_engine=True, vp_stage=vp_stage)
     else:
         return None
-
-
-def _patch_yarn_concentration_factor():
-    """Patch MCore _yarn_get_concentration_factor_from_config for None handling.
-
-    GPTModelProvider defines yarn_rotary_scaling_factor as Optional[float] = None,
-    but MCore uses hasattr() which returns True for dataclass fields set to None.
-    This causes a crash for non-YARN models. Use getattr + is not None instead.
-
-    TODO: Remove once upstream MCore merges the fix.
-    """
-    try:
-        import megatron.core.models.common.embeddings.yarn_rotary_pos_embedding as _yarn_mod
-        import megatron.core.transformer.attention as _attn_mod
-
-        _get_factor = _yarn_mod._yarn_get_concentration_factor
-
-        def _fixed_from_config(config):
-            yarn_scaling = getattr(config, "yarn_rotary_scaling_factor", None)
-            if yarn_scaling is not None:
-                return _get_factor(
-                    yarn_scaling,
-                    getattr(config, "yarn_mscale", None),
-                    getattr(config, "yarn_mscale_all_dim", None),
-                )
-            return 1.0
-
-        _yarn_mod._yarn_get_concentration_factor_from_config = _fixed_from_config
-        _attn_mod._yarn_get_concentration_factor_from_config = _fixed_from_config
-    except ImportError:
-        pass
-
-
-_patch_yarn_concentration_factor()

src/megatron/bridge/models/qwen/qwen3_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 Qwen3ForCausalLM
+from transformers import PretrainedConfig, Qwen3ForCausalLM
 
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
 from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
@@ -37,22 +39,39 @@ class Qwen3Bridge(MegatronModelBridge):
     Example:
         >>> from megatron.bridge import AutoBridge
         >>> bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen3-1.7B")
-        >>> provider = bridge.to_megatron_provider()
+        >>> model_config = bridge.get_model_config()
     """
 
-    def provider_bridge(self, hf_pretrained):
-        """Convert HuggingFace Qwen3 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 = False  # Qwen3 does NOT have QKV bias (unlike Qwen2)
-        provider.hidden_dropout = 0.0
-        provider.qk_layernorm = True  # Qwen3 uses QK layernorm
-        provider.autocast_dtype = torch.bfloat16
+    def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Convert a Hugging Face Qwen3 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=False,  # Qwen3 does NOT have QKV bias (unlike Qwen2)
+            hidden_dropout=0.0,
+            qk_layernorm=True,  # Qwen3 uses QK layernorm
+            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")
+        if config_kwargs["position_embedding_type"] == "yarn":
+            # Qwen3 YaRN configs usually omit these keys; apply the Hugging Face defaults
+            # instead of passing None to Megatron Core.
+            config_kwargs.setdefault("yarn_beta_fast", 32.0)
+            config_kwargs.setdefault("yarn_beta_slow", 1.0)
+            config_kwargs.setdefault("yarn_correction_range_round_to_int", True)
+        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 the MegatronMappingRegistry for Qwen3 parameter conversion.

src/megatron/bridge/models/transformer_config.py

@@ -145,6 +145,41 @@ def _uses_legacy_full_iteration(value: object) -> bool:
             setattr(config, padding_field, True)
 
 
+def _patch_yarn_concentration_factor() -> None:
+    """Patch MCore _yarn_get_concentration_factor_from_config for None handling.
+
+    TransformerConfig declares yarn_rotary_scaling_factor as ``float | None = None``,
+    but MCore uses hasattr(), which returns True for dataclass fields set to None.
+    This causes a crash for non-YaRN models. Use getattr + is not None instead.
+
+    TODO: Remove once upstream MCore merges the fix.
+    """
+    try:
+        import megatron.core.models.common.embeddings.yarn_rotary_pos_embedding as _yarn_mod
+        import megatron.core.transformer.attention as _attn_mod
+
+        _get_factor = _yarn_mod._yarn_get_concentration_factor
+
+        def _fixed_from_config(config):
+            yarn_scaling = getattr(config, "yarn_rotary_scaling_factor", None)
+            if yarn_scaling is not None:
+                return _get_factor(
+                    yarn_scaling,
+                    getattr(config, "yarn_mscale", None),
+                    getattr(config, "yarn_mscale_all_dim", None),
+                )
+            return 1.0
+
+        _yarn_mod._yarn_get_concentration_factor_from_config = _fixed_from_config
+        _attn_mod._yarn_get_concentration_factor_from_config = _fixed_from_config
+    except ImportError:
+        pass
+
+
+# Applied at import so every config that declares the YaRN fields below is covered.
+_patch_yarn_concentration_factor()
+
+
 @dataclass
 class TransformerConfig(MCoreTransformerConfig):
     """Megatron Core TransformerConfig with deferred post-init.
@@ -175,6 +210,17 @@ class TransformerConfig(MCoreTransformerConfig):
     gtp_weight_remat_size: int = field(init=False, default=1)
     expert_gtp_weight_remat_size: int = field(init=False, default=1)
 
+    # YaRN (Yet another RoPE extensioN) parameters, used when position_embedding_type == "yarn".
+    # Megatron Core models read these from their transformer config, so builder-backed
+    # configs declare them here rather than on the outer model config.
+    yarn_rotary_scaling_factor: float | None = None
+    yarn_original_max_position_embeddings: int | None = None
+    yarn_beta_fast: float | None = None
+    yarn_beta_slow: float | None = None
+    yarn_mscale: float | None = None
+    yarn_mscale_all_dim: float | None = None
+    yarn_correction_range_round_to_int: bool | None = None
+
     def __post_init__(self) -> None:
         """Skip MCore post_init during initial construction.
 

src/megatron/bridge/recipes/qwen/h100/qwen3.py

@@ -37,7 +37,7 @@ def qwen3_600m_pretrain_1gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _pretrain_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-0.6B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-0.6B").get_model_config()
 
     # Tokenizer (--tokenizer-model)
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-0.6B"
@@ -122,7 +122,7 @@ def qwen3_1p7b_pretrain_1gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _pretrain_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-1.7B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-1.7B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-1.7B"
@@ -207,7 +207,7 @@ def qwen3_4b_pretrain_2gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _pretrain_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-4B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-4B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-4B"
@@ -292,9 +292,7 @@ def qwen3_8b_pretrain_4gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _pretrain_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained(
-        _QWEN3_8B_MODEL_ID, revision=_QWEN3_8B_MODEL_REVISION
-    ).to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained(_QWEN3_8B_MODEL_ID, revision=_QWEN3_8B_MODEL_REVISION).get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = _QWEN3_8B_MODEL_ID
@@ -435,7 +433,7 @@ def qwen3_14b_pretrain_8gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _pretrain_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-14B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-14B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-14B"
@@ -519,7 +517,7 @@ def qwen3_32b_pretrain_16gpu_h100_bf16_config() -> ConfigContainer:
     """
     cfg = _pretrain_common()
 
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-32B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-32B").get_model_config()
 
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-32B"
 
@@ -597,7 +595,7 @@ def qwen3_600m_sft_1gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _sft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-0.6B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-0.6B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-0.6B"
@@ -741,7 +739,7 @@ def qwen3_600m_sft_8gpu_h100_bf16_yarn_128k_config() -> ConfigContainer:
     hf_path = "Qwen/Qwen3-0.6B"
 
     # Model and tokenizer use the same HF checkpoint so the chat template matches the dataset.
-    cfg.model = AutoBridge.from_hf_pretrained(hf_path).to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained(hf_path).get_model_config()
     cfg.tokenizer.tokenizer_model = hf_path
 
     # 128K context parallelism.
@@ -817,7 +815,7 @@ def qwen3_1p7b_sft_1gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _sft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-1.7B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-1.7B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-1.7B"
@@ -909,7 +907,7 @@ def qwen3_4b_sft_2gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _sft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-4B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-4B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-4B"
@@ -1001,9 +999,7 @@ def qwen3_8b_sft_4gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _sft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained(
-        _QWEN3_8B_MODEL_ID, revision=_QWEN3_8B_MODEL_REVISION
-    ).to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained(_QWEN3_8B_MODEL_ID, revision=_QWEN3_8B_MODEL_REVISION).get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = _QWEN3_8B_MODEL_ID
@@ -1149,7 +1145,7 @@ def qwen3_14b_sft_8gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _sft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-14B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-14B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-14B"
@@ -1242,7 +1238,7 @@ def qwen3_32b_sft_16gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _sft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-32B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-32B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-32B"
@@ -1344,7 +1340,7 @@ def qwen3_600m_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> C
     cfg = _peft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-0.6B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-0.6B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-0.6B"
@@ -1442,7 +1438,7 @@ def qwen3_1p7b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> C
     cfg = _peft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-1.7B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-1.7B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-1.7B"
@@ -1540,7 +1536,7 @@ def qwen3_4b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Con
     cfg = _peft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-4B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-4B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-4B"
@@ -1638,9 +1634,7 @@ def qwen3_8b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Con
     cfg = _peft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained(
-        _QWEN3_8B_MODEL_ID, revision=_QWEN3_8B_MODEL_REVISION
-    ).to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained(_QWEN3_8B_MODEL_ID, revision=_QWEN3_8B_MODEL_REVISION).get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = _QWEN3_8B_MODEL_ID
@@ -1765,7 +1759,7 @@ def qwen3_14b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Co
     cfg = _peft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-14B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-14B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-14B"
@@ -1864,7 +1858,7 @@ def qwen3_32b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Co
     cfg = _peft_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-32B").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-32B").get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-32B"