diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index be8fca56145..b097470030f 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/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".""" diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index f305a5a7668..ac90c69cde7 100644 --- a/megatron/training/arguments.py +++ b/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", diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index ffc9fe41e99..e889f0d559a 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/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}" diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index ec0e79d77ef..2837a3c4cc6 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/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()