fix(model): apply Hugging Face YaRN defaults for omitted rope_scaling keys (#6367)

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/conversion/model_bridge.py

@@ -627,6 +627,14 @@ def postprocess_hf_export_weights(self, path: Path) -> None:
         ("mscale_all_dim", "yarn_mscale_all_dim"),
     ]
 
+    # Hugging Face YaRN defaults (transformers' _compute_yarn_parameters) for keys a
+    # checkpoint omits: (megatron_param, default)
+    YARN_HF_DEFAULTS = [
+        ("yarn_beta_fast", 32.0),
+        ("yarn_beta_slow", 1.0),
+        ("yarn_correction_range_round_to_int", True),
+    ]
+
     # MLA rope scaling field mapping: (hf_rope_scaling_key, megatron_mla_param)
     # These are applied for MLA models (DeepSeek, Kimi, etc.) which use MLATransformerConfig
     # Uses direct field names without yarn_ prefix (e.g., mscale, rotary_scaling_factor)
@@ -736,6 +744,10 @@ def _hf_config_to_megatron_kwargs(self, hf_config: PretrainedConfig) -> dict[str
                         provider_kwargs[megatron_key] = value
                 if "truncate" in rope_scaling:
                     provider_kwargs["yarn_correction_range_round_to_int"] = rope_scaling["truncate"]
+                # Apply Hugging Face YaRN defaults for omitted keys. Megatron Core passes these
+                # fields to YarnRotaryEmbedding explicitly, so None would override its defaults.
+                for megatron_key, default in self.YARN_HF_DEFAULTS:
+                    provider_kwargs.setdefault(megatron_key, default)
         elif is_mla_provider:
             if rope_type not in (None, "default"):
                 logger.warning(

src/megatron/bridge/models/qwen/qwen3_bridge.py

@@ -61,14 +61,8 @@ def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[
             rope_scaling=False,
             rope_scaling_factor=1.0,
         )
-        # The shared mapping already selects "yarn" for YaRN-scaled checkpoints.
+        # The shared mapping already selects "yarn" and its Hugging Face defaults 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
 
     def hf_config_to_provider_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:

tests/unit_tests/models/test_model_bridge.py

@@ -18,7 +18,7 @@
 
 import pytest
 import torch
-from transformers import PretrainedConfig
+from transformers import PretrainedConfig, Qwen2Config
 
 from megatron.bridge.models.conversion import model_bridge as model_bridge_module
 from megatron.bridge.models.conversion import modelopt_utils
@@ -1691,3 +1691,56 @@ def __init__(self, megatron_param, hf_param):
     assert len(tasks) == 1
     assert tasks[0].global_param_name == "decoder.weight"
     assert tasks[0].megatron_module is None
+
+
+def _yarn_qwen2_config(rope_scaling):
+    return Qwen2Config(
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        max_position_embeddings=16384,
+        vocab_size=128,
+        rope_scaling=rope_scaling,
+    )
+
+
+def test_yarn_mapping_fills_hugging_face_defaults_for_omitted_keys():
+    """Omitted YaRN keys must not reach Megatron Core as None."""
+    config = _yarn_qwen2_config({"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 4096})
+
+    kwargs = DummyBridge().hf_config_to_provider_kwargs(config)
+
+    assert kwargs["position_embedding_type"] == "yarn"
+    assert kwargs["yarn_rotary_scaling_factor"] == 4.0
+    assert kwargs["yarn_original_max_position_embeddings"] == 4096
+    assert kwargs["yarn_beta_fast"] == 32.0
+    assert kwargs["yarn_beta_slow"] == 1.0
+    assert kwargs["yarn_correction_range_round_to_int"] is True
+
+
+def test_yarn_mapping_preserves_explicit_keys():
+    config = _yarn_qwen2_config(
+        {
+            "rope_type": "yarn",
+            "factor": 4.0,
+            "original_max_position_embeddings": 4096,
+            "beta_fast": 16.0,
+            "beta_slow": 2.0,
+            "truncate": False,
+        }
+    )
+
+    kwargs = DummyBridge().hf_config_to_provider_kwargs(config)
+
+    assert kwargs["yarn_beta_fast"] == 16.0
+    assert kwargs["yarn_beta_slow"] == 2.0
+    assert kwargs["yarn_correction_range_round_to_int"] is False
+
+
+def test_non_yarn_mapping_adds_no_yarn_defaults():
+    kwargs = DummyBridge().hf_config_to_provider_kwargs(_yarn_qwen2_config(None))
+
+    assert "position_embedding_type" not in kwargs
+    assert not any(key.startswith("yarn_") for key in kwargs)