[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:
co-authored by
Harry Mellor
parent
76d517fd12
commit
2bdae2a5a8
@@ -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"]
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user