[fix] Mistral-Large-3 accuracy regression on main (#57563)

Signed-off-by: jdebache <jdebache@nvidia.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
Julien Debache
2026-09-18 14:19:16 +00:00
committed by GitHub
co-authored by Harry Mellor
parent 76d517fd12
commit 2bdae2a5a8
8 changed files with 79 additions and 27 deletions
+52
View File
@@ -28,6 +28,7 @@ from vllm.transformers_utils.configs.glm5_next import (
Glm5NextTextConfig,
Glm5NextVisionConfig,
)
from vllm.transformers_utils.configs.mistral import adapt_config_dict
def test_patch_legacy_rope_type_preserves_nope_layers():
@@ -77,6 +78,57 @@ def test_patch_legacy_rope_type_normalizes_telechat3_yarn():
}
def test_mistral_yarn_apply_scale_false_disables_yarn_magnitude_scaling():
"""`yarn.apply_scale: false` must reach the DeepSeek-style attentions.
Transformers spells it `attention_factor = 1.0`, which DeepseekV2Attention
and its siblings read to select `deepseek_llama_scaling` over
`deepseek_yarn`; without it Mistral-Large-3 runs with a spurious
yarn_get_mscale(factor)^2 attention scaling.
"""
params = {
"dim": 7168,
"n_layers": 61,
"head_dim": 192,
"hidden_dim": 16384,
"n_heads": 128,
"n_kv_heads": 128,
"norm_eps": 1e-5,
"vocab_size": 131072,
"rope_theta": 10000.0,
"max_position_embeddings": 294912,
"q_lora_rank": 1536,
"kv_lora_rank": 512,
"qk_nope_head_dim": 128,
"qk_rope_head_dim": 64,
"v_head_dim": 128,
"moe": {
"num_experts": 128,
"num_experts_per_tok": 4,
"num_shared_experts": 1,
"expert_hidden_dim": 4096,
"first_k_dense_replace": 3,
"route_every_n": 1,
"routed_scale": 1.0,
"num_expert_groups": 1,
"num_expert_groups_per_tok": 1,
},
"llama_4_scaling": {"beta": 0.1, "original_max_position_embeddings": 8192},
"yarn": {
"alpha": 1,
"apply_scale": False,
"beta": 32,
"factor": 36,
"original_max_position_embeddings": 8192,
},
}
config = adapt_config_dict(params, defaults={})
assert config.architectures == ["MistralLarge3ForCausalLM"]
assert config.rope_parameters["attention_factor"] == 1.0
def test_glm5_next_accepts_deepseek_sparse_attention_layers():
layer_types = ["linear_attention", "deepseek_sparse_attention"]
+6 -6
View File
@@ -324,9 +324,9 @@ class AXK1Attention(nn.Module):
assert config.rope_parameters is not None
if config.rope_parameters["rope_type"] != "default":
config.rope_parameters["rope_type"] = (
"deepseek_yarn"
if config.rope_parameters.get("apply_yarn_scaling", True)
else "deepseek_llama_scaling"
"deepseek_llama_scaling"
if config.rope_parameters.get("attention_factor") == 1.0
else "deepseek_yarn"
)
self.rotary_emb = get_rope(
@@ -496,9 +496,9 @@ class AXK1MLAAttention(nn.Module):
assert config.rope_parameters is not None
if config.rope_parameters["rope_type"] != "default":
config.rope_parameters["rope_type"] = (
"deepseek_yarn"
if config.rope_parameters.get("apply_yarn_scaling", True)
else "deepseek_llama_scaling"
"deepseek_llama_scaling"
if config.rope_parameters.get("attention_factor") == 1.0
else "deepseek_yarn"
)
self.rotary_emb = get_rope(
+6 -6
View File
@@ -540,9 +540,9 @@ class DeepseekV2Attention(nn.Module):
)
if config.rope_parameters["rope_type"] != "default":
config.rope_parameters["rope_type"] = (
"deepseek_yarn"
if config.rope_parameters.get("apply_yarn_scaling", True)
else "deepseek_llama_scaling"
"deepseek_llama_scaling"
if config.rope_parameters.get("attention_factor") == 1.0
else "deepseek_yarn"
)
self.rotary_emb = get_rope(
@@ -1107,9 +1107,9 @@ class DeepseekV2MLAAttention(nn.Module):
if config.rope_parameters["rope_type"] != "default":
config.rope_parameters["rope_type"] = (
"deepseek_yarn"
if config.rope_parameters.get("apply_yarn_scaling", True)
else "deepseek_llama_scaling"
"deepseek_llama_scaling"
if config.rope_parameters.get("attention_factor") == 1.0
else "deepseek_yarn"
)
self.rotary_emb = get_rope(
+3 -3
View File
@@ -153,9 +153,9 @@ class DeepseekV32Attention(MLAAttention):
# DSA checkpoints may use plain ("default") or yarn-scaled RoPE.
if config.rope_parameters["rope_type"] != "default":
config.rope_parameters["rope_type"] = (
"deepseek_yarn"
if config.rope_parameters.get("apply_yarn_scaling", True)
else "deepseek_llama_scaling"
"deepseek_llama_scaling"
if config.rope_parameters.get("attention_factor") == 1.0
else "deepseek_yarn"
)
if config.rope_parameters["rope_type"] == "deepseek_yarn":
mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
+3 -3
View File
@@ -31,9 +31,9 @@ def build_deepseek_v4_rope(
if compress_ratio > 1 and rope_parameters["rope_type"] != "default":
# YaRN applies only to compressor (CSA/HCA) layers.
rope_parameters["rope_type"] = (
"deepseek_yarn"
if rope_parameters.get("apply_yarn_scaling", True)
else "deepseek_llama_scaling"
"deepseek_llama_scaling"
if rope_parameters.get("attention_factor") == 1.0
else "deepseek_yarn"
)
else:
# Sliding-window layers use plain RoPE (theta=rope_theta, no YaRN).
+3 -3
View File
@@ -32,9 +32,9 @@ def build_deepseek_v4_rope(
# YaRN applies to all compressed-KV layers (compress_ratio > 0); only
# pure sliding-window layers (compress_ratio == 0) use plain RoPE.
rope_parameters["rope_type"] = (
"deepseek_yarn"
if rope_parameters.get("apply_yarn_scaling", True)
else "deepseek_llama_scaling"
"deepseek_llama_scaling"
if rope_parameters.get("attention_factor") == 1.0
else "deepseek_yarn"
)
else:
# Sliding-window layers use plain RoPE (theta=rope_theta, no YaRN).
+3 -3
View File
@@ -496,9 +496,9 @@ class Glm5NextMLAAttention(nn.Module):
assert config.rope_parameters is not None
if config.rope_parameters["rope_type"] != "default":
config.rope_parameters["rope_type"] = (
"deepseek_yarn"
if config.rope_parameters.get("apply_yarn_scaling", True)
else "deepseek_llama_scaling"
"deepseek_llama_scaling"
if config.rope_parameters.get("attention_factor") == 1.0
else "deepseek_yarn"
)
self.rotary_emb: RotaryEmbedding | None = get_rope(
+3 -3
View File
@@ -165,9 +165,9 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
rope_parameters = dict(config.rope_parameters)
if rope_parameters["rope_type"] != "default":
rope_parameters["rope_type"] = (
"deepseek_yarn"
if rope_parameters.get("apply_yarn_scaling", True)
else "deepseek_llama_scaling"
"deepseek_llama_scaling"
if rope_parameters.get("attention_factor") == 1.0
else "deepseek_yarn"
)
self.rotary_emb = get_rope(
qk_rope_head_dim,