A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 06:45 UTC
Fix error in Dynamic NTK scaling (#41277)
Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Signed-off-by: Max de Bayser <maxdebayser@gmail.com> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
Max de Bayser committed
May 19, 2026 at 17:27 UTC
117afeea4665367a3066c1df58d4082d07fcc946
4 files changed
+45
-152
tests/models/language/pooling/test_nomic_max_model_len.py
+22
-79
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
# ruff: noqa: SIM117
4
-from typing import Any
4
5
import pytest
6
@@ -40,8 +39,8 @@ def test_default(model_info, vllm_runner):
39
# For nomic-embed-text-v2-moe the length is set to 512
40
# by sentence_bert_config.json.
41
assert model_config.max_model_len == 512
43
- else:
44
- assert model_config.max_model_len == original_max_position_embeddings
42
+ if model_info.name == "nomic-ai/nomic-embed-text-v1":
43
+ assert model_config.max_model_len == 8192
44
45
46
@pytest.mark.parametrize("model_info", MODELS)
@@ -56,10 +55,9 @@ def test_set_max_model_len_legal(model_info, vllm_runner):
55
model_config = vllm_model.llm.llm_engine.model_config
56
assert model_config.max_model_len == 256
57
59
- # set 512 < max_model_len <= 2048
58
+ # For nomic-embed-text-v2-moe the length is set to 512
59
+ # by sentence_bert_config.json.
60
if model_info.name == "nomic-ai/nomic-embed-text-v2-moe":
61
- # For nomic-embed-text-v2-moe the length is set to 512
62
- # by sentence_bert_config.json.
61
with pytest.raises(ValueError):
62
with vllm_runner(
63
model_info.name,
@@ -68,40 +66,27 @@ def test_set_max_model_len_legal(model_info, vllm_runner):
66
max_model_len=1024,
67
):
68
pass
71
- else:
72
- with vllm_runner(
73
- model_info.name,
74
- revision=model_info.revision,
75
- runner="pooling",
76
- max_model_len=1024,
77
- ) as vllm_model:
78
- model_config = vllm_model.llm.llm_engine.model_config
79
- assert model_config.max_model_len == 1024
69
+ return
70
71
+ # set 512 < max_model_len <= 2048
72
+ with vllm_runner(
73
+ model_info.name,
74
+ revision=model_info.revision,
75
+ runner="pooling",
76
+ max_model_len=1024,
77
+ ) as vllm_model:
78
+ model_config = vllm_model.llm.llm_engine.model_config
79
+ assert model_config.max_model_len == 1024
80
82
-@pytest.mark.parametrize("model_info", MODELS)
83
-def test_set_max_model_len_illegal(model_info, vllm_runner):
81
# set max_model_len > 2048
85
- with pytest.raises(ValueError):
86
- with vllm_runner(
87
- model_info.name,
88
- revision=model_info.revision,
89
- runner="pooling",
90
- max_model_len=4096,
91
- ):
92
- pass
93
-
94
- # set max_model_len > 2048 by hf_overrides
95
- hf_overrides = {"max_model_len": 4096}
96
- with pytest.raises(ValueError):
97
- with vllm_runner(
98
- model_info.name,
99
- revision=model_info.revision,
100
- runner="pooling",
101
- max_model_len=None,
102
- hf_overrides=hf_overrides,
103
- ):
104
- pass
82
+ with vllm_runner(
83
+ model_info.name,
84
+ revision=model_info.revision,
85
+ runner="pooling",
86
+ max_model_len=4096,
87
+ ) as vllm_model:
88
+ model_config = vllm_model.llm.llm_engine.model_config
89
+ assert model_config.max_model_len == 4096
90
91
92
@pytest.mark.parametrize("model_info", MODELS)
@@ -124,45 +109,3 @@ def test_use_rope_scaling_legal(model_info, vllm_runner):
109
hf_overrides=hf_overrides,
110
):
111
pass
127
-
128
-
129
-@pytest.mark.parametrize("model_info", MODELS)
130
-def test_use_rope_scaling_illegal(model_info, vllm_runner):
131
- hf_overrides: dict[str, Any] = {
132
- "rope_parameters": {
133
- "rope_theta": rope_theta,
134
- "rope_type": "yarn",
135
- "factor": factor,
136
- "original_max_position_embeddings": original_max_position_embeddings,
137
- },
138
- }
139
- # illegal max_model_len
140
- with pytest.raises(ValueError):
141
- with vllm_runner(
142
- model_info.name,
143
- revision=model_info.revision,
144
- runner="pooling",
145
- max_model_len=max_model_len + 1,
146
- hf_overrides=hf_overrides,
147
- ):
148
- pass
149
-
150
- hf_overrides = {
151
- "rope_parameters": {
152
- "rope_theta": rope_theta,
153
- "rope_type": "yarn",
154
- "factor": factor,
155
- "original_max_position_embeddings": original_max_position_embeddings,
156
- },
157
- "max_model_len": max_model_len + 1,
158
- }
159
- # illegal max_model_len by hf_overrides
160
- with pytest.raises(ValueError):
161
- with vllm_runner(
162
- model_info.name,
163
- revision=model_info.revision,
164
- runner="pooling",
165
- max_model_len=None,
166
- hf_overrides=hf_overrides,
167
- ):
168
- pass
vllm/model_executor/layers/rotary_embedding/__init__.py
+4
@@ -211,10 +211,14 @@ def get_rope(
211
)
212
elif "factor" in rope_parameters:
213
scaling_factor = rope_parameters["factor"]
214
+ max_trained_positions = rope_parameters.get(
215
+ "max_trained_positions", max_position
216
+ )
217
rotary_emb = DynamicNTKScalingRotaryEmbedding(
218
head_size,
219
rotary_dim,
220
max_position,
221
+ max_trained_positions,
222
base,
223
is_neox_style,
224
scaling_factor,
vllm/model_executor/layers/rotary_embedding/dynamic_ntk_scaling_rope.py
+8
-3
@@ -38,12 +38,14 @@ class DynamicNTKScalingRotaryEmbedding(RotaryEmbedding):
38
head_size: int,
39
rotary_dim: int,
40
max_position_embeddings: int,
41
+ max_trained_positions: int,
42
base: float,
43
is_neox_style: bool,
44
scaling_factor: float,
45
dtype: torch.dtype,
46
) -> None:
47
self.scaling_factor = scaling_factor
48
+ self.max_trained_positions = max_trained_positions
49
super().__init__(
50
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
51
)
@@ -53,13 +55,16 @@ class DynamicNTKScalingRotaryEmbedding(RotaryEmbedding):
55
# maximum length before applying the rope scaling.
56
# Thus, the maximum length after applying the rope scaling is
57
# self.max_position_embeddings * self.scaling_factor.
56
- max_len = self.max_position_embeddings * self.scaling_factor
58
base = self.base * (
58
- (self.scaling_factor * max_len / self.max_position_embeddings)
59
+ (
60
+ self.scaling_factor
61
+ * self.max_position_embeddings
62
+ / self.max_trained_positions
63
+ )
64
- (self.scaling_factor - 1)
65
) ** (self.rotary_dim / (self.rotary_dim - 2))
66
inv_freq = self._compute_inv_freq(base)
62
- t = torch.arange(max_len, dtype=torch.float)
67
+ t = torch.arange(self.max_position_embeddings, dtype=torch.float)
68
69
freqs = torch.einsum("i,j -> ij", t, inv_freq)
70
cos = freqs.cos()
vllm/model_executor/models/config.py
+11
-70
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from copy import deepcopy
3
from typing import TYPE_CHECKING
4
5
from vllm.logger import init_logger
@@ -473,80 +472,22 @@ class NomicBertModelConfig(VerifyAndUpdateConfig):
472
)
473
474
head_dim = config.hidden_size // config.num_attention_heads
476
- max_trained_positions = getattr(config, "max_trained_positions", 2048)
475
+ max_position_embeddings = getattr(config, "max_position_embeddings", 2048)
476
+ max_trained_positions = getattr(
477
+ config, "max_trained_positions", max_position_embeddings
478
+ )
479
+
480
+ rope_parameters = {
481
+ "max_trained_positions": max_trained_positions,
482
+ **(config.rope_parameters or {}),
483
+ }
484
485
config.rotary_kwargs = {
486
"head_size": head_dim,
480
- "max_position": max_trained_positions,
481
- "rope_parameters": config.rope_parameters,
487
+ "max_position": model_config.max_model_len,
488
+ "rope_parameters": rope_parameters,
489
}
490
484
- # we ignore config.rotary_scaling_factor so that for datasets shorter
485
- # than max_trained_positions 2048, the results are consistent
486
- # with SentenceTransformer.
487
- # The context extension uses vllm style rope_theta and rope_parameters.
488
- # See #17785 #18755
489
- if (
490
- not model_config.hf_overrides
491
- and model_config.original_max_model_len is None
492
- ):
493
- # Default
494
- # Reset max_model_len to max_trained_positions.
495
- # nomic-embed-text-v2-moe the length is set to 512
496
- # by sentence_bert_config.json.
497
- max_model_len_before = model_config.max_model_len
498
- max_model_len = min(model_config.max_model_len, max_trained_positions)
499
-
500
- model_config.max_model_len = model_config.get_and_verify_max_len(
501
- max_model_len
502
- )
503
-
504
- if model_config.max_model_len != max_model_len_before:
505
- logger.warning(
506
- "Nomic context extension is disabled. "
507
- "Changing max_model_len from %s to %s. "
508
- "To enable context extension, see: "
509
- "https://github.com/vllm-project/vllm/tree/main/examples/features/context_extension/context_extension_offline.py",
510
- max_model_len_before,
511
- model_config.max_model_len,
512
- )
513
- else:
514
- # We need to re-verify max_model_len to avoid lengths
515
- # greater than position_embedding.
516
- hf_text_config = model_config.hf_text_config
517
-
518
- if isinstance(model_config.hf_overrides, dict):
519
- # hf_overrides_kw
520
- max_model_len = model_config.hf_overrides.get(
521
- "max_model_len", model_config.max_model_len
522
- )
523
- else:
524
- # hf_overrides_fn
525
- # This might be overridden by sentence_bert_config.json.
526
- max_model_len = model_config.max_model_len
527
-
528
- # reset hf_text_config for recalculate_max_model_len.
529
- if hasattr(hf_text_config, "max_model_len"):
530
- delattr(hf_text_config, "max_model_len")
531
- hf_text_config.max_position_embeddings = max_trained_positions
532
- hf_text_config.rope_parameters = config.rotary_kwargs["rope_parameters"]
533
-
534
- # Update the cached derived_max_model_len to enforce the limit
535
- model_config.model_arch_config.derived_max_model_len_and_key = (
536
- float(max_trained_positions),
537
- "max_position_embeddings",
538
- )
539
-
540
- # The priority of sentence_bert_config.json is higher
541
- # than max_position_embeddings
542
- encoder_config = deepcopy(model_config.encoder_config)
543
- encoder_config.pop("max_seq_length", None)
544
- model_config.encoder_config = encoder_config
545
-
546
- model_config.max_model_len = model_config.get_and_verify_max_len(
547
- max_model_len
548
- )
549
-
491
492
class Qwen2ForProcessRewardModelConfig(VerifyAndUpdateConfig):
493
@staticmethod