From 586b69bf9e1f8dc38f5a5ab40b59d13ef0826594 Mon Sep 17 00:00:00 2001 From: zhujian Date: Wed, 20 May 2026 12:02:16 +0800 Subject: [PATCH] fix: support MLA YaRN rope scaling arguments Add original_max_position_embeddings and expose beta_fast/beta_slow for MLA YaRN RoPE, use mscale_all_dim for softmax scaling, and avoid rejecting checkpoint loads when only max_position_embeddings changes. Local-port-of: f1a85bad9e7f05c4fa80c81181e44414e1ac8ab7 Upstream-note: Based on Megatron-LM MLA/YaRN RoPE behavior; exact upstream commit not identified. --- .../core/transformer/multi_latent_attention.py | 4 ++-- .../core/transformer/transformer_config.py | 14 +++++++++++--- aiak_megatron/megatron/training/arguments.py | 9 ++++++++- aiak_megatron/megatron/training/checkpointing.py | 2 +- 4 files changed, 22 insertions(+), 7 deletions(-) diff --git a/aiak_megatron/megatron/core/transformer/multi_latent_attention.py b/aiak_megatron/megatron/core/transformer/multi_latent_attention.py index aa3bfb75..7616f0f3 100644 --- a/aiak_megatron/megatron/core/transformer/multi_latent_attention.py +++ b/aiak_megatron/megatron/core/transformer/multi_latent_attention.py @@ -88,7 +88,7 @@ def __init__( self.key_hidden_size = self.q_head_dim self.val_hidden_size = self.config.v_head_dim - mscale = _yarn_get_mscale(self.config.rotary_scaling_factor, self.config.mscale) + mscale = _yarn_get_mscale(self.config.rotary_scaling_factor, self.config.mscale_all_dim) self.softmax_scale = mscale * mscale / math.sqrt(self.q_head_dim) if self.config.rope_type == "rope": @@ -102,7 +102,7 @@ def __init__( self.config.qk_pos_emb_head_dim, rotary_base=self.config.rotary_base, scaling_factor=self.config.rotary_scaling_factor, - original_max_position_embeddings=self.config.max_position_embeddings, + original_max_position_embeddings=self.config.original_max_position_embeddings, beta_fast=self.config.beta_fast, beta_slow=self.config.beta_slow, mscale=self.config.mscale, diff --git a/aiak_megatron/megatron/core/transformer/transformer_config.py b/aiak_megatron/megatron/core/transformer/transformer_config.py index e41c3eeb..90e97ea0 100644 --- a/aiak_megatron/megatron/core/transformer/transformer_config.py +++ b/aiak_megatron/megatron/core/transformer/transformer_config.py @@ -1078,7 +1078,10 @@ class MLATransformerConfig(TransformerConfig): """Rotary scaling factor for the rotary embeddings, used by yarn.""" max_position_embeddings: int = 4096 - """Maximum position embeddings for the original model, used by yarn.""" + """This arg is not used, will be deprecated.""" + + original_max_position_embeddings: int = 4096 + """Original maximum position embeddings for the original model, used by yarn.""" beta_fast: float = 32 """Beta fast for YaRN RoPE, used by yarn.""" @@ -1086,8 +1089,13 @@ class MLATransformerConfig(TransformerConfig): beta_slow: float = 1 """Beta slow for YaRN RoPE, used by yarn.""" - mscale: float = 0.707 + mscale: float = 1.0 """Mscale for YaRN RoPE in Multi-Latent Attention, used by yarn.""" - mscale_all_dim: float = 0.707 + mscale_all_dim: float = 0.0 """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" + + def __post_init__(self): + super().__post_init__() + if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": + raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.") diff --git a/aiak_megatron/megatron/training/arguments.py b/aiak_megatron/megatron/training/arguments.py index 82b71343..6da1c0c6 100644 --- a/aiak_megatron/megatron/training/arguments.py +++ b/aiak_megatron/megatron/training/arguments.py @@ -2619,10 +2619,17 @@ def _add_mla_args(parser): help="Dimension of the head in the V projection.") group.add_argument('--rotary-scaling-factor', type=float, default=1.0, help="Rotary scaling factor for the rotary embeddings.") + group.add_argument('--beta-fast', type=float, default=32.0, + help="beta_fast for YaRN RoPE in multi-latent attention.") + group.add_argument('--beta-slow', type=float, default=1.0, + help="beta_slow for YaRN RoPE in multi-latent attention.") group.add_argument('--mscale', type=float, default=1.0, help="Mscale for YaRN RoPE in multi-latent attention.") - group.add_argument('--mscale-all-dim', type=float, default=1.0, + group.add_argument('--mscale-all-dim', type=float, default=0.0, help="Mscale all dimensions for YaRN RoPE in multi-latent attention.") + group.add_argument('--original-max-position-embeddings', type=int, default=None, + help='Original maximum number of position embeddings, used by yarn.' + 'This is the size of position embedding.') return parser diff --git a/aiak_megatron/megatron/training/checkpointing.py b/aiak_megatron/megatron/training/checkpointing.py index 2e273944..a832ca31 100644 --- a/aiak_megatron/megatron/training/checkpointing.py +++ b/aiak_megatron/megatron/training/checkpointing.py @@ -98,7 +98,7 @@ def _compare(arg_name, old_arg_name=None, default=None): _compare("num_attention_heads") _compare("add_position_embedding", default=False) if args.vocab_file: - _compare("max_position_embeddings") + # _compare("max_position_embeddings") _compare("make_vocab_size_divisible_by") if not args.use_dist_ckpt: _compare("padded_vocab_size")