Add formal YaRN RoPE config fields for hybrid models

megatron/core/transformer/transformer_config.py

@@ -226,6 +226,30 @@ class TransformerConfig(ModelParallelConfig):
     """True is rotate pairs of even and odd dimensions (RoFormer style), False is rotate pairs of
     first half and second half (LLaMa style). Default to False."""
 
+    yarn_rotary_scaling_factor: float = 1.0
+    """Scaling factor for YaRN RoPE, used when position_embedding_type == 'yarn'. Default 1.0
+    (no scaling)."""
+
+    yarn_original_max_position_embeddings: int = 4096
+    """Original (pre-extension) maximum sequence length the model was trained on, used by YaRN
+    RoPE when position_embedding_type == 'yarn'."""
+
+    yarn_beta_fast: float = 32.0
+    """Fast beta value for YaRN RoPE, used when position_embedding_type == 'yarn'."""
+
+    yarn_beta_slow: float = 1.0
+    """Slow beta value for YaRN RoPE, used when position_embedding_type == 'yarn'."""
+
+    yarn_mscale: float = 1.0
+    """Mscale value for YaRN RoPE, used when position_embedding_type == 'yarn'."""
+
+    yarn_mscale_all_dim: float = 0.0
+    """Mscale-all-dim value for YaRN RoPE, used when position_embedding_type == 'yarn'."""
+
+    yarn_correction_range_round_to_int: bool = True
+    """Whether to round the YaRN correction dimension range bounds to integers, used when
+    position_embedding_type == 'yarn'."""
+
     window_size: Optional[Tuple[int, int]] = None
     """If not None, then will use sliding window attention. The size of the window is specified by
     the numbers inside the tuple; -1 is special value meaning "infinite window size"."""

megatron/training/arguments.py

@@ -2022,6 +2022,16 @@ def _add_inference_args(parser):
 
 def _add_network_size_args(parser):
     exclude = [
+        # YaRN RoPE: CLI flags are added manually in _add_network_size_args
+        # (--yarn-*), so exclude these fields from auto-generation to avoid
+        # duplicate argparse option strings.
+        "yarn_rotary_scaling_factor",
+        "yarn_original_max_position_embeddings",
+        "yarn_beta_fast",
+        "yarn_beta_slow",
+        "yarn_mscale",
+        "yarn_mscale_all_dim",
+        "yarn_correction_range_round_to_int",
         # cannot provide callables over CLI
         "timers",
         "finalize_model_grads_func",

tests/unit_tests/models/test_hybrid_model.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
 
+import dataclasses
 import os
 from datetime import timedelta
 from itertools import accumulate
@@ -563,21 +564,26 @@ def test_dynamic_inference_padding_with_fp8(self):
 
 def _make_yarn_config(**kwargs):
     """Build a TransformerConfig with yarn positional embedding attributes."""
+    # Yarn-specific attributes are formal TransformerConfig fields and are passed directly to
+    # the constructor (overridable via kwargs).
+    yarn_defaults = dict(
+        yarn_rotary_scaling_factor=2.0,
+        yarn_original_max_position_embeddings=4,
+        yarn_beta_fast=32.0,
+        yarn_beta_slow=1.0,
+        yarn_mscale=1.0,
+        yarn_mscale_all_dim=0.0,
+        yarn_correction_range_round_to_int=True,
+    )
+    yarn_defaults.update({k: kwargs.pop(k) for k in list(kwargs) if k in yarn_defaults})
     cfg = TransformerConfig(
         num_layers=3,  # 1 Mamba layer, 1 attention layer, 1 MLP layer
         hidden_size=256,
         num_attention_heads=4,
         use_cpu_initialization=True,
+        **yarn_defaults,
         **kwargs,
     )
-    # Yarn-specific attributes are set dynamically on the config (not TransformerConfig fields).
-    cfg.yarn_rotary_scaling_factor = 2.0
-    cfg.yarn_original_max_position_embeddings = 4
-    cfg.yarn_beta_fast = 32.0
-    cfg.yarn_beta_slow = 1.0
-    cfg.yarn_mscale = 1.0
-    cfg.yarn_mscale_all_dim = 0.0
-    cfg.yarn_correction_range_round_to_int = True
     return cfg
 
 
@@ -608,6 +614,17 @@ def test_constructor(self):
         # YaRN creates a YarnRotaryEmbedding rather than a plain RotaryEmbedding.
         assert isinstance(self.model.rotary_pos_emb, YarnRotaryEmbedding)
 
+    def test_config_values_flow_into_embedding(self):
+        # The yarn_* config fields must be plumbed through to the YarnRotaryEmbedding instance.
+        emb = self.model.rotary_pos_emb
+        assert emb.scaling_factor == 2.0
+        assert emb.original_max_position_embeddings == 4
+        assert emb.beta_fast == 32.0
+        assert emb.beta_slow == 1.0
+        assert emb.mscale == 1.0
+        assert emb.mscale_all_dim == 0.0
+        assert emb.correction_range_round_to_int is True
+
     def test_forward(self):
         sequence_length = self.model.max_sequence_length
         micro_batch_size = 2
@@ -667,3 +684,21 @@ def test_inference(self):
                 # StaticInferenceContext always sets materialize_only_last_token_logits=True.
                 assert logits.shape[1] == 1
                 assert logits.shape[2] == self.model.vocab_size
+
+
+def test_yarn_config_fields_are_formal_dataclass_fields():
+    """The yarn_* attributes consumed by HybridModel must be real TransformerConfig fields
+    (with defaults), so `--position-embedding-type yarn` works without dynamic patching."""
+    field_defaults = {f.name: f.default for f in dataclasses.fields(TransformerConfig)}
+    expected = {
+        "yarn_rotary_scaling_factor": 1.0,
+        "yarn_original_max_position_embeddings": 4096,
+        "yarn_beta_fast": 32.0,
+        "yarn_beta_slow": 1.0,
+        "yarn_mscale": 1.0,
+        "yarn_mscale_all_dim": 0.0,
+        "yarn_correction_range_round_to_int": True,
+    }
+    for name, default in expected.items():
+        assert name in field_defaults, f"{name} is not a TransformerConfig field"
+        assert field_defaults[name] == default, f"unexpected default for {name}"

tests/unit_tests/models/test_hybrid_moe_model.py

@@ -319,6 +319,13 @@
     "use_transformer_engine_op_fuser": False,
     "moe_single_grouped_weight": False,
     "moe_single_grouped_bias": False,
+    "yarn_rotary_scaling_factor": 1.0,
+    "yarn_original_max_position_embeddings": 4096,
+    "yarn_beta_fast": 32.0,
+    "yarn_beta_slow": 1.0,
+    "yarn_mscale": 1.0,
+    "yarn_mscale_all_dim": 0.0,
+    "yarn_correction_range_round_to_int": True,
 }
 # Fields to ignore entirely (ephemeral, environment-specific, very large).
 SKIP_FIELDS = set()