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 15, 08:30 UTC
[Bugfix] Fix nomic max_model_len (#18755)
wang.yuqi committed
May 28, 2025 at 11:29 UTC
3e9ce609bd0636301b3dd95deeea4b7b6df3f654
4 files changed
+242
-3
examples/offline_inference/context_extension.py
new
+46
@@ -0,0 +1,46 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from vllm import LLM, SamplingParams
4
+
5
+rope_theta = 1000000
6
+original_max_position_embeddings = 32768
7
+factor = 4.0
8
+
9
+# Use yarn to extend context
10
+hf_overrides = {
11
+ "rope_theta": rope_theta,
12
+ "rope_scaling": {
13
+ "rope_type": "yarn",
14
+ "factor": factor,
15
+ "original_max_position_embeddings": original_max_position_embeddings,
16
+ },
17
+ "max_model_len": int(original_max_position_embeddings * factor),
18
+}
19
+
20
+llm = LLM(model="Qwen/Qwen3-0.6B", hf_overrides=hf_overrides)
21
+
22
+sampling_params = SamplingParams(
23
+ temperature=0.8,
24
+ top_p=0.95,
25
+ max_tokens=128,
26
+)
27
+
28
+conversation = [
29
+ {"role": "system", "content": "You are a helpful assistant"},
30
+ {"role": "user", "content": "Hello"},
31
+ {"role": "assistant", "content": "Hello! How can I assist you today?"},
32
+]
33
+outputs = llm.chat(conversation, sampling_params, use_tqdm=False)
34
+
35
+
36
+def print_outputs(outputs):
37
+ print("\nGenerated Outputs:\n" + "-" * 80)
38
+ for output in outputs:
39
+ prompt = output.prompt
40
+ generated_text = output.outputs[0].text
41
+ print(f"Prompt: {prompt!r}\n")
42
+ print(f"Generated text: {generated_text!r}")
43
+ print("-" * 80)
44
+
45
+
46
+print_outputs(outputs)
tests/models/language/pooling/test_nomic_max_model_len.py
new
+130
@@ -0,0 +1,130 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# ruff: noqa: SIM117
3
+import pytest
4
+
5
+from ...utils import EmbedModelInfo
6
+
7
+MODELS = [
8
+ EmbedModelInfo("nomic-ai/nomic-embed-text-v1"),
9
+ #EmbedModelInfo("nomic-ai/nomic-embed-text-v1.5"),
10
+ #EmbedModelInfo("nomic-ai/CodeRankEmbed"),
11
+ EmbedModelInfo("nomic-ai/nomic-embed-text-v2-moe"),
12
+ #EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-long"),
13
+]
14
+
15
+rope_theta = 1000
16
+factor = 4.0
17
+original_max_position_embeddings = 2048
18
+max_model_len = int(original_max_position_embeddings * factor)
19
+
20
+
21
+@pytest.mark.parametrize("model_info", MODELS)
22
+def test_default(model_info, vllm_runner):
23
+ with vllm_runner(model_info.name, task="embed",
24
+ max_model_len=None) as vllm_model:
25
+ model_config = vllm_model.model.llm_engine.model_config
26
+ if model_info.name == "nomic-ai/nomic-embed-text-v2-moe":
27
+ # For nomic-embed-text-v2-moe the length is set to 512
28
+ # by sentence_bert_config.json.
29
+ assert model_config.max_model_len == 512
30
+ else:
31
+ assert (
32
+ model_config.max_model_len == original_max_position_embeddings)
33
+
34
+
35
+@pytest.mark.parametrize("model_info", MODELS)
36
+def test_set_max_model_len_legal(model_info, vllm_runner):
37
+ # set max_model_len <= 512
38
+ with vllm_runner(model_info.name, task="embed",
39
+ max_model_len=256) as vllm_model:
40
+ model_config = vllm_model.model.llm_engine.model_config
41
+ assert model_config.max_model_len == 256
42
+
43
+ # set 512 < max_model_len <= 2048
44
+ if model_info.name == "nomic-ai/nomic-embed-text-v2-moe":
45
+ # For nomic-embed-text-v2-moe the length is set to 512
46
+ # by sentence_bert_config.json.
47
+ with pytest.raises(ValueError):
48
+ with vllm_runner(model_info.name, task="embed",
49
+ max_model_len=1024):
50
+ pass
51
+ else:
52
+ with vllm_runner(model_info.name, task="embed",
53
+ max_model_len=1024) as vllm_model:
54
+ model_config = vllm_model.model.llm_engine.model_config
55
+ assert model_config.max_model_len == 1024
56
+
57
+
58
+@pytest.mark.parametrize("model_info", MODELS)
59
+def test_set_max_model_len_illegal(model_info, vllm_runner):
60
+ # set max_model_len > 2048
61
+ with pytest.raises(ValueError):
62
+ with vllm_runner(model_info.name, task="embed", max_model_len=4096):
63
+ pass
64
+
65
+ # set max_model_len > 2048 by hf_overrides
66
+ hf_overrides = {"max_model_len": 4096}
67
+ with pytest.raises(ValueError):
68
+ with vllm_runner(model_info.name,
69
+ task="embed",
70
+ max_model_len=None,
71
+ hf_overrides=hf_overrides):
72
+ pass
73
+
74
+
75
+@pytest.mark.parametrize("model_info", MODELS)
76
+def test_use_rope_scaling_legal(model_info, vllm_runner):
77
+ hf_overrides = {
78
+ "rope_theta": rope_theta,
79
+ "rope_scaling": {
80
+ "rope_type": "yarn",
81
+ "factor": factor,
82
+ "original_max_position_embeddings":
83
+ original_max_position_embeddings
84
+ },
85
+ "max_model_len": max_model_len
86
+ }
87
+
88
+ with vllm_runner(model_info.name,
89
+ task="embed",
90
+ max_model_len=None,
91
+ hf_overrides=hf_overrides):
92
+ pass
93
+
94
+
95
+@pytest.mark.parametrize("model_info", MODELS)
96
+def test_use_rope_scaling_illegal(model_info, vllm_runner):
97
+ hf_overrides = {
98
+ "rope_theta": rope_theta,
99
+ "rope_scaling": {
100
+ "rope_type": "yarn",
101
+ "factor": factor,
102
+ "original_max_position_embeddings":
103
+ original_max_position_embeddings
104
+ }
105
+ }
106
+ # illegal max_model_len
107
+ with pytest.raises(ValueError):
108
+ with vllm_runner(model_info.name,
109
+ task="embed",
110
+ max_model_len=max_model_len + 1,
111
+ hf_overrides=hf_overrides):
112
+ pass
113
+
114
+ hf_overrides = {
115
+ "rope_theta": rope_theta,
116
+ "rope_scaling": {
117
+ "rope_type": "yarn",
118
+ "factor": factor,
119
+ "original_max_position_embeddings":
120
+ original_max_position_embeddings
121
+ },
122
+ "max_model_len": max_model_len + 1
123
+ }
124
+ # illegal max_model_len by hf_overrides
125
+ with pytest.raises(ValueError):
126
+ with vllm_runner(model_info.name,
127
+ task="embed",
128
+ max_model_len=None,
129
+ hf_overrides=hf_overrides):
130
+ pass
vllm/config.py
+14
@@ -571,6 +571,7 @@ class ModelConfig:
571
572
sliding_window = None
573
574
+ self.original_max_model_len = self.max_model_len
575
self.max_model_len = _get_and_verify_max_len(
576
hf_config=self.hf_text_config,
577
max_model_len=self.max_model_len,
@@ -4471,6 +4472,19 @@ class VllmConfig:
4472
self.compilation_config.init_with_cudagraph_sizes(
4473
batch_size_capture_list)
4474
4475
+ def recalculate_max_model_len(self, max_model_len: int):
4476
+ model_config = self.model_config
4477
+ max_model_len = _get_and_verify_max_len(
4478
+ hf_config=model_config.hf_text_config,
4479
+ max_model_len=max_model_len,
4480
+ disable_sliding_window=model_config.disable_sliding_window,
4481
+ sliding_window_len=model_config.get_hf_config_sliding_window(),
4482
+ spec_target_max_model_len=model_config.spec_target_max_model_len,
4483
+ encoder_config=model_config.encoder_config)
4484
+ self.model_config.max_model_len = max_model_len
4485
+ self.scheduler_config.max_model_len = max_model_len
4486
+ self.compute_hash()
4487
+
4488
def __str__(self):
4489
return (
4490
f"model={self.model_config.model!r},"
vllm/model_executor/models/bert_with_rope.py
+52
-3
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
from collections.abc import Iterable
3
+from copy import deepcopy
4
from typing import Optional
5
6
import torch
@@ -10,6 +11,7 @@ from vllm.attention import Attention, AttentionType
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import CacheConfig, VllmConfig
13
from vllm.distributed import get_tensor_model_parallel_world_size
14
+from vllm.logger import init_logger
15
from vllm.model_executor.layers.activation import (get_act_and_mul_fn,
16
get_act_fn)
17
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
@@ -27,6 +29,8 @@ from vllm.model_executor.models.interfaces import SupportsQuant
29
from vllm.model_executor.models.utils import WeightsMapper
30
from vllm.sequence import IntermediateTensors
31
32
+logger = init_logger(__name__)
33
+
34
35
class BertWithRopeEmbedding(nn.Module):
36
@@ -513,10 +517,11 @@ class NomicBertModel(BertWithRope):
517
518
head_dim = config.hidden_size // config.num_attention_heads
519
rotary_emb_dim = head_dim * config.rotary_emb_fraction
520
+ max_trained_positions = getattr(config, "max_trained_positions", 2048)
521
config.rotary_kwargs = {
522
"head_size": head_dim,
523
"rotary_dim": rotary_emb_dim,
519
- "max_position": config.max_trained_positions,
524
+ "max_position": max_trained_positions,
525
"base": getattr(config, "rope_theta", config.rotary_emb_base),
526
"rope_scaling": getattr(config, "rope_scaling", None)
527
}
@@ -525,8 +530,52 @@ class NomicBertModel(BertWithRope):
530
# than max_trained_positions 2048, the results are consistent
531
# with SentenceTransformer.
532
# The context extension uses vllm style rope_theta and rope_scaling.
528
- # See #17785
529
-
533
+ # See #17785 #18755
534
+ if (not vllm_config.model_config.hf_overrides
535
+ and vllm_config.model_config.original_max_model_len is None):
536
+ # Default
537
+ # Reset max_model_len to max_trained_positions.
538
+ # nomic-embed-text-v2-moe the length is set to 512
539
+ # by sentence_bert_config.json.
540
+ max_model_len_before = vllm_config.model_config.max_model_len
541
+ max_model_len = min(vllm_config.model_config.max_model_len,
542
+ max_trained_positions)
543
+
544
+ vllm_config.recalculate_max_model_len(max_model_len)
545
+ logger.warning(
546
+ "Nomic context extension is disabled. "
547
+ "Changing max_model_len from %s to %s. "
548
+ "To enable context extension, see: "
549
+ "https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/context_extension.html",
550
+ max_model_len_before, vllm_config.model_config.max_model_len)
551
+ else:
552
+ # We need to re-verify max_model_len to avoid lengths
553
+ # greater than position_embedding.
554
+ model_config = vllm_config.model_config
555
+ hf_text_config = model_config.hf_text_config
556
+
557
+ if isinstance(model_config.hf_overrides, dict):
558
+ # hf_overrides_kw
559
+ max_model_len = model_config.hf_overrides.get(
560
+ "max_model_len", vllm_config.model_config.max_model_len)
561
+ else:
562
+ # hf_overrides_fn
563
+ # This might be overridden by sentence_bert_config.json.
564
+ max_model_len = vllm_config.model_config.max_model_len
565
+
566
+ # reset hf_text_config for recalculate_max_model_len.
567
+ if hasattr(hf_text_config, "max_model_len"):
568
+ delattr(hf_text_config, "max_model_len")
569
+ hf_text_config.max_position_embeddings = max_trained_positions
570
+ hf_text_config.rope_scaling = config.rotary_kwargs["rope_scaling"]
571
+
572
+ # The priority of sentence_bert_config.json is higher
573
+ # than max_position_embeddings
574
+ encoder_config = deepcopy(model_config.encoder_config)
575
+ encoder_config.pop("max_seq_length", None)
576
+ model_config.encoder_config = encoder_config
577
+
578
+ vllm_config.recalculate_max_model_len(max_model_len)
579
return config
580
581