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 14, 05:30 UTC
[CI] Add ci_envs for convenient local testing (#24630)
Signed-off-by: wang.yuqi <noooop@126.com>
wang.yuqi committed
Sep 12, 2025 at 16:52 UTC
d21a36f5f9569949dd6c313deed609d77e393850
4 files changed
+98
-18
tests/ci_envs.py
new
+45
@@ -0,0 +1,45 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+These envs only work for a small part of the tests, fix what you need!
5
+"""
6
+
7
+import os
8
+from typing import TYPE_CHECKING, Any, Callable, Optional
9
+
10
+if TYPE_CHECKING:
11
+ VLLM_CI_NO_SKIP: bool = False
12
+ VLLM_CI_DTYPE: Optional[str] = None
13
+ VLLM_CI_HEAD_DTYPE: Optional[str] = None
14
+ VLLM_CI_HF_DTYPE: Optional[str] = None
15
+
16
+environment_variables: dict[str, Callable[[], Any]] = {
17
+ # A model family has many models with the same architecture.
18
+ # By default, a model family tests only one model.
19
+ # Through this flag, all models can be tested.
20
+ "VLLM_CI_NO_SKIP": lambda: bool(int(os.getenv("VLLM_CI_NO_SKIP", "0"))),
21
+ # Allow changing the dtype used by vllm in tests
22
+ "VLLM_CI_DTYPE": lambda: os.getenv("VLLM_CI_DTYPE", None),
23
+ # Allow changing the head dtype used by vllm in tests
24
+ "VLLM_CI_HEAD_DTYPE": lambda: os.getenv("VLLM_CI_HEAD_DTYPE", None),
25
+ # Allow changing the head dtype used by transformers in tests
26
+ "VLLM_CI_HF_DTYPE": lambda: os.getenv("VLLM_CI_HF_DTYPE", None),
27
+}
28
+
29
+
30
+def __getattr__(name: str):
31
+ # lazy evaluation of environment variables
32
+ if name in environment_variables:
33
+ return environment_variables[name]()
34
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
35
+
36
+
37
+def __dir__():
38
+ return list(environment_variables.keys())
39
+
40
+
41
+def is_set(name: str):
42
+ """Check if an environment variable is explicitly set."""
43
+ if name in environment_variables:
44
+ return name in os.environ
45
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
tests/models/language/generation_ppl_test/ppl_utils.py
+14
-5
@@ -7,6 +7,7 @@ import pytest
7
import torch
8
from datasets import load_dataset
9
10
+import tests.ci_envs as ci_envs
11
from tests.models.utils import (GenerateModelInfo,
12
TokensTextLogprobsPromptLogprobs)
13
from vllm.logprobs import Logprob
@@ -26,19 +27,26 @@ def wikitext_ppl_test(hf_runner,
27
28
# A model family has many models with the same architecture,
29
# and we don't need to test each one.
29
- if not model_info.enable_test:
30
+ if not ci_envs.VLLM_CI_NO_SKIP and not model_info.enable_test:
31
pytest.skip("Skipping test.")
32
33
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
34
35
# Allow vllm to test using the given dtype, such as float32
36
vllm_extra_kwargs = vllm_extra_kwargs or {}
36
- vllm_extra_kwargs["dtype"] = model_info.dtype
37
+ vllm_extra_kwargs["dtype"] = ci_envs.VLLM_CI_DTYPE or model_info.dtype
38
39
# Allow vllm to test using hf_overrides
40
if model_info.hf_overrides is not None:
41
vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides
42
43
+ # Allow changing the head dtype used by vllm in tests
44
+ if ci_envs.VLLM_CI_HEAD_DTYPE is not None:
45
+ if "hf_overrides" not in vllm_extra_kwargs:
46
+ vllm_extra_kwargs["hf_overrides"] = {}
47
+ vllm_extra_kwargs["hf_overrides"][
48
+ "head_dtype"] = ci_envs.VLLM_CI_HEAD_DTYPE
49
+
50
with vllm_runner(model_info.name,
51
gpu_memory_utilization=0.7,
52
max_model_len=max_length,
@@ -46,7 +54,7 @@ def wikitext_ppl_test(hf_runner,
54
enforce_eager=True,
55
**vllm_extra_kwargs) as vllm_model:
56
# Use max_num_seqs=1 to avoid OOM,
49
- # and batch different requests together.
57
+ # and avoid batch different requests together.
58
59
model_config = vllm_model.llm.llm_engine.model_config
60
@@ -91,12 +99,13 @@ def wikitext_ppl_test(hf_runner,
99
n_tokens += len(token_log_probs)
100
vllm_ppl = float(torch.exp(nll_sum / n_tokens))
101
vllm_dtype = model_config.dtype
102
+ head_dtype = model_config.head_dtype
103
104
# Accelerate ppl test by setting Transformers ppl score to a constant
105
if model_info.hf_ppl is None:
106
with hf_runner(
107
model_info.name,
99
- dtype=model_info.hf_dtype,
108
+ dtype=ci_envs.VLLM_CI_HF_DTYPE or model_info.hf_dtype,
109
) as hf_model:
110
nll_sum = torch.tensor(0., dtype=torch.float32, device="cpu")
111
n_tokens = 0
@@ -121,7 +130,7 @@ def wikitext_ppl_test(hf_runner,
130
131
differ = (vllm_ppl - hf_ppl) / hf_ppl
132
print("Model:", model_info.name)
124
- print("VLLM:", vllm_dtype, vllm_ppl)
133
+ print("VLLM:", f"dtype:{vllm_dtype}", f"head_dtype:{head_dtype}", vllm_ppl)
134
print("Transformers:", hf_dtype, hf_ppl)
135
print("Difference (%):", differ * 100)
136
tests/models/language/pooling_mteb_test/mteb_utils.py
+30
-9
@@ -11,6 +11,7 @@ import pytest
11
import requests
12
import torch
13
14
+import tests.ci_envs as ci_envs
15
from tests.models.utils import (EmbedModelInfo, RerankModelInfo,
16
check_embeddings_close)
17
@@ -168,7 +169,7 @@ def mteb_test_embed_models(hf_runner,
169
atol=MTEB_EMBED_TOL):
170
# A model family has many models with the same architecture,
171
# and we don't need to test each one.
171
- if not model_info.enable_test:
172
+ if not ci_envs.VLLM_CI_NO_SKIP and not model_info.enable_test:
173
pytest.skip("Skipping test.")
174
175
# Test embed_dims, isnan and whether to use normalize
@@ -176,12 +177,19 @@ def mteb_test_embed_models(hf_runner,
177
178
# Allow vllm to test using the given dtype, such as float32
179
vllm_extra_kwargs = vllm_extra_kwargs or {}
179
- vllm_extra_kwargs["dtype"] = model_info.dtype
180
+ vllm_extra_kwargs["dtype"] = ci_envs.VLLM_CI_DTYPE or model_info.dtype
181
182
# Allow vllm to test using hf_overrides
183
if model_info.hf_overrides is not None:
184
vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides
185
186
+ # Allow changing the head dtype used by vllm in tests
187
+ if ci_envs.VLLM_CI_HEAD_DTYPE is not None:
188
+ if "hf_overrides" not in vllm_extra_kwargs:
189
+ vllm_extra_kwargs["hf_overrides"] = {}
190
+ vllm_extra_kwargs["hf_overrides"][
191
+ "head_dtype"] = ci_envs.VLLM_CI_HEAD_DTYPE
192
+
193
with vllm_runner(model_info.name,
194
runner="pooling",
195
max_model_len=None,
@@ -202,6 +210,7 @@ def mteb_test_embed_models(hf_runner,
210
vllm_main_score = run_mteb_embed_task(VllmMtebEncoder(vllm_model),
211
MTEB_EMBED_TASKS)
212
vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype
213
+ head_dtype = model_config.head_dtype
214
215
# Test embed_dims, isnan and whether to use normalize
216
vllm_outputs = vllm_model.embed(example_prompts,
@@ -211,9 +220,11 @@ def mteb_test_embed_models(hf_runner,
220
# Accelerate mteb test by setting
221
# SentenceTransformers mteb score to a constant
222
if model_info.mteb_score is None:
214
- with hf_runner(model_info.name,
215
- is_sentence_transformer=True,
216
- dtype=model_info.hf_dtype) as hf_model:
223
+ with hf_runner(
224
+ model_info.name,
225
+ is_sentence_transformer=True,
226
+ dtype=ci_envs.VLLM_CI_HF_DTYPE or model_info.hf_dtype,
227
+ ) as hf_model:
228
229
# e.g. setting default parameters for the encode method of hf_runner
230
if hf_model_callback is not None:
@@ -236,7 +247,8 @@ def mteb_test_embed_models(hf_runner,
247
st_dtype = "Constant"
248
249
print("Model:", model_info.name)
239
- print("VLLM:", vllm_dtype, vllm_main_score)
250
+ print("VLLM:", f"dtype:{vllm_dtype}", f"head_dtype:{head_dtype}",
251
+ vllm_main_score)
252
print("SentenceTransformers:", st_dtype, st_main_score)
253
print("Difference:", st_main_score - vllm_main_score)
254
@@ -319,17 +331,24 @@ def mteb_test_rerank_models(hf_runner,
331
atol=MTEB_RERANK_TOL):
332
# A model family has many models with the same architecture,
333
# and we don't need to test each one.
322
- if not model_info.enable_test:
334
+ if not ci_envs.VLLM_CI_NO_SKIP and not model_info.enable_test:
335
pytest.skip("Skipping test.")
336
337
# Allow vllm to test using the given dtype, such as float32
338
vllm_extra_kwargs = vllm_extra_kwargs or {}
327
- vllm_extra_kwargs["dtype"] = model_info.dtype
339
+ vllm_extra_kwargs["dtype"] = ci_envs.VLLM_CI_DTYPE or model_info.dtype
340
341
# Allow vllm to test using hf_overrides
342
if model_info.hf_overrides is not None:
343
vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides
344
345
+ # Allow changing the head dtype used by vllm in tests
346
+ if ci_envs.VLLM_CI_HEAD_DTYPE is not None:
347
+ if "hf_overrides" not in vllm_extra_kwargs:
348
+ vllm_extra_kwargs["hf_overrides"] = {}
349
+ vllm_extra_kwargs["hf_overrides"][
350
+ "head_dtype"] = ci_envs.VLLM_CI_HEAD_DTYPE
351
+
352
with vllm_runner(model_info.name,
353
runner="pooling",
354
max_model_len=None,
@@ -355,6 +374,7 @@ def mteb_test_rerank_models(hf_runner,
374
tasks=MTEB_RERANK_TASKS,
375
languages=MTEB_RERANK_LANGS)
376
vllm_dtype = model_config.dtype
377
+ head_dtype = model_config.head_dtype
378
379
# Accelerate mteb test by setting
380
# SentenceTransformers mteb score to a constant
@@ -366,7 +386,8 @@ def mteb_test_rerank_models(hf_runner,
386
st_dtype = "Constant"
387
388
print("Model:", model_info.name)
369
- print("VLLM:", vllm_dtype, vllm_main_score)
389
+ print("VLLM:", f"dtype:{vllm_dtype}", f"head_dtype:{head_dtype}",
390
+ vllm_main_score)
391
print("SentenceTransformers:", st_dtype, st_main_score)
392
print("Difference:", st_main_score - vllm_main_score)
393
vllm/config/__init__.py
+9
-4
@@ -1775,16 +1775,21 @@ class ModelConfig:
1775
such as the lm_head in a generation model,
1776
or the score or classifier in a classification model.
1777
1778
- The default head_dtype based on runner_type.\n
1778
+ `head_dtype` currently only supports pooling models.\n
1779
- The pooling model defaults to using fp32 head,
1780
- you can use --hf-overrides '{"head_dtype": "model"}' to disable it.\n
1781
- - The generate model defaults to not using fp32 head,
1782
- you can use --hf-overrides '{"head_dtype": "float32"}' to enable it.
1780
+ you can use --hf-overrides '{"head_dtype": "model"}' to disable it.
1781
"""
1782
+
1783
head_dtype = _get_head_dtype(config=self.hf_config,
1784
dtype=self.dtype,
1785
runner_type=self.runner_type)
1786
1787
+ if self.runner_type != "pooling" and head_dtype != self.dtype:
1788
+ logger.warning_once(
1789
+ "`head_dtype` currently only supports pooling models."
1790
+ "fallback to model dtype [%s].", self.dtype)
1791
+ return self.dtype
1792
+
1793
if head_dtype not in current_platform.supported_dtypes:
1794
logger.warning_once(
1795
"The current platform does not support [%s] head dtype, "