megatron_patch/fixes/yarn_args/fix_yarn_args.patch

diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py
index 724ddfc0..2faffdac 100644
--- a/megatron/core/transformer/multi_latent_attention.py
+++ b/megatron/core/transformer/multi_latent_attention.py
@@ -89,7 +89,7 @@ class MultiLatentAttention(Attention):
                 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/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index b15b1abb..31fcadb9 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -1014,6 +1014,9 @@ class MLATransformerConfig(TransformerConfig):
     """Rotary scaling factor for the rotary embeddings, used by yarn."""
 
     max_position_embeddings: int = 4096
+    """NOT USED."""
+
+    original_max_position_embeddings: int = 4096
     """Maximum position embeddings for the original model, used by yarn."""
 
     beta_fast: float = 32
diff --git a/tools/checkpoint/loader_mixtral_hf.py b/tools/checkpoint/loader_mixtral_hf.py
index 131d6dc6..567aeb28 100644
--- a/tools/checkpoint/loader_mixtral_hf.py
+++ b/tools/checkpoint/loader_mixtral_hf.py
@@ -42,7 +42,7 @@ def load_args_from_checkpoint(args):
     args.tokenizer_type = "Llama2Tokenizer"
     args.disable_bias_linear = True
 
-    args.max_position_embeddings = mixtral_config.max_position_embeddings
+    args.max_position_embeddings = mixtral_config.original_max_position_embeddings
     args.hidden_size = mixtral_config.hidden_size
     args.num_attention_heads = mixtral_config.num_attention_heads
     args.num_layers = mixtral_config.num_hidden_layers