iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix] Align vLLM YaRN with Transformers and stop re-scaling max_model_len (#56446)

Browse Source 
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Harry Mellor
2026-09-11 13:01:38 +00:00
committed by GitHub
co-authored by Claude Opus 5
parent a5b714eee4
commit c191787a68
11 changed files with 133 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_config.py
+48
View File
@@ -1853,6 +1853,54 @@ def test_get_and_verify_max_len_with_nope_layers(
assert hf_config.rope_parameters["full_attention"] is None
@pytest.mark.parametrize(
("rope_type", "factor", "expected_max_len"),
[
# TeleChat3-36B-Thinking: 32768 already scaled from 8192 by 4
("yarn", 4.0, 32768),
# sarvam-105b: declares factor 40 but only serves 131072 of it
("deepseek_yarn", 40.0, 32768),
("deepseek_llama_scaling", 40.0, 32768),
# Non-YaRN scaling still multiplies
("linear", 4.0, 131072),
],
)
def test_get_and_verify_max_len_yarn_is_already_scaled(
rope_type, factor, expected_max_len
):
"""YaRN variants must not re-apply `factor` to max_position_embeddings.
Transformers treats max_position_embeddings as the final context length
for every YaRN variant, so scaling it again overstates the limit and lets
requests past the end of the cos/sin cache.
"""
from transformers import PretrainedConfig
from vllm.config.model import _get_and_verify_max_len
from vllm.transformers_utils.model_arch_config_convertor import (
ModelArchConfigConvertorBase,
)
hf_config = PretrainedConfig(max_position_embeddings=32768)
hf_config.rope_parameters = {
"rope_type": rope_type,
"factor": factor,
"original_max_position_embeddings": 8192,
}
model_arch_config = ModelArchConfigConvertorBase(hf_config, hf_config).convert()
actual_max_len = _get_and_verify_max_len(
hf_config=hf_config,
model_arch_config=model_arch_config,
tokenizer_config=None,
max_model_len=None,
disable_sliding_window=False,
sliding_window=None,
)
assert actual_max_len == expected_max_len
class MockConfig:
"""Simple mock object for testing maybe_pull_model_tokenizer_for_runai"""
tests/transformers_utils/test_config.py
+26
View File
@@ -6,6 +6,7 @@ only get the `eos_token_id` from the tokenizer as defined by
`BaseRenderer.get_eos_token_id`.
"""
import math
from types import SimpleNamespace
from typing import cast
from unittest.mock import MagicMock, patch
@@ -52,6 +53,31 @@ def test_patch_legacy_rope_type_preserves_nope_layers():
}
def test_patch_legacy_rope_type_normalizes_telechat3_yarn():
"""TeleChat3's RoPE is YaRN with 0.07 in place of the usual 0.1.
Encoding that as a precomputed attention_factor keeps the config
plain YaRN, which Transformers and every "yarn" guard understand.
`mscale` cannot express it: Transformers only applies mscale when
mscale_all_dim is also truthy.
"""
rope_parameters = {
"type": "telechat3-yarn",
"rope_type": "telechat3-yarn",
"factor": 4.0,
"original_max_position_embeddings": 8192,
}
patch_legacy_rope_type(rope_parameters)
assert rope_parameters == {
"rope_type": "yarn",
"factor": 4.0,
"original_max_position_embeddings": 8192,
"attention_factor": pytest.approx(0.07 * math.log(4.0) + 1.0),
}
def test_glm5_next_accepts_deepseek_sparse_attention_layers():
layer_types = ["linear_attention", "deepseek_sparse_attention"]
vllm/config/model.py
+11 -4
View File
@@ -2465,13 +2465,20 @@ def _get_and_verify_max_len(
# loading HF config
rope_type = rp["rope_type"]
if rope_type not in ("su", "longrope", "llama3"):
# YaRN variants leave max_position_embeddings already scaled, as
# Transformers' _compute_yarn_parameters assumes, so `factor` must
# not be applied to it again.
if rope_type not in (
"su",
"longrope",
"llama3",
"yarn",
"deepseek_yarn",
"deepseek_llama_scaling",
):
# NOTE: rope_type == "default" does not define factor https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/modeling_rope_utils.py
# NOTE: This assumes all layer types have the same scaling factor.
scaling_factor = rp.get("factor", scaling_factor)
if rope_type == "yarn":
derived_max_model_len = rp["original_max_position_embeddings"]
if scaling_factor is None:
# Fallback the factor to 1.0 if a user assigned `null`
logger.warning_once(
vllm/model_executor/layers/rotary_embedding/__init__.py
+3 -37
View File
@@ -24,7 +24,6 @@ from .mrope import MRotaryEmbedding
from .mrope_interleaved import MRotaryEmbeddingInterleaved
from .ntk_scaling_rope import NTKScalingRotaryEmbedding
from .phi3_long_rope_scaled_rope import Phi3LongRoPEScaledRotaryEmbedding
from .telechat3_scaling_rope import TeleChat3RoPEScaledRotaryEmbedding
from .yarn_scaling_rope import YaRNScalingRotaryEmbedding
_ROPE_DICT: dict[tuple[Any, ...], RotaryEmbedding] = {}
@@ -246,16 +245,15 @@ def get_rope(
for k, v in rope_parameters.items()
if k
in (
"extrapolation_factor",
"attn_factor",
"beta_fast",
"beta_slow",
"apply_yarn_scaling",
"mscale",
"mscale_all_dim",
"attention_factor",
"truncate",
)
}
if "mrope_section" in rope_parameters:
extra_kwargs.pop("apply_yarn_scaling", None)
rotary_emb = MRotaryEmbedding(
head_size,
rotary_dim,
@@ -288,8 +286,6 @@ def get_rope(
for k, v in rope_parameters.items()
if k
in (
"extrapolation_factor",
"attn_factor",
"beta_fast",
"beta_slow",
"mscale",
@@ -346,36 +342,6 @@ def get_rope(
)
else:
raise ValueError("Pangu mrope lacks necessary parameters.")
elif scaling_type == "telechat3-yarn":
scaling_factor = rope_parameters["factor"]
if "original_max_position_embeddings" in rope_parameters:
original_max_position = rope_parameters["original_max_position_embeddings"]
scaling_factor = max_position / original_max_position
else:
original_max_position = max_position
extra_kwargs = {
k: v
for k, v in rope_parameters.items()
if k
in (
"extrapolation_factor",
"attn_factor",
"beta_fast",
"beta_slow",
"mscale",
"mscale_all_dim",
)
}
rotary_emb = TeleChat3RoPEScaledRotaryEmbedding(
head_size,
rotary_dim,
original_max_position,
base,
is_neox_style,
scaling_factor,
dtype,
**extra_kwargs,
)
else:
raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
_ROPE_DICT[key] = rotary_emb
vllm/model_executor/layers/rotary_embedding/common.py
+2 -2
View File
@@ -70,10 +70,10 @@ def yarn_linear_ramp_mask(
return ramp_func
def yarn_get_mscale(scale: float = 1) -> float:
def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
if scale <= 1:
return 1.0
return 0.1 * math.log(scale) + 1.0
return 0.1 * mscale * math.log(scale) + 1.0
def _flashinfer_rotary_embedding(
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+3 -9
View File
@@ -39,8 +39,6 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
scaling_factor: float,
dtype: torch.dtype,
*,
extrapolation_factor: float = 1,
attn_factor: float = 1,
beta_fast: int = 32,
beta_slow: int = 1,
mscale: float = 1,
@@ -48,15 +46,12 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
init_cache: bool = True,
) -> None:
self.scaling_factor = scaling_factor
self.extrapolation_factor = extrapolation_factor
self.attn_factor = attn_factor
self.beta_fast = beta_fast
self.beta_slow = beta_slow
# Get n-d magnitude scaling corrected for interpolation.
self.mscale = float(
yarn_get_mscale(self.scaling_factor, float(mscale))
/ yarn_get_mscale(self.scaling_factor, float(mscale_all_dim))
* attn_factor
)
self.use_flashinfer = (
self.enabled()
@@ -96,10 +91,9 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
self.max_position_embeddings,
)
# Get n-d rotational scaling corrected for extrapolation
inv_freq_mask = (
1
- yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float)
) * self.extrapolation_factor
inv_freq_mask = 1 - yarn_linear_ramp_mask(
low, high, self.rotary_dim // 2, dtype=torch.float
)
inv_freq = (
inv_freq_interpolation * (1 - inv_freq_mask)
+ inv_freq_extrapolation * inv_freq_mask
vllm/model_executor/layers/rotary_embedding/mrope.py
+12 -5
View File
@@ -252,21 +252,28 @@ class MRotaryEmbedding(RotaryEmbeddingBase):
# YaRN parameters.
*,
scaling_factor: float | None = None,
extrapolation_factor: float = 1,
attn_factor: float = 1,
beta_fast: int = 32,
beta_slow: int = 1,
mscale: float | None = None,
mscale_all_dim: float | None = None,
attention_factor: float | None = None,
truncate: bool = True,
) -> None:
self.scaling_factor = scaling_factor
self.extrapolation_factor = extrapolation_factor
self.attn_factor = attn_factor
self.beta_fast = beta_fast
self.beta_slow = beta_slow
self.truncate = truncate
if self.scaling_factor is not None:
# Get n-d magnitude scaling corrected for interpolation
self.mscale = float(yarn_get_mscale(self.scaling_factor) * attn_factor)
if attention_factor is not None:
self.mscale = float(attention_factor)
elif mscale and mscale_all_dim:
self.mscale = float(
yarn_get_mscale(self.scaling_factor, mscale)
/ yarn_get_mscale(self.scaling_factor, mscale_all_dim)
)
else:
self.mscale = float(yarn_get_mscale(self.scaling_factor))
else:
self.mscale = 1.0
vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py
-57
View File
@@ -1,57 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
import torch
from .base import RotaryEmbedding
from .yarn_scaling_rope import YaRNScalingRotaryEmbedding
class TeleChat3RoPEScaledRotaryEmbedding(YaRNScalingRotaryEmbedding):
"""TeleChat3 uses a variant of YaRN method.
To achieve code reuse as much as possible, we have rewritten the
`get_mscale` method in the initialization function
"""
def __init__(
self,
head_size: int,
rotary_dim: int,
max_position_embeddings: int,
base: int,
is_neox_style: bool,
scaling_factor: float,
dtype: torch.dtype,
*,
extrapolation_factor: float = 1,
attn_factor: float = 1,
beta_fast: int = 32,
beta_slow: int = 1,
truncate: bool = True,
) -> None:
self.scaling_factor = scaling_factor
self.extrapolation_factor = extrapolation_factor
self.attn_factor = attn_factor
self.beta_fast = beta_fast
self.beta_slow = beta_slow
self.truncate = truncate
def get_mscale(scale, mscale=1):
if scale <= 1:
return 1.0
return 0.07 * mscale * math.log(scale) + 1.0
self.mscale = float(get_mscale(self.scaling_factor) * attn_factor)
# Initialization must be performed after mscale, otherwise mscale is useless
RotaryEmbedding.__init__(
self,
head_size,
rotary_dim,
max_position_embeddings,
base,
is_neox_style,
dtype,
)
vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
+15 -14
View File
@@ -23,25 +23,27 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
scaling_factor: float,
dtype: torch.dtype,
*,
extrapolation_factor: float = 1,
attn_factor: float = 1,
beta_fast: int = 32,
beta_slow: int = 1,
apply_yarn_scaling: bool = True,
mscale: float | None = None,
mscale_all_dim: float | None = None,
attention_factor: float | None = None,
truncate: bool = True,
) -> None:
self.scaling_factor = scaling_factor
self.extrapolation_factor = extrapolation_factor
self.attn_factor = attn_factor
self.beta_fast = beta_fast
self.beta_slow = beta_slow
self.truncate = truncate
# Get n-d magnitude scaling corrected for interpolation
self.mscale = (
float(yarn_get_mscale(self.scaling_factor) * attn_factor)
if apply_yarn_scaling
else float(attn_factor)
)
if attention_factor is not None:
self.mscale = float(attention_factor)
elif mscale and mscale_all_dim:
self.mscale = float(
yarn_get_mscale(self.scaling_factor, mscale)
/ yarn_get_mscale(self.scaling_factor, mscale_all_dim)
)
else:
self.mscale = float(yarn_get_mscale(self.scaling_factor))
super().__init__(
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
)
@@ -62,10 +64,9 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
self.truncate,
)
# Get n-d rotational scaling corrected for extrapolation
inv_freq_mask = (
1
- yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float)
) * self.extrapolation_factor
inv_freq_mask = 1 - yarn_linear_ramp_mask(
low, high, self.rotary_dim // 2, dtype=torch.float
)
inv_freq = (
inv_freq_interpolation * (1 - inv_freq_mask)
+ inv_freq_extrapolation * inv_freq_mask
vllm/transformers_utils/config.py
+9
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
import os
from collections.abc import Callable, Iterator, Sequence
from contextlib import contextmanager
@@ -539,6 +540,14 @@ def patch_legacy_rope_type(rope_parameters: dict[str, Any] | None) -> None:
)
rope_parameters["rope_type"] = "default"
logger.warning("Replacing legacy rope_type 'mrope' with 'default'")
elif rope_parameters["rope_type"] == "telechat3-yarn":
# TeleChat3 is YaRN with 0.07 in place of YaRN's 0.1 attention
# scaling coefficient. Precompute it so the config is plain YaRN.
factor = rope_parameters["factor"]
rope_parameters["rope_type"] = "yarn"
rope_parameters["attention_factor"] = 0.07 * math.log(factor) + 1.0
rope_parameters.pop("type", None)
logger.warning("Replacing rope_type 'telechat3-yarn' with 'yarn'")
# Handle nested rope_parameters in interleaved sliding attention models
if is_rope_parameters_nested(rope_parameters):
vllm/transformers_utils/configs/mistral.py
+4 -6
View File
@@ -2,9 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any
from packaging.version import Version
from transformers import PretrainedConfig, WhisperConfig
from transformers import __version__ as TRANSFORMERS_VERSION
from vllm.logger import init_logger
@@ -119,7 +117,6 @@ def _remap_mistral_yarn_args(config: dict) -> dict:
"original_max_position_embeddings": ("original_max_position_embeddings", int),
"beta": ("beta_fast", float),
"alpha": ("beta_slow", float),
"apply_scale": ("apply_yarn_scaling", bool),
}
yarn_config = config.get("yarn") or {}
@@ -136,9 +133,10 @@ def _remap_mistral_yarn_args(config: dict) -> dict:
# Cast to remove Transformers > v5 type warnings
config["rope_parameters"][new_name] = cast(yarn_config.pop(old_name))
# Ignore apply_yarn_scaling in Transformers > v5 RoPE validation to remove warnings
if Version(TRANSFORMERS_VERSION) >= Version("5.3.0.dev0"):
config["ignore_keys_at_rope_validation"] = {"apply_yarn_scaling"}
# `apply_scale: false` means no magnitude correction, which Transformers
# spells as an explicit attention_factor of 1.
if not yarn_config.pop("apply_scale", True):
config["rope_parameters"]["attention_factor"] = 1.0
assert len(yarn_config) == 0, f"Unparsed yarn config: {yarn_config}"