diff --git a/tests/test_config.py b/tests/test_config.py index 827dbcd4f8..3982507b6b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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""" diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index ba1f0ef8dd..23a0f44a5d 100644 --- a/tests/transformers_utils/test_config.py +++ b/tests/transformers_utils/test_config.py @@ -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"] diff --git a/vllm/config/model.py b/vllm/config/model.py index e47f3eee96..89978c35d0 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -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( diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py index ab06448957..e23a5128b8 100644 --- a/vllm/model_executor/layers/rotary_embedding/__init__.py +++ b/vllm/model_executor/layers/rotary_embedding/__init__.py @@ -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 diff --git a/vllm/model_executor/layers/rotary_embedding/common.py b/vllm/model_executor/layers/rotary_embedding/common.py index 0a392812b6..006dba17ec 100644 --- a/vllm/model_executor/layers/rotary_embedding/common.py +++ b/vllm/model_executor/layers/rotary_embedding/common.py @@ -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( diff --git a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py index bfaf81ad00..f5e6472e23 100644 --- a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py @@ -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 diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 2cbffab8dd..a7fe51c465 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -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 diff --git a/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py deleted file mode 100644 index dd2fb9c320..0000000000 --- a/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py +++ /dev/null @@ -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, - ) diff --git a/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py index f01ca1e231..01fbf0b2c5 100644 --- a/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py +++ b/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py @@ -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 diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 93e3ad4de8..d34e391a26 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -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): diff --git a/vllm/transformers_utils/configs/mistral.py b/vllm/transformers_utils/configs/mistral.py index 2b07966914..3ec94eb0be 100644 --- a/vllm/transformers_utils/configs/mistral.py +++ b/vllm/transformers_utils/configs/mistral.py @@ -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}"