From 1463315a7b78fbce6b2166d61972ae55c9cdc643 Mon Sep 17 00:00:00 2001 From: Guihong Li Date: Tue, 23 Jun 2026 15:08:52 -0700 Subject: [PATCH 1/3] Add formal YaRN RoPE config fields for hybrid models HybridModel already builds YarnRotaryEmbedding from getattr(self.config, "yarn_*") (#4244), but those attributes were never declared on TransformerConfig, so --position-embedding-type yarn raised AttributeError and YaRN could not be configured. Declaring the 7 yarn_* fields on TransformerConfig fixes this and also makes the --yarn-* CLI flags appear automatically: _add_network_size_args builds a TransformerConfig argument group via ArgumentGroupFactory(TransformerConfig), which derives one CLI flag per dataclass field. No manual argument plumbing is needed -- adding the flags by hand duplicates the auto-generated ones and raises argparse 'conflicting option string'. - transformer_config.py: declare the 7 yarn_* fields with safe defaults (yarn_rotary_scaling_factor default 1.0 is a no-op). - test_hybrid_model.py: drive the yarn test config via the constructor and add tests for the dataclass fields and config->embedding plumbing. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Guihong Li --- .../core/transformer/transformer_config.py | 24 +++++++++ tests/unit_tests/models/test_hybrid_model.py | 51 ++++++++++++++++--- 2 files changed, 67 insertions(+), 8 deletions(-) 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/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}" From 8f4b4a7d90c6b3042984e4aeb62d3b512cb1bcf3 Mon Sep 17 00:00:00 2001 From: Guihong Li Date: Wed, 24 Jun 2026 13:04:58 -0700 Subject: [PATCH 2/3] test: add yarn_* fields to Mamba MoE golden config The hybrid MoE golden-config drift test (test_hybrid_moe_model.py) pins the exact set of TransformerConfig fields. Adding the yarn_* fields trips its [ADDED ARGS] check, so register them in GOLDEN_CONFIG with their default values. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Guihong Li --- tests/unit_tests/models/test_hybrid_moe_model.py | 7 +++++++ 1 file changed, 7 insertions(+) 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() From f6337e43cd2b421d6663f2762ef89b24871fcc7d Mon Sep 17 00:00:00 2001 From: Guihong Li Date: Wed, 24 Jun 2026 15:51:30 -0700 Subject: [PATCH 3/3] args: exclude yarn_* fields from CLI auto-generation main adds the --yarn-* CLI flags manually in _add_network_size_args. Now that the yarn_* fields are declared on TransformerConfig, ArgumentGroupFactory would also auto-generate those flags, producing duplicate argparse option strings and breaking every test that builds the parser. Exclude the yarn_* fields from auto-generation so the manual flags remain the single source; the fields still provide config storage + defaults consumed by gpt_model/hybrid_model. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Guihong Li --- megatron/training/arguments.py | 10 ++++++++++ 1 file changed, 10 insertions(+) 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",