96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
import tempfile
|
|
from time import time
|
|
|
|
import pytest
|
|
|
|
from vllm import LLM, envs
|
|
from vllm.platforms import current_platform
|
|
from vllm.sampling_params import SamplingParams
|
|
|
|
if not envs.VLLM_USE_V1:
|
|
pytest.skip(
|
|
"Skipping V1 tests. Rerun with `VLLM_USE_V1=1` to test.",
|
|
allow_module_level=True,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("model_name", ["D4nt3/Qwen2.5-two-layers"])
|
|
@pytest.mark.skipif(not current_platform.is_tpu(),
|
|
reason="This test needs a TPU")
|
|
def test_sampler_compilation(model_name: str, monkeypatch):
|
|
"""
|
|
Check that no recompilation happens despite changing sampling parameters.
|
|
We can't read XLA metrics from the engine process, hence we measure time.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
monkeypatch.setenv("VLLM_XLA_CACHE_PATH", temp_dir)
|
|
# Compiling model init may still take some time, enforce_eager to skip.
|
|
llm = LLM(model_name,
|
|
enforce_eager=True,
|
|
max_num_seqs=16,
|
|
max_model_len=1024,
|
|
gpu_memory_utilization=0.5)
|
|
prompts = [
|
|
"A robot may not injure a human being",
|
|
"It is only with the heart that one can see rightly;",
|
|
]
|
|
# First inference should be slow
|
|
sampling_params = SamplingParams(
|
|
temperature=0.7,
|
|
# top_p=0.6, # TODO too slow!
|
|
top_k=10,
|
|
min_p=0.2,
|
|
max_tokens=16)
|
|
s = time()
|
|
_ = llm.generate(prompts, sampling_params)
|
|
run1 = time() - s
|
|
|
|
# Second request with different params, but for which we
|
|
# compiled for in previous eager iteration.
|
|
sampling_params = SamplingParams(temperature=0.1,
|
|
top_k=12,
|
|
min_p=0.8,
|
|
max_tokens=24)
|
|
s = time()
|
|
_ = llm.generate(prompts, sampling_params)
|
|
run2 = time() - s
|
|
# Much faster after compiling
|
|
assert run1 * 0.1 > run2
|
|
print("TIMES", run1, run2)
|
|
|
|
# Third request with min_p set to "None". It will not trigger
|
|
# recompilation as a default 0 value will be used.
|
|
sampling_params = SamplingParams(max_tokens=24, temperature=0.0)
|
|
s = time()
|
|
_ = llm.generate(prompts, sampling_params)
|
|
run3 = time() - s
|
|
assert run1 * 0.1 > run3
|
|
print("TIMES", run1, run3)
|
|
|
|
|
|
@pytest.mark.parametrize("model_name", ["Qwen/Qwen2.5-1.5B-Instruct"])
|
|
@pytest.mark.skipif(not current_platform.is_tpu(),
|
|
reason="This test needs a TPU")
|
|
def test_sampler_different(model_name: str):
|
|
"""
|
|
Test significantly different sampling params to assert the model produces
|
|
different results.
|
|
"""
|
|
llm = LLM(
|
|
model_name,
|
|
enforce_eager=True,
|
|
max_num_seqs=1,
|
|
max_model_len=64,
|
|
# TODO: setting to 0.5 or it will go OOM
|
|
gpu_memory_utilization=0.5)
|
|
prompts = [
|
|
"Write a short story about a robot that dreams for the first time."
|
|
]
|
|
sampling_params = SamplingParams(temperature=0.9, min_p=0.2, max_tokens=64)
|
|
output = llm.generate(prompts, sampling_params)
|
|
|
|
sampling_params = SamplingParams(temperature=0.1, min_p=0.8, max_tokens=64)
|
|
output2 = llm.generate(prompts, sampling_params)
|
|
assert output[0].outputs[0].text != output2[0].outputs[0].text
|