chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1 @@
|
||||
{"transcriptions": ["There is no clear relationship between the barking and the music, as they seem to be independent of each other.", "(B) To indicate that language cannot express clearly, satirizing the inversion of black and white in the world"], "token_ids": [[3862, 374, 902, 2797, 5025, 1948, 279, 293, 33452, 323, 279, 4627, 11, 438, 807, 2803, 311, 387, 9489, 315, 1817, 1008, 13, 151645], [5349, 8, 2014, 13216, 429, 4128, 4157, 3158, 9355, 11, 7578, 404, 4849, 279, 46488, 315, 3691, 323, 4158, 304, 279, 1879, 151645, 151671]]}
|
||||
@@ -0,0 +1 @@
|
||||
{"transcriptions": ["There is no clear relationship between the barking and the music, as they seem to be independent of each other."], "token_ids": [[3862, 374, 902, 2797, 5025, 1948, 279, 293, 33452, 323, 279, 4627, 11, 438, 807, 2803, 311, 387, 9489, 315, 1817, 1008, 13, 151645]]}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1 @@
|
||||
{"transcriptions": ["This track is an energetic Eurodance / Dance‑Pop anthem that blends the bright, melodic sensibilities of mainstream pop with the driving, club‑ready pulse of classic Eurodance. The duration of the piece is ", "**Verse 1**\nMidnight cravings in bloom, lights flicker in the room, pepperoni dreams arise, pizza party on your skies\n\n**Verse 2**\nCheese melts on the crust, in flavor we trust, boxes stacked to the"], "token_ids": [[1986, 3754, 374, 458, 44855, 19461, 98875, 378, 107, 14, 378, 107, 35, 681, 55964, 11598, 55564, 429, 57843, 279, 9906, 11, 10581, 52760, 6097, 13450, 315, 20729, 2420, 448, 279, 9842, 11, 6335, 55964, 2307, 27235, 315, 11416, 19461, 98875, 13, 220, 576, 8090, 315, 279, 6573, 374, 220], [334, 68043, 220, 16, 1019, 33648, 9287, 88828, 304, 51454, 11, 12711, 28347, 261, 304, 279, 3054, 11, 24353, 20783, 18707, 30789, 11, 22502, 4614, 389, 697, 49293, 271, 334, 68043, 220, 17, 1019, 26843, 2367, 98091, 389, 279, 39612, 11, 304, 17172, 582, 6950, 11, 14697, 41315, 311, 279]]}
|
||||
@@ -0,0 +1 @@
|
||||
{"transcriptions": ["This track is an energetic Eurodance / Dance‑Pop anthem that blends the bright, melodic sensibilities of mainstream pop with the driving, club‑ready pulse of classic Eurodance. The duration of the piece is "], "token_ids": [[1986, 3754, 374, 458, 44855, 19461, 98875, 378, 107, 14, 378, 107, 35, 681, 55964, 11598, 55564, 429, 57843, 279, 9906, 11, 10581, 52760, 6097, 13450, 315, 20729, 2420, 448, 279, 9842, 11, 6335, 55964, 2307, 27235, 315, 11416, 19461, 98875, 13, 220, 576, 8090, 315, 279, 6573, 374, 220]]}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1 @@
|
||||
[[[0.0006361007690429688, 0.99951171875], [0.81884765625, 0.1812744140625], [0.025543212890625, 0.974609375], [0.0004382133483886719, 0.99951171875]]]
|
||||
@@ -0,0 +1,46 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Pytest configuration for vLLM language generation tests."""
|
||||
|
||||
import os
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
"""Early ROCm configuration that must happen before test collection."""
|
||||
if not current_platform.is_rocm():
|
||||
return
|
||||
|
||||
# Disable skinny GEMM on ROCm to avoid non-deterministic results
|
||||
# from atomic reductions in wvSplitKrc kernel.
|
||||
# See: https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975
|
||||
os.environ["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
|
||||
warnings.warn(
|
||||
"ROCm: Set VLLM_ROCM_USE_SKINNY_GEMM=0 to avoid non-deterministic "
|
||||
"results from skinny GEMM atomic reductions",
|
||||
UserWarning,
|
||||
stacklevel=1,
|
||||
)
|
||||
|
||||
|
||||
def pytest_sessionstart(session):
|
||||
"""Configure ROCm-specific settings before test session starts."""
|
||||
if not current_platform.is_rocm():
|
||||
return
|
||||
|
||||
# Disable Flash/MemEfficient SDP on ROCm to avoid HF Transformers
|
||||
# accuracy issues: https://github.com/vllm-project/vllm/issues/30167
|
||||
# TODO: Remove once ROCm SDP accuracy issues are resolved on HuggingFace
|
||||
torch.backends.cuda.enable_flash_sdp(False)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(False)
|
||||
torch.backends.cuda.enable_math_sdp(True)
|
||||
warnings.warn(
|
||||
"ROCm: Disabled flash_sdp and mem_efficient_sdp, enabled math_sdp "
|
||||
"to avoid HuggingFace Transformers accuracy issues",
|
||||
UserWarning,
|
||||
stacklevel=1,
|
||||
)
|
||||
@@ -0,0 +1,232 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from packaging.version import Version
|
||||
from transformers import __version__ as TRANSFORMERS_VERSION
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....utils import large_gpu_mark
|
||||
from ...registry import HF_EXAMPLE_MODELS
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
# Models that require embedding scaling for prompt_embeds test
|
||||
EMBED_SCALING_MODELS = {
|
||||
"openbmb/MiniCPM4.1-8B",
|
||||
}
|
||||
|
||||
# This list contains the model that are using AITER kernel.
|
||||
# Skip model that are not using AITER tests.
|
||||
# When more AITER kernels are added, this list will not be
|
||||
# needed as all the models will be calling AITER kernels
|
||||
# in parts of the operators
|
||||
AITER_MODEL_LIST = [
|
||||
"meta-llama/Llama-3.2-1B-Instruct",
|
||||
"openbmb/MiniCPM3-4B",
|
||||
"Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"TitanML/tiny-mixtral",
|
||||
"Qwen/Qwen3-8B",
|
||||
]
|
||||
|
||||
|
||||
# @maybe_test_rocm_aiter
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
pytest.param(
|
||||
"bigscience/bloom-560m", # bloom - testing alibi slopes
|
||||
marks=[
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.slow_test,
|
||||
pytest.mark.cpu_model,
|
||||
],
|
||||
),
|
||||
pytest.param(
|
||||
"openai-community/gpt2", # gpt2
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
pytest.param("Milos/slovak-gpt-j-405M"), # gptj
|
||||
pytest.param("bigcode/tiny_starcoder_py"), # gpt_bigcode
|
||||
pytest.param("EleutherAI/pythia-70m"), # gpt_neox
|
||||
pytest.param(
|
||||
"google/gemma-1.1-2b-it", # gemma
|
||||
marks=[
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.cpu_model,
|
||||
pytest.mark.slow_test,
|
||||
],
|
||||
),
|
||||
pytest.param(
|
||||
"google/gemma-2-2b-it", # test hybrid attention
|
||||
marks=[pytest.mark.cpu_model],
|
||||
),
|
||||
pytest.param(
|
||||
"zai-org/chatglm3-6b", # chatglm (text-only)
|
||||
),
|
||||
pytest.param(
|
||||
"meta-llama/Llama-3.2-1B-Instruct", # llama
|
||||
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
|
||||
),
|
||||
pytest.param(
|
||||
"openbmb/MiniCPM4.1-8B", # minicpm
|
||||
marks=[pytest.mark.core_model, large_gpu_mark(min_gb=48)],
|
||||
),
|
||||
pytest.param(
|
||||
"facebook/opt-125m", # opt
|
||||
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
|
||||
),
|
||||
pytest.param(
|
||||
"microsoft/phi-2", # phi
|
||||
marks=[pytest.mark.core_model, pytest.mark.slow_test],
|
||||
),
|
||||
pytest.param(
|
||||
"Qwen/Qwen2.5-0.5B-Instruct", # qwen2
|
||||
marks=[
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.cpu_model,
|
||||
pytest.mark.slow_test,
|
||||
],
|
||||
),
|
||||
pytest.param(
|
||||
"Qwen/Qwen3-8B", # qwen (text-only)
|
||||
),
|
||||
pytest.param("stabilityai/stablelm-3b-4e1t"), # stablelm
|
||||
pytest.param("bigcode/starcoder2-3b"), # starcoder2
|
||||
pytest.param(
|
||||
"TitanML/tiny-mixtral", # mixtral
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
pytest.param("swiss-ai/Apertus-8B-Instruct-2509"), # apertus
|
||||
pytest.param(
|
||||
"naver-hyperclovax/HyperCLOVAX-SEED-Think-14B", # hyperclovax
|
||||
marks=[large_gpu_mark(min_gb=32)],
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("max_tokens", [32])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.parametrize(
|
||||
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
|
||||
)
|
||||
@pytest.mark.parametrize("use_prompt_embeds", [True, False])
|
||||
def test_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
use_rocm_aiter: bool,
|
||||
use_prompt_embeds: bool,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
if use_rocm_aiter and (model in AITER_MODEL_LIST):
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
if model == "TitanML/tiny-mixtral":
|
||||
# Untrained model: near-uniform logits make argmax sensitive to
|
||||
# AITER's bfloat16 rounding error. Route the plain rms_norm and the
|
||||
# fused MoE (whose near-uniform router logits flip expert selection
|
||||
# under ~1 ULP drift) through the native kernels for this model.
|
||||
# See ROCm/aiter#3806 for the tracking issue and minimal repro.
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_AITER_RMSNORM", "0")
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
|
||||
elif use_rocm_aiter and model not in AITER_MODEL_LIST:
|
||||
# Skip model that are not using AITER tests.
|
||||
# When more AITER kernels are added, this list will not be
|
||||
# needed as all the models will be calling AITER kernels
|
||||
# in parts of the operators
|
||||
pytest.skip(f"Skipping '{model}' model test with AITER kernel.")
|
||||
|
||||
if model == "bigcode/starcoder2-3b":
|
||||
# Replace example.txt's Test1 (an NL prompt) with a code prompt:
|
||||
# starcoder2-3b is a code model, so NL prompts give near-uniform
|
||||
# digit logits where HF<->vLLM bf16 drift can reorder top-K.
|
||||
example_prompts = list(example_prompts)
|
||||
example_prompts[1] = (
|
||||
"def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - "
|
||||
)
|
||||
|
||||
with hf_runner(
|
||||
model,
|
||||
revision=model_info.revision,
|
||||
trust_remote_code=model_info.trust_remote_code,
|
||||
) as hf_model:
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
prompt_embeds: list[torch.Tensor] | None = [] if use_prompt_embeds else None
|
||||
|
||||
for prompt in example_prompts:
|
||||
token_ids = hf_model.tokenizer(prompt, return_tensors="pt").input_ids.to(
|
||||
hf_model.model.device
|
||||
)
|
||||
if prompt_embeds is not None:
|
||||
embed = hf_model.model.get_input_embeddings()(token_ids)
|
||||
|
||||
if "gemma" in model.lower() and (
|
||||
Version(TRANSFORMERS_VERSION) < Version("5.3.0.dev0")
|
||||
):
|
||||
# For Gemma 1/2 models with Transformers 5.4.0+, the prompt
|
||||
# embeddings are normalised in `get_prompt_embeddings`,
|
||||
# like Gemma 3. For older versions, we need to manually normalise.
|
||||
embed_scale = hf_model.config.hidden_size**0.5
|
||||
normalizer = torch.tensor(embed_scale, dtype=embed.dtype)
|
||||
embed *= normalizer
|
||||
|
||||
# MiniCPM models apply scale_emb to embeddings internally.
|
||||
# vLLM expects pre-scaled embeddings when using inputs_embeds.
|
||||
if model in EMBED_SCALING_MODELS:
|
||||
config = hf_model.model.config
|
||||
embed = embed * config.scale_emb
|
||||
|
||||
prompt_embeds.append(embed.squeeze(0))
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
tokenizer_name=model_info.tokenizer or model,
|
||||
tokenizer_mode=model_info.tokenizer_mode,
|
||||
revision=model_info.revision,
|
||||
trust_remote_code=model_info.trust_remote_code,
|
||||
# Remove the effects of batch variance on ROCm since batch invariance
|
||||
# is not yet supported.
|
||||
# See: https://github.com/vllm-project/vllm/issues/27433
|
||||
max_num_seqs=1 if current_platform.is_rocm() else 2,
|
||||
enable_prompt_embeds=use_prompt_embeds,
|
||||
compilation_config={"cudagraph_capture_sizes": [1, 2]},
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
if prompt_embeds is not None:
|
||||
vllm_outputs_from_embeds = vllm_model.generate_greedy_logprobs(
|
||||
prompt_embeds, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
if prompt_embeds is not None:
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=vllm_outputs,
|
||||
outputs_1_lst=vllm_outputs_from_embeds,
|
||||
name_0="vllm",
|
||||
name_1="vllm_from_embeds",
|
||||
)
|
||||
|
||||
if use_rocm_aiter:
|
||||
# this is to ensure that vllm engine
|
||||
# has deallocated the memory before running the next
|
||||
# unit tests. On ROCm, when using AITER
|
||||
# the memory might not be deallocated completely
|
||||
# before running the next test case
|
||||
torch.accelerator.synchronize()
|
||||
@@ -0,0 +1,27 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
MODELS = ["google/gemma-2b", "google/gemma-2-2b", "google/gemma-3-4b-it"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
def test_dummy_loader(vllm_runner, monkeypatch, model: str) -> None:
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
with vllm_runner(
|
||||
model,
|
||||
load_format="dummy",
|
||||
) as llm:
|
||||
if model == "google/gemma-3-4b-it":
|
||||
normalizers = llm.llm.collective_rpc(
|
||||
lambda self: self.model_runner.model.language_model.model.normalizer.cpu().item() # noqa: E501
|
||||
)
|
||||
config = llm.llm.llm_engine.model_config.hf_config.text_config
|
||||
else:
|
||||
normalizers = llm.llm.collective_rpc(
|
||||
lambda self: self.model_runner.model.model.normalizer.cpu().item()
|
||||
)
|
||||
config = llm.llm.llm_engine.model_config.hf_config
|
||||
assert np.allclose(normalizers, config.hidden_size**0.5, rtol=2e-3)
|
||||
@@ -0,0 +1,42 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
MODELS = [
|
||||
# TODO(sang): Sliding window should be tested separately.
|
||||
"ibm/PowerLM-3b",
|
||||
"ibm/PowerMoE-3b",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.cpu_model
|
||||
def test_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
with hf_runner(model, dtype=dtype) as hf_model:
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
with vllm_runner(model, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
@@ -0,0 +1,915 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager, nullcontext
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.registry import HF_EXAMPLE_MODELS
|
||||
from tests.utils import multi_gpu_test
|
||||
from vllm import LLM
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
|
||||
|
||||
from ...utils import check_logprobs_close, check_outputs_equal
|
||||
|
||||
# Mark all tests as hybrid
|
||||
pytestmark = pytest.mark.hybrid_model
|
||||
|
||||
# NOTE: The first model in each list is taken as the primary model,
|
||||
# meaning that it will be used in all tests in this file
|
||||
# The rest of the models will only be tested by test_models
|
||||
|
||||
APC_MULTIPLY_BY = 300
|
||||
|
||||
SSM_MODELS = [
|
||||
"state-spaces/mamba-130m-hf",
|
||||
"tiiuae/falcon-mamba-tiny-dev",
|
||||
# mamba2-codestral in transformers is broken pending:
|
||||
# https://github.com/huggingface/transformers/pull/40861
|
||||
# "yujiepan/mamba2-codestral-v0.1-tiny-random",
|
||||
]
|
||||
|
||||
HYBRID_MODELS = [
|
||||
"ai21labs/Jamba-tiny-dev",
|
||||
"pfnet/plamo-2-1b",
|
||||
"Zyphra/Zamba2-1.2B-instruct",
|
||||
"ibm-granite/granite-4.0-tiny-preview",
|
||||
"tiiuae/Falcon-H1-0.5B-Base",
|
||||
"LiquidAI/LFM2-1.2B",
|
||||
"tiny-random/qwen3-next-moe",
|
||||
]
|
||||
|
||||
HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL = {
|
||||
"LiquidAI/LFM2-1.2B",
|
||||
"tiny-random/qwen3-next-moe",
|
||||
}
|
||||
|
||||
FULL_CUDA_GRAPH_MODELS = [
|
||||
"ai21labs/Jamba-tiny-dev",
|
||||
"pfnet/plamo-2-1b",
|
||||
"Zyphra/Zamba2-1.2B-instruct",
|
||||
]
|
||||
|
||||
FP32_STATE_MODELS = [
|
||||
"state-spaces/mamba-130m-hf",
|
||||
"Zyphra/Zamba2-1.2B-instruct",
|
||||
]
|
||||
|
||||
# Avoid OOM
|
||||
MAX_NUM_SEQS = 4
|
||||
|
||||
ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
|
||||
|
||||
|
||||
def _set_conv_state_layout(monkeypatch, layout: str) -> None:
|
||||
"""Set conv state layout env var and clear cache to pick up new value."""
|
||||
from vllm.model_executor.layers.mamba import mamba_utils
|
||||
|
||||
monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", layout)
|
||||
mamba_utils.get_conv_state_layout.cache_clear()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
def test_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
try:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
with hf_runner(model) as hf_model:
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
extra_kwargs = {}
|
||||
if model in HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL:
|
||||
extra_kwargs["enable_chunked_prefill"] = True
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_num_seqs=MAX_NUM_SEQS,
|
||||
attention_backend=ATTN_BACKEND,
|
||||
**extra_kwargs,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"])
|
||||
def test_batching(
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
conv_state_layout: str,
|
||||
) -> None:
|
||||
try:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
_set_conv_state_layout(monkeypatch, conv_state_layout)
|
||||
|
||||
for_loop_outputs = []
|
||||
with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
|
||||
for prompt in example_prompts:
|
||||
(single_output,) = vllm_model.generate_greedy_logprobs(
|
||||
[prompt], max_tokens, num_logprobs
|
||||
)
|
||||
for_loop_outputs.append(single_output)
|
||||
|
||||
batched_outputs = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=for_loop_outputs,
|
||||
outputs_1_lst=batched_outputs,
|
||||
name_0="for_loop_vllm",
|
||||
name_1="batched_vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
|
||||
@pytest.mark.parametrize("max_tokens", [10])
|
||||
@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"])
|
||||
def test_chunked_prefill_with_parallel_sampling(
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
conv_state_layout: str,
|
||||
) -> None:
|
||||
"""
|
||||
Tests chunked prefill in conjunction with n > 1.
|
||||
|
||||
In this case, prefill is populated with decoding tokens and
|
||||
we test that it doesn't fail.
|
||||
|
||||
This test might fail if cache is not allocated correctly for n > 1
|
||||
decoding steps inside a chunked prefill forward pass
|
||||
(where we have both prefill and decode together)
|
||||
"""
|
||||
_set_conv_state_layout(monkeypatch, conv_state_layout)
|
||||
|
||||
sampling_params = SamplingParams(n=3, temperature=1, seed=0, max_tokens=max_tokens)
|
||||
with vllm_runner(
|
||||
model,
|
||||
enable_chunked_prefill=True,
|
||||
# forces prefill chunks with decoding
|
||||
max_num_batched_tokens=MAX_NUM_SEQS * 3,
|
||||
max_num_seqs=MAX_NUM_SEQS,
|
||||
attention_backend=ATTN_BACKEND,
|
||||
) as vllm_model:
|
||||
vllm_model.generate(example_prompts, sampling_params)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
|
||||
@pytest.mark.parametrize("max_tokens", [20])
|
||||
@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"])
|
||||
def test_mamba_cache_cg_padding(
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
conv_state_layout: str,
|
||||
) -> None:
|
||||
"""
|
||||
This test is for verifying that mamba cache is padded to CG captured
|
||||
batch size. If it's not, a torch RuntimeError will be raised because
|
||||
tensor dimensions aren't compatible.
|
||||
"""
|
||||
_set_conv_state_layout(monkeypatch, conv_state_layout)
|
||||
|
||||
vllm_config = EngineArgs(model=model, trust_remote_code=True).create_engine_config()
|
||||
cudagraph_dispatcher = CudagraphDispatcher(vllm_config)
|
||||
cudagraph_dispatcher.initialize_cudagraph_keys(
|
||||
vllm_config.compilation_config.cudagraph_mode
|
||||
)
|
||||
while (
|
||||
len(example_prompts)
|
||||
== cudagraph_dispatcher.dispatch(len(example_prompts))[1].num_tokens
|
||||
):
|
||||
example_prompts.append(example_prompts[0])
|
||||
|
||||
try:
|
||||
with vllm_runner(model) as vllm_model:
|
||||
vllm_model.generate_greedy(example_prompts, max_tokens)
|
||||
except RuntimeError:
|
||||
pytest.fail(
|
||||
"Couldn't run batch size which is not equal to a Cuda Graph "
|
||||
"captured batch size. "
|
||||
"Could be related to mamba cache not padded correctly"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
|
||||
def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
) -> None:
|
||||
"""
|
||||
This test is for verifying that the hybrid inner state management doesn't
|
||||
collapse in case where the number of incoming requests and
|
||||
finished_requests_ids is larger than the maximum mamba block capacity.
|
||||
|
||||
This could generally happen due to the fact that hybrid does support
|
||||
statelessness mechanism where it can clean up new incoming requests in
|
||||
a single step.
|
||||
"""
|
||||
try:
|
||||
with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
|
||||
vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
|
||||
except ValueError:
|
||||
pytest.fail(
|
||||
"Hybrid inner state wasn't cleaned up properly between"
|
||||
"steps finished requests registered unnecessarily "
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
|
||||
def test_state_cleanup(
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
) -> None:
|
||||
"""
|
||||
This test is for verifying that the Hybrid state is cleaned up between
|
||||
steps.
|
||||
|
||||
If it's not cleaned, an error would be expected.
|
||||
"""
|
||||
try:
|
||||
with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
|
||||
for _ in range(10):
|
||||
vllm_model.generate_greedy([example_prompts[0]] * 100, 1)
|
||||
except ValueError:
|
||||
pytest.fail(
|
||||
"Hybrid inner state wasn't cleaned up between states, "
|
||||
"could be related to finished_requests_ids"
|
||||
)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
def test_distributed_correctness(
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
with vllm_runner(
|
||||
model, tensor_parallel_size=1, max_num_seqs=MAX_NUM_SEQS
|
||||
) as vllm_model:
|
||||
vllm_outputs_tp_1 = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
with vllm_runner(
|
||||
model, tensor_parallel_size=2, max_num_seqs=MAX_NUM_SEQS
|
||||
) as vllm_model:
|
||||
vllm_outputs_tp_2 = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=vllm_outputs_tp_1,
|
||||
outputs_1_lst=vllm_outputs_tp_2,
|
||||
name_0="vllm_tp_1",
|
||||
name_1="vllm_tp_2",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", FULL_CUDA_GRAPH_MODELS)
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
def test_full_cuda_graph(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
try:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
with hf_runner(model) as hf_model:
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
with vllm_runner(
|
||||
model, max_num_seqs=MAX_NUM_SEQS, attention_backend=ATTN_BACKEND
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", FP32_STATE_MODELS)
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.parametrize(
|
||||
"cache_dtype_param", ["mamba_ssm_cache_dtype", "mamba_cache_dtype"]
|
||||
)
|
||||
def test_fp32_cache_state(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
cache_dtype_param: str,
|
||||
) -> None:
|
||||
try:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
with hf_runner(model) as hf_model:
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
with vllm_runner(
|
||||
model, max_num_seqs=MAX_NUM_SEQS, **{cache_dtype_param: "float32"}
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
# Helper functions for the APC tests
|
||||
def _get_vllm_runner_params(
|
||||
model: str,
|
||||
max_model_len: int,
|
||||
tensor_parallel_size: int = 1,
|
||||
):
|
||||
return {
|
||||
"model_name": model,
|
||||
"enable_chunked_prefill": True,
|
||||
"enable_prefix_caching": False,
|
||||
"max_model_len": max_model_len,
|
||||
"tensor_parallel_size": tensor_parallel_size,
|
||||
"gpu_memory_utilization": 0.4,
|
||||
"attention_backend": ATTN_BACKEND,
|
||||
}
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _owned_vllm_runner(vllm_runner, kwargs):
|
||||
with vllm_runner(**kwargs) as runner:
|
||||
yield runner
|
||||
|
||||
|
||||
def _get_vLLM_output(
|
||||
vllm_runner,
|
||||
kwargs,
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs,
|
||||
num_repetitions=1,
|
||||
vllm_model=None,
|
||||
):
|
||||
runner_context = (
|
||||
_owned_vllm_runner(vllm_runner, kwargs)
|
||||
if vllm_model is None
|
||||
else nullcontext(vllm_model)
|
||||
)
|
||||
with runner_context as runner:
|
||||
outs = []
|
||||
for _ in range(num_repetitions):
|
||||
if num_logprobs < 0:
|
||||
vllm_output = runner.generate_greedy(prompts, max_tokens)
|
||||
else:
|
||||
vllm_output = runner.generate_greedy_logprobs(
|
||||
prompts, max_tokens, num_logprobs
|
||||
)
|
||||
outs.append(vllm_output)
|
||||
|
||||
return outs, vllm_model
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [HYBRID_MODELS[0]])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("n_repetitions", [2])
|
||||
# If num_logprobs is set to -1, then the stringent version
|
||||
# of the test is executed using `check_outputs_equal`
|
||||
# instead of `check_logprobs_close`
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1])
|
||||
def test_apc_single_prompt(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
n_repetitions: int,
|
||||
num_logprobs: int,
|
||||
tensor_parallel_size: int,
|
||||
) -> None:
|
||||
try:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
compare_operator: Callable = (
|
||||
check_logprobs_close if num_logprobs > 0 else check_outputs_equal # type: ignore
|
||||
)
|
||||
|
||||
# Sample prompts.
|
||||
generated_prompts = [APC_MULTIPLY_BY * example_prompts[0]]
|
||||
|
||||
max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
|
||||
vllm_runner_kwargs = _get_vllm_runner_params(
|
||||
model, max_model_len, tensor_parallel_size=tensor_parallel_size
|
||||
)
|
||||
vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
|
||||
vllm_outputs_no_cache, _ = _get_vLLM_output(
|
||||
vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
vllm_runner_kwargs["enable_prefix_caching"] = True
|
||||
vllm_outputs_cache_rep, _ = _get_vLLM_output(
|
||||
vllm_runner,
|
||||
vllm_runner_kwargs,
|
||||
generated_prompts,
|
||||
max_tokens,
|
||||
num_logprobs,
|
||||
n_repetitions,
|
||||
)
|
||||
|
||||
for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
|
||||
# In the first repetition, the caches are filled
|
||||
# In the second repetition, these caches are reused
|
||||
|
||||
compare_operator(
|
||||
outputs_0_lst=vllm_outputs_no_cache[0],
|
||||
outputs_1_lst=vllm_outputs_cache_itn,
|
||||
name_0="vllm_no_cache",
|
||||
name_1=f"vllm_cache_it_{r_idx + 1}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [HYBRID_MODELS[0]])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("n_repetitions", [2])
|
||||
# If num_logprobs is set to -1, then the stringent version
|
||||
# of the test is executed using `check_outputs_equal`
|
||||
# instead of `check_logprobs_close`
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1])
|
||||
def test_apc_single_prompt_block_align_alignment(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
n_repetitions: int,
|
||||
num_logprobs: int,
|
||||
tensor_parallel_size: int,
|
||||
) -> None:
|
||||
try:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
compare_operator: Callable = (
|
||||
check_logprobs_close if num_logprobs > 0 else check_outputs_equal # type: ignore
|
||||
)
|
||||
|
||||
# Sample prompts. This custom prompt is used, as it causes the most issues
|
||||
generated_prompts = ["The president of the United States is " * APC_MULTIPLY_BY]
|
||||
|
||||
max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
|
||||
vllm_runner_kwargs = _get_vllm_runner_params(
|
||||
model, max_model_len, tensor_parallel_size=tensor_parallel_size
|
||||
)
|
||||
vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
|
||||
|
||||
vllm_outputs_no_cache, _ = _get_vLLM_output(
|
||||
vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
vllm_runner_kwargs["enable_prefix_caching"] = True
|
||||
with vllm_runner(**vllm_runner_kwargs) as vllm_model:
|
||||
# Retrieve the default mamba state block size
|
||||
vllm_config = vllm_model.llm.llm_engine.vllm_config
|
||||
mamba_block_size = vllm_config.cache_config.mamba_block_size
|
||||
|
||||
# In case the hybrid model does not have the
|
||||
# "mamba_block_size" assume a fixed constant
|
||||
if mamba_block_size is None:
|
||||
mamba_block_size = 512
|
||||
|
||||
mamba_block_size_multiplier = 10
|
||||
for offsets in [-3, 3, mamba_block_size // 4 + 3, mamba_block_size // 2 - 3]:
|
||||
vllm_runner_kwargs["max_num_batched_tokens"] = (
|
||||
mamba_block_size_multiplier * mamba_block_size - offsets
|
||||
)
|
||||
vllm_outputs_cache_rep, _ = _get_vLLM_output(
|
||||
vllm_runner,
|
||||
vllm_runner_kwargs,
|
||||
generated_prompts,
|
||||
max_tokens,
|
||||
num_logprobs,
|
||||
n_repetitions,
|
||||
)
|
||||
|
||||
# Check alignment of the output logits when using APC
|
||||
for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
|
||||
# In the first repetition, the caches are filled
|
||||
# In the second repetition, these caches are reused
|
||||
|
||||
compare_operator(
|
||||
outputs_0_lst=vllm_outputs_no_cache[0],
|
||||
outputs_1_lst=vllm_outputs_cache_itn,
|
||||
name_0="vllm_no_cache",
|
||||
name_1=f"vllm_cache_it_{r_idx + 1}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [HYBRID_MODELS[0]])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("n_repetitions", [2])
|
||||
# If num_logprobs is set to -1, then the stringent version
|
||||
# of the test is executed using `check_outputs_equal`
|
||||
# instead of `check_logprobs_close`
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1])
|
||||
def test_apc_multiple_prompts_all_cached_outputs(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
n_repetitions: int,
|
||||
num_logprobs: int,
|
||||
tensor_parallel_size: int,
|
||||
) -> None:
|
||||
try:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
compare_operator: Callable = (
|
||||
check_logprobs_close if num_logprobs > 0 else check_outputs_equal # type: ignore
|
||||
)
|
||||
|
||||
# Sample prompts.
|
||||
generated_prompts = [APC_MULTIPLY_BY * prompt for prompt in example_prompts]
|
||||
|
||||
max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
|
||||
vllm_runner_kwargs = _get_vllm_runner_params(
|
||||
model, max_model_len, tensor_parallel_size=tensor_parallel_size
|
||||
)
|
||||
vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
|
||||
# Reduce the effects of batch variance on ROCm since batch invariance is not
|
||||
# yet supported. See: https://github.com/vllm-project/vllm/issues/27433
|
||||
if current_platform.is_rocm():
|
||||
vllm_runner_kwargs["max_num_seqs"] = 4
|
||||
|
||||
vllm_outputs_no_cache, _ = _get_vLLM_output(
|
||||
vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
vllm_runner_kwargs["enable_prefix_caching"] = True
|
||||
vllm_outputs_cache_rep, _ = _get_vLLM_output(
|
||||
vllm_runner,
|
||||
vllm_runner_kwargs,
|
||||
generated_prompts,
|
||||
max_tokens,
|
||||
num_logprobs,
|
||||
n_repetitions,
|
||||
)
|
||||
|
||||
for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
|
||||
# In the first repetition, the caches are filled
|
||||
# In the second repetition, these caches are reused
|
||||
|
||||
compare_operator(
|
||||
outputs_0_lst=vllm_outputs_no_cache[0],
|
||||
outputs_1_lst=vllm_outputs_cache_itn,
|
||||
name_0="vllm_no_cache",
|
||||
name_1=f"vllm_cache_it_{r_idx + 1}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [HYBRID_MODELS[0]])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("n_repetitions", [2])
|
||||
# If num_logprobs is set to -1, then the stringent version
|
||||
# of the test is executed using `check_outputs_equal`
|
||||
# instead of `check_logprobs_close`
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1])
|
||||
def test_apc_multiple_prompts_block_align_alignment(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
n_repetitions: int,
|
||||
num_logprobs: int,
|
||||
tensor_parallel_size: int,
|
||||
) -> None:
|
||||
try:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
compare_operator: Callable = (
|
||||
check_logprobs_close if num_logprobs > 0 else check_outputs_equal # type: ignore
|
||||
)
|
||||
|
||||
# Sample prompts. This custom prompt is used, as it causes the most issues
|
||||
prompt_text = "The president of the United States is "
|
||||
prompt_offsets = [0, 3, 7, 13, 17, 22, 25, 31]
|
||||
generated_prompts = [
|
||||
prompt_text[offset:] * APC_MULTIPLY_BY for offset in prompt_offsets
|
||||
]
|
||||
|
||||
max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
|
||||
vllm_runner_kwargs = _get_vllm_runner_params(
|
||||
model, max_model_len, tensor_parallel_size
|
||||
)
|
||||
vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
|
||||
|
||||
vllm_outputs_no_cache, _ = _get_vLLM_output(
|
||||
vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
vllm_runner_kwargs["enable_prefix_caching"] = True
|
||||
with vllm_runner(**vllm_runner_kwargs) as vllm_model:
|
||||
# Retrieve the default mamba state block size
|
||||
vllm_config = vllm_model.llm.llm_engine.vllm_config
|
||||
mamba_block_size = vllm_config.cache_config.mamba_block_size
|
||||
|
||||
# In case the hybrid model does not have the
|
||||
# "mamba_block_size" assume a fixed constant
|
||||
if mamba_block_size is None:
|
||||
mamba_block_size = 512
|
||||
|
||||
mamba_block_size_multiplier = 10
|
||||
for offsets in [-3, 3, mamba_block_size // 4 + 3, mamba_block_size // 2 - 3]:
|
||||
vllm_runner_kwargs["max_num_batched_tokens"] = (
|
||||
mamba_block_size_multiplier * mamba_block_size - offsets
|
||||
)
|
||||
vllm_outputs_cache_rep, _ = _get_vLLM_output(
|
||||
vllm_runner,
|
||||
vllm_runner_kwargs,
|
||||
generated_prompts,
|
||||
max_tokens,
|
||||
num_logprobs,
|
||||
n_repetitions,
|
||||
)
|
||||
|
||||
# Check alignment of the output logits when using APC
|
||||
for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
|
||||
# In the first repetition, the caches are filled
|
||||
# In the second repetition, these caches are reused
|
||||
|
||||
compare_operator(
|
||||
outputs_0_lst=vllm_outputs_no_cache[0],
|
||||
outputs_1_lst=vllm_outputs_cache_itn,
|
||||
name_0="vllm_no_cache",
|
||||
name_1=f"vllm_cache_it_{r_idx + 1}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [HYBRID_MODELS[0]])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("n_repetitions", [2])
|
||||
# If num_logprobs is set to -1, then the stringent version
|
||||
# of the test is executed using `check_outputs_equal`
|
||||
# instead of `check_logprobs_close`
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1])
|
||||
def test_apc_multiple_prompts_partial_cached_outputs(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
n_repetitions: int,
|
||||
num_logprobs: int,
|
||||
tensor_parallel_size: int,
|
||||
) -> None:
|
||||
try:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
compare_operator: Callable = (
|
||||
check_logprobs_close if num_logprobs > 0 else check_outputs_equal # type: ignore
|
||||
)
|
||||
|
||||
# Sample prompts.
|
||||
generated_prompts = [APC_MULTIPLY_BY * prompt for prompt in example_prompts]
|
||||
|
||||
max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
|
||||
vllm_runner_kwargs = _get_vllm_runner_params(
|
||||
model, max_model_len, tensor_parallel_size=tensor_parallel_size
|
||||
)
|
||||
vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
|
||||
|
||||
vllm_outputs_no_cache, _ = _get_vLLM_output(
|
||||
vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
# Cache only part of all the prompts
|
||||
vllm_runner_kwargs["enable_prefix_caching"] = True
|
||||
with _owned_vllm_runner(vllm_runner, vllm_runner_kwargs) as vllm_model:
|
||||
vllm_outputs_partial_cache, _ = _get_vLLM_output(
|
||||
vllm_runner,
|
||||
vllm_runner_kwargs,
|
||||
generated_prompts[:3],
|
||||
max_tokens,
|
||||
num_logprobs,
|
||||
vllm_model=vllm_model,
|
||||
)
|
||||
|
||||
compare_operator(
|
||||
outputs_0_lst=vllm_outputs_no_cache[0][:3],
|
||||
outputs_1_lst=vllm_outputs_partial_cache[0],
|
||||
name_0="vllm_no_cache",
|
||||
name_1="vllm_partial_cache",
|
||||
)
|
||||
|
||||
vllm_outputs_cache_rep, _ = _get_vLLM_output(
|
||||
vllm_runner,
|
||||
vllm_runner_kwargs,
|
||||
generated_prompts,
|
||||
max_tokens,
|
||||
num_logprobs,
|
||||
n_repetitions,
|
||||
vllm_model=vllm_model,
|
||||
)
|
||||
|
||||
for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
|
||||
# In the first repetition, the caches are filled
|
||||
# In the second repetition, these caches are reused
|
||||
|
||||
compare_operator(
|
||||
outputs_0_lst=vllm_outputs_no_cache[0],
|
||||
outputs_1_lst=vllm_outputs_cache_itn,
|
||||
name_0="vllm_no_cache",
|
||||
name_1=f"vllm_cache_it_{r_idx + 1}",
|
||||
)
|
||||
|
||||
|
||||
# Test that outputs match whether prefix caching is enabled or not for mamba.
|
||||
@pytest.mark.parametrize("model", ["tiiuae/falcon-mamba-7b"])
|
||||
def test_same_mamba_output_apc_on_vs_off(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
) -> None:
|
||||
num_logprobs = 5
|
||||
prompts = [
|
||||
"hello what is one plus one what is one plus one what is one plus one the answer is", # noqa: E501
|
||||
"hello what is one plus one what is one plus one what is one plus one the answer is", # noqa: E501
|
||||
]
|
||||
max_tokens = 20
|
||||
max_model_len = max(len(p) for p in prompts) + max_tokens + 64
|
||||
|
||||
base_kwargs = _get_vllm_runner_params(model, max_model_len)
|
||||
base_kwargs.update(
|
||||
enforce_eager=True, block_size=16, seed=42, gpu_memory_utilization=0.8
|
||||
)
|
||||
|
||||
# No prefix caching
|
||||
kwargs_no_apc = {**base_kwargs, "enable_prefix_caching": False}
|
||||
with vllm_runner(**kwargs_no_apc) as vllm_model:
|
||||
outputs_no_apc, _ = _get_vLLM_output(
|
||||
vllm_runner,
|
||||
kwargs_no_apc,
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
vllm_model=vllm_model,
|
||||
)
|
||||
# With prefix caching
|
||||
kwargs_with_apc = {
|
||||
**base_kwargs,
|
||||
"enable_prefix_caching": True,
|
||||
"mamba_block_size": 16,
|
||||
}
|
||||
with vllm_runner(**kwargs_with_apc) as vllm_model:
|
||||
outputs_with_apc, _ = _get_vLLM_output(
|
||||
vllm_runner,
|
||||
kwargs_with_apc,
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
vllm_model=vllm_model,
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=outputs_no_apc[0],
|
||||
outputs_1_lst=outputs_with_apc[0],
|
||||
name_0="vllm_no_apc",
|
||||
name_1="vllm_with_apc",
|
||||
)
|
||||
|
||||
|
||||
# we have to use a real large model to get reasonable results
|
||||
# the model can't be a hybrid model as we need block_size 16
|
||||
@pytest.mark.parametrize("model", ["tiiuae/falcon-mamba-7b"])
|
||||
def test_apc_common_prefix_same_batch(
|
||||
model: str,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
# Required to put the two requests in the same batch
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
llm = LLM(
|
||||
model=model,
|
||||
enforce_eager=True,
|
||||
block_size=16,
|
||||
mamba_block_size=16,
|
||||
enable_prefix_caching=True,
|
||||
seed=42,
|
||||
attention_backend=ATTN_BACKEND,
|
||||
)
|
||||
prompts = [
|
||||
"hello what is one plus one what is one plus one what is one plus one the answer is", # noqa: E501
|
||||
"hello what is one plus one what is one plus one what is one plus one the answer is", # noqa: E501
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=20)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
for output in outputs:
|
||||
assert "two" in output.outputs[0].text
|
||||
@@ -0,0 +1,352 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import copy
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.tokenizers.mistral import MistralTokenizer
|
||||
from vllm.tool_parsers.mistral_tool_parser import (
|
||||
MistralToolCall,
|
||||
MistralToolParser,
|
||||
)
|
||||
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
MODELS = [
|
||||
"mistralai/Mistral-7B-Instruct-v0.3",
|
||||
]
|
||||
|
||||
MISTRAL_FORMAT_MODELS = [
|
||||
"mistralai/Mistral-7B-Instruct-v0.3",
|
||||
# uses the v3-Tekken tokenizer
|
||||
"mistralai/Ministral-8B-Instruct-2410",
|
||||
# Mistral-Nemo is too big for CI, but passes locally
|
||||
# "mistralai/Mistral-Nemo-Instruct-2407"
|
||||
]
|
||||
|
||||
SAMPLING_PARAMS = SamplingParams(max_tokens=512, temperature=0.0, logprobs=5)
|
||||
SYMBOLIC_LANG_PROMPTS = [
|
||||
"勇敢な船乗りについての詩を書く", # japanese
|
||||
"寫一首關於勇敢的水手的詩", # chinese
|
||||
"ပုံပြင်လေးပြောပြပါ်:\n", # burmese
|
||||
"Repeat the phrase 'URGENCY🌶️':\nURGENCY🌶️\nURGENCY🌶️\n", # see https://github.com/vllm-project/vllm/pull/9625
|
||||
]
|
||||
|
||||
# for function calling
|
||||
TOOLS = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The city to find the weather for, e.g. "
|
||||
"'San Francisco'",
|
||||
},
|
||||
"state": {
|
||||
"type": "string",
|
||||
"description": "the two-letter abbreviation for the state that "
|
||||
"the city is in, e.g. 'CA' which would mean 'California'",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"description": "The unit to fetch the temperature in",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["city", "state", "unit"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "rewrite",
|
||||
"description": "Rewrites text",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"required": [],
|
||||
"properties": {
|
||||
"text": {
|
||||
"type": "string",
|
||||
"description": "The input text to rewrite.",
|
||||
}
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
MSGS = [
|
||||
{"role": "system", "content": "You are an assistant."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Could you please rewrite the below article? \n\n My English needs "
|
||||
"improving, maybe I make errors.",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "bbc5b7ede",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "rewrite",
|
||||
"arguments": '{"text":"My English needs improving, maybe '
|
||||
'I make errors."}',
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": '{"action":"rewrite","outcome":"My English needs improving, maybe '
|
||||
'I make errors."}',
|
||||
"tool_call_id": "bbc5b7ede",
|
||||
"name": "rewrite",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "---\n\nMy English needs improving, maybe I make errors",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"Can you tell me what the temperate will be in Dallas, in fahrenheit?"
|
||||
),
|
||||
},
|
||||
]
|
||||
|
||||
SAMPLE_JSON_SCHEMA = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"age": {"type": "integer"},
|
||||
"skills": {
|
||||
"type": "array",
|
||||
"items": {"type": "string", "maxLength": 10},
|
||||
"minItems": 3,
|
||||
},
|
||||
"work_history": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"company": {"type": "string"},
|
||||
"duration": {"type": "number"},
|
||||
"position": {"type": "string"},
|
||||
},
|
||||
"required": ["company", "position"],
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["name", "age", "skills", "work_history"],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
def test_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
# TODO(sang): Sliding window should be tested separately.
|
||||
with hf_runner(model, dtype=dtype) as hf_model:
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
with vllm_runner(model, dtype=dtype, tokenizer_mode="mistral") as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MISTRAL_FORMAT_MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
def test_mistral_format(
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
tokenizer_mode="mistral",
|
||||
load_format="mistral",
|
||||
config_format="mistral",
|
||||
) as mistral_format_model:
|
||||
mistral_format_outputs = mistral_format_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
tokenizer_mode="hf",
|
||||
load_format="safetensors",
|
||||
config_format="hf",
|
||||
) as hf_format_model:
|
||||
hf_format_outputs = hf_format_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_format_outputs,
|
||||
outputs_1_lst=mistral_format_outputs,
|
||||
name_0="hf",
|
||||
name_1="mistral",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MISTRAL_FORMAT_MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
def test_mistral_symbolic_languages(vllm_runner, model: str, dtype: str) -> None:
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_model_len=8192,
|
||||
tokenizer_mode="mistral",
|
||||
config_format="mistral",
|
||||
load_format="mistral",
|
||||
) as vllm_model:
|
||||
for prompt in SYMBOLIC_LANG_PROMPTS:
|
||||
msg = {"role": "user", "content": prompt}
|
||||
outputs = vllm_model.llm.chat([msg], sampling_params=SAMPLING_PARAMS)
|
||||
assert "�" not in outputs[0].outputs[0].text.strip()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MISTRAL_FORMAT_MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
def test_mistral_function_calling(vllm_runner, model: str, dtype: str) -> None:
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
tokenizer_mode="mistral",
|
||||
config_format="mistral",
|
||||
load_format="mistral",
|
||||
) as vllm_model:
|
||||
msgs = copy.deepcopy(MSGS)
|
||||
outputs = vllm_model.llm.chat(
|
||||
msgs, tools=TOOLS, sampling_params=SAMPLING_PARAMS
|
||||
)
|
||||
|
||||
tokenizer = vllm_model.llm.get_tokenizer()
|
||||
tool_parser = MistralToolParser(tokenizer)
|
||||
|
||||
model_output = outputs[0].outputs[0].text.strip()
|
||||
assert model_output.startswith(tool_parser.bot_token), model_output
|
||||
parsed_message = tool_parser.extract_tool_calls(model_output, None)
|
||||
|
||||
assert parsed_message.tools_called
|
||||
|
||||
assert MistralToolCall.is_valid_id(parsed_message.tool_calls[0].id)
|
||||
assert parsed_message.tool_calls[0].function.name == "get_current_weather"
|
||||
assert (
|
||||
parsed_message.tool_calls[0].function.arguments
|
||||
== '{"city": "Dallas", "state": "TX", "unit": "fahrenheit"}'
|
||||
) # noqa
|
||||
assert parsed_message.content is None
|
||||
|
||||
|
||||
def test_mistral_function_call_nested_json():
|
||||
"""Ensure that the function-name regex captures the entire outermost
|
||||
JSON block, including nested braces."""
|
||||
|
||||
# Create a minimal stub tokenizer that provides the few attributes the
|
||||
# parser accesses (`version` and `get_vocab`).
|
||||
class _StubMistralTokenizer(MistralTokenizer):
|
||||
version = 11 # Satisfy the version check
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def get_vocab():
|
||||
# Provide the special TOOL_CALLS token expected by the parser.
|
||||
return {"[TOOL_CALLS]": 0}
|
||||
|
||||
tokenizer = _StubMistralTokenizer()
|
||||
parser = MistralToolParser(tokenizer)
|
||||
|
||||
# Craft a model output featuring nested JSON inside the arguments.
|
||||
args_dict = {
|
||||
"city": "Dallas",
|
||||
"state": "TX",
|
||||
"unit": "fahrenheit",
|
||||
"sub_dict": {"foo": "bar", "inner": {"x": 1, "y": 2}},
|
||||
}
|
||||
|
||||
model_output = f"{parser.bot_token}get_current_weather{json.dumps(args_dict)}"
|
||||
|
||||
parsed = parser.extract_tool_calls(model_output, None)
|
||||
|
||||
# Assertions: the tool call is detected and the full nested JSON is parsed
|
||||
# without truncation.
|
||||
assert parsed.tools_called
|
||||
|
||||
assert MistralToolCall.is_valid_id(parsed.tool_calls[0].id)
|
||||
assert parsed.tool_calls[0].function.name == "get_current_weather"
|
||||
assert json.loads(parsed.tool_calls[0].function.arguments) == args_dict
|
||||
# No additional content outside the tool call should be returned.
|
||||
assert parsed.content is None
|
||||
|
||||
# multiple calls
|
||||
multiple_args_dict = [
|
||||
{
|
||||
"city": "Dallas",
|
||||
"state": "TX",
|
||||
"unit": "fahrenheit",
|
||||
"sub_dict": {"foo": "bar", "inner": {"x": 1, "y": 2}},
|
||||
},
|
||||
{},
|
||||
{"a": 0},
|
||||
{"a": 1, "b": "c"},
|
||||
]
|
||||
names = ["get_current_weather", "get_current_weather_2", "random", "random_2"]
|
||||
|
||||
model_output = "".join(
|
||||
[
|
||||
f"{parser.bot_token}{name}{json.dumps(args)}"
|
||||
for name, args in zip(names, multiple_args_dict)
|
||||
]
|
||||
)
|
||||
|
||||
parsed = parser.extract_tool_calls(model_output, None)
|
||||
|
||||
# Assertions: the tool call is detected and the full nested JSON is parsed
|
||||
# without truncation.
|
||||
assert parsed.tools_called
|
||||
assert len(parsed.tool_calls) == len(multiple_args_dict)
|
||||
|
||||
for i, tool_call in enumerate(parsed.tool_calls):
|
||||
assert MistralToolCall.is_valid_id(tool_call.id)
|
||||
assert tool_call.function.name == names[i]
|
||||
assert json.loads(tool_call.function.arguments) == multiple_args_dict[i]
|
||||
# No additional content outside the tool call should be returned.
|
||||
assert parsed.content is None
|
||||
@@ -0,0 +1,109 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....utils import large_gpu_test
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
MODELS = [
|
||||
"microsoft/Phi-3.5-MoE-instruct",
|
||||
]
|
||||
|
||||
|
||||
def test_phimoe_routing_function():
|
||||
from vllm.model_executor.models.phimoe import phimoe_routing_function
|
||||
|
||||
test_case = {
|
||||
0: {
|
||||
"hidden_states": torch.tensor(
|
||||
[1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.float32, requires_grad=False
|
||||
).view(4, 2),
|
||||
"gating_output": torch.tensor(
|
||||
[0.1, 0.2, 0.3, 0.4], dtype=torch.float32, requires_grad=False
|
||||
),
|
||||
"topk": 2,
|
||||
"renormalize": False,
|
||||
},
|
||||
1: {
|
||||
"hidden_states": torch.tensor(
|
||||
[1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.float32, requires_grad=False
|
||||
).view(4, 2),
|
||||
"gating_output": torch.tensor(
|
||||
[0.4, 0.2, 0.3, 0.4], dtype=torch.float32, requires_grad=False
|
||||
),
|
||||
"topk": 2,
|
||||
"renormalize": False,
|
||||
},
|
||||
}
|
||||
|
||||
ground_truth = {
|
||||
0: {
|
||||
"topk_weights": torch.tensor(
|
||||
[1.0, 1.0], dtype=torch.float32, requires_grad=False
|
||||
),
|
||||
"topk_ids": torch.tensor([3, 2], dtype=torch.long, requires_grad=False),
|
||||
},
|
||||
1: {
|
||||
"topk_weights": torch.tensor(
|
||||
[0.5, 1.0], dtype=torch.float32, requires_grad=False
|
||||
),
|
||||
"topk_ids": torch.tensor([0, 3], dtype=torch.long, requires_grad=False),
|
||||
},
|
||||
}
|
||||
|
||||
for test_id in test_case:
|
||||
topk_weights, topk_ids = phimoe_routing_function(**test_case[test_id])
|
||||
assert torch.allclose(topk_weights, ground_truth[test_id]["topk_weights"])
|
||||
assert torch.equal(topk_ids, ground_truth[test_id]["topk_ids"])
|
||||
|
||||
|
||||
# There is a known issue that triggers `AttributeError: 'DynamicCache'
|
||||
# object has no attribute 'seen_tokens'` when running:
|
||||
# `tests/models/language/generation/test_phimoe.py::test_models
|
||||
# [5-64-bfloat16-microsoft/Phi-3.5-MoE-instruct]`
|
||||
# This issue is being investigated and tracked in:
|
||||
# https://huggingface.co/microsoft/Phi-3.5-MoE-instruct/discussions/58
|
||||
# It is platform-agnostic. Therefore, we skip this test on all platforms for now.
|
||||
@pytest.mark.skip(
|
||||
reason="Skipping due to known issue: "
|
||||
"'DynamicCache' object has no attribute 'seen_tokens'. See: "
|
||||
"https://huggingface.co/microsoft/Phi-3.5-MoE-instruct/discussions/58 "
|
||||
"for details.",
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
condition=current_platform.is_cpu(),
|
||||
reason="This test takes a lot time to run on CPU, "
|
||||
"and vllm CI's disk space is not enough for this model.",
|
||||
)
|
||||
@large_gpu_test(min_gb=80)
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
def test_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
with hf_runner(model, dtype=dtype) as hf_model:
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
|
||||
with vllm_runner(model, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
)
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
@@ -0,0 +1,128 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from https://huggingface.co/docs/transformers/perplexity
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
from datasets import load_dataset
|
||||
|
||||
import tests.ci_envs as ci_envs
|
||||
from tests.models.utils import (
|
||||
GenerateModelInfo,
|
||||
TokensTextLogprobsPromptLogprobs,
|
||||
get_vllm_extra_kwargs,
|
||||
)
|
||||
from vllm.logprobs import Logprob
|
||||
|
||||
# See #24485
|
||||
PPL_TOL = 0.01
|
||||
MAX_LENGTH = 1024
|
||||
|
||||
|
||||
@torch.inference_mode
|
||||
def wikitext_ppl_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model_info: GenerateModelInfo,
|
||||
max_length=MAX_LENGTH,
|
||||
vllm_extra_kwargs=None,
|
||||
atol=PPL_TOL,
|
||||
):
|
||||
vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
|
||||
|
||||
dataset = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="test")
|
||||
|
||||
with vllm_runner(
|
||||
model_info.name,
|
||||
gpu_memory_utilization=0.7,
|
||||
max_model_len=max_length,
|
||||
max_num_seqs=1,
|
||||
**vllm_extra_kwargs,
|
||||
) as vllm_model:
|
||||
# Use max_num_seqs=1 to avoid OOM,
|
||||
# and avoid batch different requests together.
|
||||
|
||||
model_config = vllm_model.llm.llm_engine.model_config
|
||||
|
||||
# Confirm whether vllm is using the correct architecture
|
||||
if model_info.architecture:
|
||||
assert model_info.architecture in model_config.architectures
|
||||
|
||||
max_length = min(model_config.max_model_len - 1, max_length)
|
||||
stride = max_length
|
||||
|
||||
tokenizer = vllm_model.llm.get_tokenizer()
|
||||
tokens = tokenizer.encode("\n\n".join(dataset["text"]))
|
||||
n_tokens = len(tokens)
|
||||
|
||||
chunks = []
|
||||
for begin_loc in range(0, n_tokens, stride):
|
||||
end_loc = min(begin_loc + max_length, n_tokens)
|
||||
chunks.append(tokens[begin_loc:end_loc])
|
||||
|
||||
outputs = vllm_model.generate_greedy_logprobs(
|
||||
prompts=chunks,
|
||||
max_tokens=1,
|
||||
num_logprobs=None,
|
||||
num_prompt_logprobs=0,
|
||||
use_tqdm=False,
|
||||
)
|
||||
nll_sum = torch.tensor(0.0, dtype=torch.float32, device="cpu")
|
||||
n_tokens = 0
|
||||
for output in outputs:
|
||||
output = cast(TokensTextLogprobsPromptLogprobs, output)
|
||||
token_datas = cast(list[dict[int, Logprob] | None], output[3])
|
||||
|
||||
assert token_datas[0] is None
|
||||
token_log_probs = []
|
||||
for token_data in token_datas[1:]:
|
||||
assert token_data is not None
|
||||
assert len(token_data) == 1
|
||||
token_log_prob = list(token_data.values())[0].logprob
|
||||
token_log_probs.append(token_log_prob)
|
||||
|
||||
neg_log_likelihood = -torch.tensor(
|
||||
token_log_probs, dtype=torch.float32, device="cpu"
|
||||
).sum()
|
||||
nll_sum += neg_log_likelihood
|
||||
n_tokens += len(token_log_probs)
|
||||
vllm_ppl = float(torch.exp(nll_sum / n_tokens))
|
||||
vllm_dtype = model_config.dtype
|
||||
head_dtype = model_config.head_dtype
|
||||
|
||||
# Accelerate ppl test by setting Transformers ppl score to a constant
|
||||
if model_info.hf_ppl is None:
|
||||
with hf_runner(
|
||||
model_info.name,
|
||||
dtype=ci_envs.VLLM_CI_HF_DTYPE or model_info.hf_dtype,
|
||||
) as hf_model:
|
||||
nll_sum = torch.tensor(0.0, dtype=torch.float32, device="cpu")
|
||||
n_tokens = 0
|
||||
for chunk in chunks:
|
||||
inputs = hf_model.wrap_device({"input_ids": torch.tensor([chunk])})
|
||||
input_ids = inputs["input_ids"]
|
||||
outputs = hf_model.model(input_ids, labels=input_ids)
|
||||
neg_log_likelihood = outputs.loss
|
||||
|
||||
neg_log_likelihood = neg_log_likelihood.to(torch.float32).cpu()
|
||||
|
||||
num_loss_tokens = len(chunk) - 1
|
||||
nll_sum += neg_log_likelihood * num_loss_tokens
|
||||
n_tokens += num_loss_tokens
|
||||
|
||||
hf_ppl = float(torch.exp(nll_sum / n_tokens))
|
||||
hf_dtype = next(hf_model.model.parameters()).dtype
|
||||
else:
|
||||
hf_ppl = model_info.hf_ppl
|
||||
hf_dtype = "Constant"
|
||||
|
||||
differ = (vllm_ppl - hf_ppl) / hf_ppl
|
||||
print("Model:", model_info.name)
|
||||
print("VLLM:", f"dtype:{vllm_dtype}", f"head_dtype:{head_dtype}", vllm_ppl)
|
||||
print("Transformers:", hf_dtype, hf_ppl)
|
||||
print("Difference (%):", differ * 100)
|
||||
|
||||
# PPL the smaller, the better
|
||||
# We are not concerned that the vllm PPL is less than Transformers,
|
||||
# so we only perform one-sided testing.
|
||||
assert differ < atol
|
||||
@@ -0,0 +1,18 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
from tests.models.utils import GenerateModelInfo
|
||||
|
||||
from .ppl_utils import wikitext_ppl_test
|
||||
|
||||
MODELS = [
|
||||
GenerateModelInfo("google/gemma-2b", hf_ppl=21.48524284362793),
|
||||
GenerateModelInfo("google/gemma-2-2b", hf_ppl=102.59290313720703),
|
||||
GenerateModelInfo("google/gemma-3-4b-it", hf_ppl=27.79648208618164),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_ppl(hf_runner, vllm_runner, model_info: GenerateModelInfo):
|
||||
wikitext_ppl_test(hf_runner, vllm_runner, model_info)
|
||||
@@ -0,0 +1,14 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
from tests.models.utils import GenerateModelInfo
|
||||
|
||||
from .ppl_utils import wikitext_ppl_test
|
||||
|
||||
MODELS = [GenerateModelInfo("openai-community/gpt2-large", hf_ppl=19.457056045532227)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_ppl(hf_runner, vllm_runner, model_info: GenerateModelInfo):
|
||||
wikitext_ppl_test(hf_runner, vllm_runner, model_info)
|
||||
@@ -0,0 +1,27 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.utils import GenerateModelInfo
|
||||
|
||||
from .ppl_utils import wikitext_ppl_test
|
||||
|
||||
MODELS = [
|
||||
# for Qwen3
|
||||
GenerateModelInfo("Qwen/Qwen3-0.6B", hf_ppl=23.864173889160156),
|
||||
GenerateModelInfo("Qwen/Qwen3-0.6B-FP8", hf_ppl=24.313045501708984),
|
||||
# for Qwen3.5
|
||||
GenerateModelInfo("Qwen/Qwen3.5-0.8B", hf_ppl=19.38858413696289),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_ppl(hf_runner, vllm_runner, model_info: GenerateModelInfo):
|
||||
vllm_extra_kwargs = {}
|
||||
if model_info.name == "Qwen/Qwen3.5-0.8B":
|
||||
vllm_extra_kwargs["language_model_only"] = True
|
||||
|
||||
wikitext_ppl_test(
|
||||
hf_runner, vllm_runner, model_info, vllm_extra_kwargs=vllm_extra_kwargs
|
||||
)
|
||||
@@ -0,0 +1,29 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Pytest configuration for vLLM language generation tests."""
|
||||
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def pytest_sessionstart(session):
|
||||
"""Configure ROCm-specific settings before test session starts."""
|
||||
if not current_platform.is_rocm():
|
||||
return
|
||||
|
||||
# Disable Flash/MemEfficient SDP on ROCm to avoid HF Transformers
|
||||
# accuracy issues: https://github.com/vllm-project/vllm/issues/30167
|
||||
# TODO: Remove once ROCm SDP accuracy issues are resolved on HuggingFace
|
||||
torch.backends.cuda.enable_flash_sdp(False)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(False)
|
||||
torch.backends.cuda.enable_math_sdp(True)
|
||||
torch.set_float32_matmul_precision("high")
|
||||
warnings.warn(
|
||||
"ROCm: Disabled flash_sdp and mem_efficient_sdp, enabled math_sdp "
|
||||
"to avoid HuggingFace Transformers accuracy issues",
|
||||
UserWarning,
|
||||
stacklevel=1,
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from collections.abc import Sequence
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
|
||||
from tests.conftest import HfRunner
|
||||
from tests.models.utils import EmbedModelInfo, check_embeddings_close, matryoshka_fy
|
||||
|
||||
|
||||
def run_embedding_correctness_test(
|
||||
hf_model: "HfRunner",
|
||||
inputs: list[str],
|
||||
vllm_outputs: Sequence[list[float]],
|
||||
dimensions: int | None = None,
|
||||
):
|
||||
hf_outputs = hf_model.encode(inputs)
|
||||
if dimensions:
|
||||
hf_outputs = matryoshka_fy(hf_outputs, dimensions)
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=hf_outputs,
|
||||
embeddings_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
tol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
def correctness_test_embed_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model_info: EmbedModelInfo,
|
||||
example_prompts,
|
||||
vllm_extra_kwargs=None,
|
||||
hf_model_callback=None,
|
||||
):
|
||||
pytest.skip("Debug only, ci prefers to use mteb test.")
|
||||
|
||||
# The example_prompts has ending "\n", for example:
|
||||
# "Write a short story about a robot that dreams for the first time.\n"
|
||||
# sentence_transformers will strip the input texts, see:
|
||||
# https://github.com/UKPLab/sentence-transformers/blob/v3.1.1/sentence_transformers/models/Transformer.py#L159
|
||||
# This makes the input_ids different between hf_model and vllm_model.
|
||||
# So we need to strip the input texts to avoid test failing.
|
||||
example_prompts = [str(s).strip() for s in example_prompts]
|
||||
|
||||
vllm_extra_kwargs = vllm_extra_kwargs or {}
|
||||
vllm_extra_kwargs["dtype"] = model_info.dtype
|
||||
|
||||
if model_info.hf_overrides is not None:
|
||||
vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides
|
||||
|
||||
with vllm_runner(
|
||||
model_info.name, runner="pooling", max_model_len=None, **vllm_extra_kwargs
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.embed(example_prompts)
|
||||
|
||||
with hf_runner(
|
||||
model_info.name,
|
||||
dtype=model_info.hf_dtype,
|
||||
is_sentence_transformer=True,
|
||||
) as hf_model:
|
||||
if hf_model_callback is not None:
|
||||
hf_model_callback(hf_model)
|
||||
|
||||
run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs)
|
||||
|
||||
|
||||
async def run_client_embeddings(
|
||||
client: openai.AsyncOpenAI,
|
||||
model_name: str,
|
||||
queries: list[str],
|
||||
instruction: str = "",
|
||||
) -> list[list[float]]:
|
||||
outputs = await client.embeddings.create(
|
||||
model=model_name,
|
||||
input=[instruction + q for q in queries],
|
||||
)
|
||||
return [data.embedding for data in outputs.data]
|
||||
@@ -0,0 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoModel
|
||||
|
||||
from tests.models.utils import check_embeddings_close
|
||||
from vllm import TokensPrompt
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["Qwen/Qwen3-Embedding-0.6B"],
|
||||
)
|
||||
@torch.inference_mode
|
||||
def test_embed_models(hf_runner, vllm_runner, model: str):
|
||||
chunk_size = 10
|
||||
n_prompt_tokens = [55, 56, 57]
|
||||
token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
max_model_len=128,
|
||||
max_num_batched_tokens=chunk_size,
|
||||
enforce_eager=True,
|
||||
# `enable_chunked_prefill`: Set to `False` instead of `None` in VllmRunner
|
||||
enable_chunked_prefill=True,
|
||||
enable_prefix_caching=True,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_embed(
|
||||
[TokensPrompt(prompt_token_ids=t) for t in token_prompts],
|
||||
)
|
||||
|
||||
with hf_runner(
|
||||
model,
|
||||
auto_cls=AutoModel,
|
||||
) as hf_model:
|
||||
hf_outputs = []
|
||||
for token_prompt in token_prompts:
|
||||
inputs = hf_model.wrap_device({"input_ids": torch.tensor([token_prompt])})
|
||||
input_ids = inputs["input_ids"]
|
||||
output = hf_model.model(input_ids)
|
||||
hf_outputs.append(output.last_hidden_state.cpu().float()[0])
|
||||
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=hf_output,
|
||||
embeddings_1_lst=vllm_output,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
tol=1e-2,
|
||||
)
|
||||
@@ -0,0 +1,113 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoModelForSequenceClassification
|
||||
|
||||
from tests.models.language.pooling.embed_utils import run_embedding_correctness_test
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["jason9693/Qwen2.5-1.5B-apeach"],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_classify_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
# example_prompts is too short for testing prefix_caching
|
||||
example_prompts = [s * 10 for s in example_prompts]
|
||||
|
||||
with vllm_runner(
|
||||
model, max_model_len=512, dtype=dtype, enable_prefix_caching=True
|
||||
) as vllm_model:
|
||||
vllm_config = vllm_model.llm.llm_engine.vllm_config
|
||||
cache_config = vllm_config.cache_config
|
||||
assert cache_config.enable_prefix_caching
|
||||
|
||||
# First Run
|
||||
vllm_model.classify(example_prompts)
|
||||
|
||||
# assert prefix_caching works
|
||||
pooling_outputs = vllm_model.llm.encode(
|
||||
example_prompts, pooling_task="classify"
|
||||
)
|
||||
for output in pooling_outputs:
|
||||
assert output.num_cached_tokens > 0
|
||||
vllm_outputs = [req_output.outputs.data for req_output in pooling_outputs]
|
||||
|
||||
with hf_runner(
|
||||
model, dtype=dtype, auto_cls=AutoModelForSequenceClassification
|
||||
) as hf_model:
|
||||
hf_outputs = hf_model.classify(example_prompts)
|
||||
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
hf_output = torch.tensor(hf_output)
|
||||
vllm_output = torch.tensor(vllm_output)
|
||||
|
||||
assert torch.allclose(
|
||||
hf_output, vllm_output, 1e-3 if dtype == "float" else 1e-2
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["Qwen/Qwen3-Embedding-0.6B"],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_embed_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
):
|
||||
# example_prompts is too short for testing prefix_caching
|
||||
example_prompts = [str(s).strip() * 10 for s in example_prompts]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
max_model_len=None,
|
||||
enable_prefix_caching=True,
|
||||
) as vllm_model:
|
||||
vllm_config = vllm_model.llm.llm_engine.vllm_config
|
||||
cache_config = vllm_config.cache_config
|
||||
assert cache_config.enable_prefix_caching
|
||||
|
||||
# First Run
|
||||
vllm_model.embed(example_prompts)
|
||||
|
||||
# assert prefix_caching works
|
||||
pooling_outputs = vllm_model.llm.encode(example_prompts, pooling_task="embed")
|
||||
for output in pooling_outputs:
|
||||
assert output.num_cached_tokens > 0
|
||||
vllm_outputs = [req_output.outputs.data for req_output in pooling_outputs]
|
||||
|
||||
with hf_runner(
|
||||
model,
|
||||
is_sentence_transformer=True,
|
||||
) as hf_model:
|
||||
run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"intfloat/e5-small",
|
||||
"Alibaba-NLP/gte-Qwen2-1.5B-instruct", # is_causal == False
|
||||
"papluca/xlm-roberta-base-language-detection",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_non_causal_models(
|
||||
hf_runner, vllm_runner, example_prompts, model: str, dtype: str
|
||||
) -> None:
|
||||
with vllm_runner(model, max_model_len=512, dtype=dtype) as vllm_model:
|
||||
vllm_config = vllm_model.llm.llm_engine.vllm_config
|
||||
cache_config = vllm_config.cache_config
|
||||
assert not cache_config.enable_prefix_caching
|
||||
@@ -0,0 +1,212 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
from .embed_utils import run_client_embeddings
|
||||
|
||||
MODEL_NAME = "BAAI/bge-m3"
|
||||
MAX_MODEL_LEN = 512
|
||||
|
||||
|
||||
# Example from https://huggingface.co/BAAI/bge-m3
|
||||
sentences_1 = ["What is BGE M3?", "Definition of BM25"]
|
||||
sentences_2 = [
|
||||
"BGE M3 is an embedding model supporting dense retrieval, "
|
||||
"lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set "
|
||||
"of documents based on the query terms appearing in each document",
|
||||
]
|
||||
|
||||
similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]]
|
||||
lexical_score_reference = [0.19554901123046875, 0.0]
|
||||
colbert_score_reference = [0.7797, 0.4620]
|
||||
SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=SUPPORTED_TASKS)
|
||||
def pooling_task(request):
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(pooling_task):
|
||||
args = [
|
||||
"--max-model-len",
|
||||
str(MAX_MODEL_LEN),
|
||||
"--hf-overrides",
|
||||
'{"architectures": ["BgeM3EmbeddingModel"]}',
|
||||
"--pooler-config.task",
|
||||
pooling_task,
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_embedding(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "embed":
|
||||
with pytest.raises(openai.InternalServerError):
|
||||
await run_client_embeddings(
|
||||
client,
|
||||
MODEL_NAME,
|
||||
sentences_1,
|
||||
)
|
||||
return
|
||||
|
||||
embeddings_list_1 = await run_client_embeddings(
|
||||
client,
|
||||
MODEL_NAME,
|
||||
sentences_1,
|
||||
)
|
||||
embeddings_list_2 = await run_client_embeddings(
|
||||
client,
|
||||
MODEL_NAME,
|
||||
sentences_2,
|
||||
)
|
||||
|
||||
embeddings_1 = torch.tensor(embeddings_list_1)
|
||||
embeddings_2 = torch.tensor(embeddings_list_2)
|
||||
similarity = embeddings_1 @ embeddings_2.T
|
||||
|
||||
# reference values from BAAI/bge-m3 documentation
|
||||
reference = torch.tensor(similarity_reference)
|
||||
|
||||
assert torch.allclose(similarity, reference, rtol=0.01)
|
||||
|
||||
|
||||
async def tokenize(client: openai.AsyncOpenAI, sentences: list[str]) -> list[list[int]]:
|
||||
futures = []
|
||||
for sentence in sentences:
|
||||
futures.append(
|
||||
client.post(
|
||||
"../tokenize",
|
||||
body={"model": MODEL_NAME, "prompt": sentence},
|
||||
cast_to=httpx.Response,
|
||||
)
|
||||
)
|
||||
return [(await future).json()["tokens"] for future in futures]
|
||||
|
||||
|
||||
async def sparse_embeddings(
|
||||
client: openai.AsyncOpenAI, sentences: list[str]
|
||||
) -> list[dict[int, float]]:
|
||||
all_tokens = await tokenize(client, sentences)
|
||||
result = await client.post(
|
||||
"../pooling",
|
||||
body={"model": MODEL_NAME, "input": sentences, "task": "token_classify"},
|
||||
cast_to=httpx.Response,
|
||||
)
|
||||
all_embeddings = [data["data"] for data in result.json()["data"]]
|
||||
|
||||
ret = []
|
||||
|
||||
for sent_tokens, sent_emb in zip(all_tokens, all_embeddings):
|
||||
token_embs = dict[int, float]()
|
||||
if sent_tokens[0] == 0:
|
||||
sent_tokens = sent_tokens[1:]
|
||||
for token, val in zip(sent_tokens, sent_emb):
|
||||
token_embs[token] = max(val, token_embs.get(token, 0.0))
|
||||
ret.append(token_embs)
|
||||
return ret
|
||||
|
||||
|
||||
# Based on https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L129
|
||||
def compute_lexical_matching_score(
|
||||
lw1: dict[int, float], lw2: dict[int, float]
|
||||
) -> float:
|
||||
scores = 0.0
|
||||
for token, weight in lw1.items():
|
||||
if token in lw2:
|
||||
scores += weight * lw2[token]
|
||||
return scores
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "token_classify":
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
await sparse_embeddings(client, sentences_1)
|
||||
return
|
||||
|
||||
embeddings_1 = await sparse_embeddings(client, sentences_1)
|
||||
embeddings_2 = await sparse_embeddings(client, sentences_2)
|
||||
|
||||
lexical_scores_1_0_x_2_0 = compute_lexical_matching_score(
|
||||
embeddings_1[0], embeddings_2[0]
|
||||
)
|
||||
assert lexical_scores_1_0_x_2_0 == pytest.approx(
|
||||
lexical_score_reference[0], rel=0.01
|
||||
)
|
||||
|
||||
lexical_scores_1_0_x_1_1 = compute_lexical_matching_score(
|
||||
embeddings_1[0], embeddings_1[1]
|
||||
)
|
||||
assert lexical_scores_1_0_x_1_1 == pytest.approx(
|
||||
lexical_score_reference[1], rel=0.01
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task):
|
||||
if pooling_task != "token_classify":
|
||||
return
|
||||
|
||||
client = server.get_async_client()
|
||||
embeddings = await sparse_embeddings(client, ["Hi"])
|
||||
assert len(embeddings) == 1
|
||||
assert 2673 in embeddings[0]
|
||||
assert embeddings[0][2673] == pytest.approx(0.26710861921310425, rel=0.01)
|
||||
|
||||
|
||||
# https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L163
|
||||
def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
|
||||
token_scores = torch.einsum("in,jn->ij", q_reps, p_reps)
|
||||
scores, _ = token_scores.max(-1)
|
||||
scores = torch.sum(scores) / q_reps.size(0)
|
||||
return scores
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_multi_vector(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "token_embed":
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
await client.post(
|
||||
"../pooling",
|
||||
body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
|
||||
cast_to=httpx.Response,
|
||||
)
|
||||
return
|
||||
|
||||
result_1 = await client.post(
|
||||
"../pooling",
|
||||
body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
|
||||
cast_to=httpx.Response,
|
||||
)
|
||||
embeddings_1 = [torch.tensor(data["data"]) for data in result_1.json()["data"]]
|
||||
|
||||
result_2 = await client.post(
|
||||
"../pooling",
|
||||
body={"model": MODEL_NAME, "input": sentences_2, "task": "token_embed"},
|
||||
cast_to=httpx.Response,
|
||||
)
|
||||
embeddings_2 = [torch.tensor(data["data"]) for data in result_2.json()["data"]]
|
||||
|
||||
colbert_score_1_0_x_2_0 = colbert_score(embeddings_1[0], embeddings_2[0])
|
||||
assert colbert_score_1_0_x_2_0 == pytest.approx(
|
||||
colbert_score_reference[0], rel=0.01
|
||||
)
|
||||
colbert_score_1_0_x_2_1 = colbert_score(embeddings_1[0], embeddings_2[1])
|
||||
assert colbert_score_1_0_x_2_1 == pytest.approx(
|
||||
colbert_score_reference[1], rel=0.01
|
||||
)
|
||||
@@ -0,0 +1,51 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoModelForSequenceClassification
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
pytest.param(
|
||||
"jason9693/Qwen2.5-1.5B-apeach",
|
||||
marks=[
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.cpu_model,
|
||||
pytest.mark.slow_test,
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"] if current_platform.is_rocm() else ["float"])
|
||||
def test_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
with vllm_runner(model, max_model_len=512, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.classify(example_prompts)
|
||||
|
||||
with hf_runner(
|
||||
model, dtype=dtype, auto_cls=AutoModelForSequenceClassification
|
||||
) as hf_model:
|
||||
hf_outputs = hf_model.classify(example_prompts)
|
||||
|
||||
# check logits difference
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
hf_output = torch.tensor(hf_output)
|
||||
vllm_output = torch.tensor(vllm_output)
|
||||
|
||||
# the tolerance value of 1e-2 is selected based on the
|
||||
# half datatype tests in
|
||||
# tests/models/language/pooling/test_embedding.py
|
||||
assert torch.allclose(
|
||||
hf_output,
|
||||
vllm_output,
|
||||
rtol=2e-3 if dtype == "float" else 1e-2,
|
||||
)
|
||||
@@ -0,0 +1,420 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for ColBERT late interaction scoring.
|
||||
|
||||
Tests are parametrized across multiple ColBERT backbones to ensure the
|
||||
generic ColBERT support works with different encoder architectures.
|
||||
"""
|
||||
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.utils import wait_for_rocm_memory_to_settle
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Model definitions: (model_name, colbert_dim, extra vllm_runner kwargs)
|
||||
# -----------------------------------------------------------------------
|
||||
COLBERT_MODELS = {
|
||||
"bert": {
|
||||
"model": "answerdotai/answerai-colbert-small-v1",
|
||||
"colbert_dim": 96,
|
||||
"max_model_len": 512,
|
||||
"extra_kwargs": {},
|
||||
"hf_comparison": {
|
||||
"weights_file": "model.safetensors",
|
||||
"weights_key": "linear.weight",
|
||||
"trust_remote_code": False,
|
||||
"model_cls": "BertModel",
|
||||
},
|
||||
},
|
||||
"modernbert": {
|
||||
"model": "lightonai/GTE-ModernColBERT-v1",
|
||||
"colbert_dim": 128,
|
||||
"max_model_len": 299,
|
||||
"extra_kwargs": {
|
||||
"hf_overrides": {
|
||||
"architectures": ["ColBERTModernBertModel"],
|
||||
},
|
||||
},
|
||||
"hf_comparison": {
|
||||
"weights_file": "1_Dense/model.safetensors",
|
||||
"weights_key": "linear.weight",
|
||||
"trust_remote_code": False,
|
||||
"model_cls": "AutoModel",
|
||||
},
|
||||
},
|
||||
"jina": {
|
||||
"model": "jinaai/jina-colbert-v2",
|
||||
"colbert_dim": 128,
|
||||
"max_model_len": 8192,
|
||||
"extra_kwargs": {
|
||||
"hf_overrides": {
|
||||
"architectures": ["ColBERTJinaRobertaModel"],
|
||||
},
|
||||
},
|
||||
"hf_comparison": {
|
||||
"weights_file": "model.safetensors",
|
||||
"weights_key": "linear.weight",
|
||||
"trust_remote_code": True,
|
||||
"model_cls": "AutoModel",
|
||||
},
|
||||
},
|
||||
"lfm2": {
|
||||
"model": "LiquidAI/LFM2-ColBERT-350M",
|
||||
"colbert_dim": 128,
|
||||
"max_model_len": 511,
|
||||
"extra_kwargs": {
|
||||
"hf_overrides": {
|
||||
"architectures": ["ColBERTLfm2Model"],
|
||||
},
|
||||
},
|
||||
"hf_comparison": {
|
||||
"weights_file": "1_Dense/model.safetensors",
|
||||
"weights_key": "linear.weight",
|
||||
"trust_remote_code": False,
|
||||
"model_cls": "AutoModel",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
TEXTS_1 = [
|
||||
"What is the capital of France?",
|
||||
"What is the capital of Germany?",
|
||||
]
|
||||
|
||||
TEXTS_2 = [
|
||||
"The capital of France is Paris.",
|
||||
"The capital of Germany is Berlin.",
|
||||
]
|
||||
|
||||
DTYPE = "half"
|
||||
|
||||
|
||||
def _load_hf_model(model_name: str, hf_spec: dict, device: torch.device):
|
||||
"""Load HF model on the given device with a compatible attention impl."""
|
||||
from transformers import AutoModel, BertModel
|
||||
|
||||
cls = BertModel if hf_spec["model_cls"] == "BertModel" else AutoModel
|
||||
trust = hf_spec.get("trust_remote_code", False)
|
||||
|
||||
# Flash / Triton kernels require GPU tensors; fall back to eager on CPU.
|
||||
extra = {}
|
||||
if device.type == "cpu":
|
||||
extra["attn_implementation"] = "eager"
|
||||
|
||||
model = cls.from_pretrained(
|
||||
model_name,
|
||||
trust_remote_code=trust,
|
||||
**extra,
|
||||
).to(device)
|
||||
model.eval()
|
||||
|
||||
# Transformers 5.0 weight materialization can clear non-persistent
|
||||
# buffers (e.g. rotary inv_freq) that were registered with
|
||||
# persistent=False. Re-compute them so the model produces valid output.
|
||||
for mod in model.modules():
|
||||
if hasattr(mod, "_compute_inv_freq") and hasattr(mod, "inv_freq"):
|
||||
mod.inv_freq = mod._compute_inv_freq(device=device)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device):
|
||||
"""Download and return the ColBERT linear projection weight."""
|
||||
from huggingface_hub import hf_hub_download
|
||||
from safetensors.torch import load_file
|
||||
|
||||
path = hf_hub_download(model_name, filename=hf_spec["weights_file"])
|
||||
weights = load_file(path)
|
||||
return weights[hf_spec["weights_key"]].to(device)
|
||||
|
||||
|
||||
def _compute_hf_colbert_embeddings(model, tokenizer, linear_weight, texts, device):
|
||||
"""Run HF model + projection and return L2-normalised token embeddings."""
|
||||
import torch.nn.functional as F
|
||||
|
||||
embeddings = []
|
||||
for text in texts:
|
||||
inputs = tokenizer(text, return_tensors="pt").to(device)
|
||||
with torch.no_grad():
|
||||
hidden = model(**inputs).last_hidden_state.float()
|
||||
projected = F.linear(hidden, linear_weight.float())
|
||||
normalised = F.normalize(projected, p=2, dim=-1)
|
||||
embeddings.append(normalised.squeeze(0).cpu())
|
||||
return embeddings
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _hf_colbert_model(model_name: str, hf_spec: dict, device: torch.device):
|
||||
"""Load the HF backbone + ColBERT projection, freeing the GPU on exit.
|
||||
|
||||
These live outside any runner context, so without explicit cleanup ROCm
|
||||
keeps the VRAM resident and the next backend parametrization (or test)
|
||||
OOMs on startup.
|
||||
"""
|
||||
hf_model = _load_hf_model(model_name, hf_spec, device)
|
||||
linear_weight = _load_projection_weight(model_name, hf_spec, device)
|
||||
try:
|
||||
yield hf_model, linear_weight
|
||||
finally:
|
||||
del hf_model, linear_weight
|
||||
cleanup_dist_env_and_memory()
|
||||
wait_for_rocm_memory_to_settle()
|
||||
|
||||
|
||||
def _assert_embeddings_close(vllm_outputs, hf_embeddings):
|
||||
"""Assert that vLLM and HuggingFace embeddings match."""
|
||||
for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)):
|
||||
vllm_emb = torch.as_tensor(vllm_out).float()
|
||||
|
||||
assert hf_emb.shape == vllm_emb.shape, (
|
||||
f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}"
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
vllm_emb,
|
||||
hf_emb,
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
msg=f"Embedding mismatch for text {i}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="module")
|
||||
def colbert_spec(request):
|
||||
"""Return the model spec dict for the current parametrization."""
|
||||
return COLBERT_MODELS[request.param]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def colbert_model_name(colbert_spec):
|
||||
return colbert_spec["model"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def colbert_dim(colbert_spec):
|
||||
return colbert_spec["colbert_dim"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def colbert_max_model_len(colbert_spec):
|
||||
return colbert_spec["max_model_len"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def colbert_extra_kwargs(colbert_spec):
|
||||
return colbert_spec["extra_kwargs"]
|
||||
|
||||
|
||||
def test_colbert_token_embed(
|
||||
vllm_runner,
|
||||
colbert_model_name,
|
||||
colbert_dim,
|
||||
colbert_max_model_len,
|
||||
colbert_extra_kwargs,
|
||||
):
|
||||
"""Test that ColBERT model produces token embeddings."""
|
||||
with vllm_runner(
|
||||
colbert_model_name,
|
||||
runner="pooling",
|
||||
dtype=DTYPE,
|
||||
max_model_len=colbert_max_model_len,
|
||||
enforce_eager=True,
|
||||
**colbert_extra_kwargs,
|
||||
) as vllm_model:
|
||||
outputs = vllm_model.token_embed([TEXTS_1[0]])
|
||||
|
||||
assert len(outputs) == 1
|
||||
emb = torch.as_tensor(outputs[0])
|
||||
assert emb.dim() == 2
|
||||
assert emb.shape[1] == colbert_dim
|
||||
assert emb.shape[0] > 1
|
||||
|
||||
|
||||
def test_colbert_late_interaction_1_to_1(
|
||||
vllm_runner,
|
||||
colbert_model_name,
|
||||
colbert_max_model_len,
|
||||
colbert_extra_kwargs,
|
||||
):
|
||||
"""Test ColBERT late interaction scoring with 1:1 query-document pair."""
|
||||
with vllm_runner(
|
||||
colbert_model_name,
|
||||
runner="pooling",
|
||||
dtype=DTYPE,
|
||||
max_model_len=colbert_max_model_len,
|
||||
enforce_eager=True,
|
||||
**colbert_extra_kwargs,
|
||||
) as vllm_model:
|
||||
q_outputs = vllm_model.token_embed([TEXTS_1[0]])
|
||||
d_outputs = vllm_model.token_embed([TEXTS_2[0]])
|
||||
|
||||
q_emb = torch.as_tensor(q_outputs[0])
|
||||
d_emb = torch.as_tensor(d_outputs[0])
|
||||
|
||||
manual_score = compute_maxsim_score(q_emb, d_emb).item()
|
||||
|
||||
vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2[0])
|
||||
|
||||
assert len(vllm_scores) == 1
|
||||
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
|
||||
|
||||
|
||||
def test_colbert_late_interaction_1_to_N(
|
||||
vllm_runner,
|
||||
colbert_model_name,
|
||||
colbert_max_model_len,
|
||||
colbert_extra_kwargs,
|
||||
):
|
||||
"""Test ColBERT late interaction scoring with 1:N query-documents."""
|
||||
with vllm_runner(
|
||||
colbert_model_name,
|
||||
runner="pooling",
|
||||
dtype=DTYPE,
|
||||
max_model_len=colbert_max_model_len,
|
||||
enforce_eager=True,
|
||||
**colbert_extra_kwargs,
|
||||
) as vllm_model:
|
||||
q_outputs = vllm_model.token_embed([TEXTS_1[0]])
|
||||
d_outputs = vllm_model.token_embed(TEXTS_2)
|
||||
|
||||
q_emb = torch.as_tensor(q_outputs[0])
|
||||
|
||||
manual_scores = []
|
||||
for d_out in d_outputs:
|
||||
d_emb = torch.as_tensor(d_out)
|
||||
manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
|
||||
|
||||
vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2)
|
||||
|
||||
assert len(vllm_scores) == 2
|
||||
for i in range(2):
|
||||
assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
|
||||
|
||||
|
||||
def test_colbert_late_interaction_N_to_N(
|
||||
vllm_runner,
|
||||
colbert_model_name,
|
||||
colbert_max_model_len,
|
||||
colbert_extra_kwargs,
|
||||
):
|
||||
"""Test ColBERT late interaction scoring with N:N query-documents."""
|
||||
with vllm_runner(
|
||||
colbert_model_name,
|
||||
runner="pooling",
|
||||
dtype=DTYPE,
|
||||
max_model_len=colbert_max_model_len,
|
||||
enforce_eager=True,
|
||||
**colbert_extra_kwargs,
|
||||
) as vllm_model:
|
||||
q_outputs = vllm_model.token_embed(TEXTS_1)
|
||||
d_outputs = vllm_model.token_embed(TEXTS_2)
|
||||
|
||||
manual_scores = []
|
||||
for q_out, d_out in zip(q_outputs, d_outputs):
|
||||
q_emb = torch.as_tensor(q_out)
|
||||
d_emb = torch.as_tensor(d_out)
|
||||
manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
|
||||
|
||||
vllm_scores = vllm_model.score(TEXTS_1, TEXTS_2)
|
||||
|
||||
assert len(vllm_scores) == 2
|
||||
for i in range(2):
|
||||
assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
|
||||
|
||||
|
||||
def test_colbert_relevance_ordering(
|
||||
vllm_runner,
|
||||
colbert_model_name,
|
||||
colbert_max_model_len,
|
||||
colbert_extra_kwargs,
|
||||
):
|
||||
"""Test that ColBERT scores relevant documents higher than irrelevant."""
|
||||
query = "What is machine learning?"
|
||||
documents = [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"Python is a programming language.",
|
||||
"Deep learning uses neural networks.",
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
colbert_model_name,
|
||||
runner="pooling",
|
||||
dtype=DTYPE,
|
||||
max_model_len=colbert_max_model_len,
|
||||
enforce_eager=True,
|
||||
**colbert_extra_kwargs,
|
||||
) as vllm_model:
|
||||
scores = vllm_model.score(query, documents)
|
||||
|
||||
assert len(scores) == 3
|
||||
assert scores[0] > scores[1], "ML doc should score higher than Python doc"
|
||||
assert scores[2] > scores[1], "DL doc should score higher than Python doc"
|
||||
|
||||
|
||||
def test_colbert_embed_not_supported(
|
||||
vllm_runner,
|
||||
colbert_model_name,
|
||||
colbert_max_model_len,
|
||||
colbert_extra_kwargs,
|
||||
):
|
||||
"""Test that ColBERT model does not support 'embed' task."""
|
||||
with (
|
||||
vllm_runner(
|
||||
colbert_model_name,
|
||||
runner="pooling",
|
||||
dtype=DTYPE,
|
||||
max_model_len=colbert_max_model_len,
|
||||
enforce_eager=True,
|
||||
**colbert_extra_kwargs,
|
||||
) as vllm_model,
|
||||
pytest.raises(ValueError, match="Embedding API is not supported"),
|
||||
):
|
||||
vllm_model.embed([TEXTS_1[0]])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", list(COLBERT_MODELS.keys()))
|
||||
def test_colbert_hf_comparison(vllm_runner, backend):
|
||||
"""Test that vLLM ColBERT embeddings match HuggingFace for each backend."""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
spec = COLBERT_MODELS[backend]
|
||||
hf_spec = spec["hf_comparison"]
|
||||
extra_kwargs = spec["extra_kwargs"]
|
||||
model_name = spec["model"]
|
||||
assert isinstance(model_name, str)
|
||||
assert isinstance(hf_spec, dict)
|
||||
assert isinstance(extra_kwargs, dict)
|
||||
test_texts = [TEXTS_1[0], TEXTS_2[0]]
|
||||
|
||||
with vllm_runner(
|
||||
model_name,
|
||||
runner="pooling",
|
||||
dtype="float32",
|
||||
max_model_len=spec["max_model_len"],
|
||||
enforce_eager=True,
|
||||
**extra_kwargs,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_embed(test_texts)
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
hf_tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_name,
|
||||
trust_remote_code=hf_spec.get("trust_remote_code", False),
|
||||
)
|
||||
with _hf_colbert_model(model_name, hf_spec, device) as (hf_model, linear_weight):
|
||||
hf_embeddings = _compute_hf_colbert_embeddings(
|
||||
hf_model,
|
||||
hf_tokenizer,
|
||||
linear_weight,
|
||||
test_texts,
|
||||
device,
|
||||
)
|
||||
|
||||
_assert_embeddings_close(vllm_outputs, hf_embeddings)
|
||||
@@ -0,0 +1,89 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
from ...utils import check_embeddings_close
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
# Be careful of the order of models, decoder-only models should be
|
||||
# placed before encoder-only models, otherwise `Qwen2.5-0.5B-Instruct`
|
||||
# case won't pass because gte-Qwen2-1.5B-instruct will cache custom
|
||||
# model code with bidirectional attention.
|
||||
# [Decoder-only]
|
||||
pytest.param(
|
||||
"BAAI/bge-multilingual-gemma2",
|
||||
marks=[pytest.mark.core_model, pytest.mark.slow_test],
|
||||
),
|
||||
pytest.param(
|
||||
"intfloat/e5-mistral-7b-instruct",
|
||||
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
|
||||
),
|
||||
pytest.param(
|
||||
"ssmits/Qwen2-7B-Instruct-embed-base", marks=[pytest.mark.cpu_model]
|
||||
),
|
||||
# [Encoder-only]
|
||||
pytest.param(
|
||||
"BAAI/bge-base-en-v1.5",
|
||||
marks=[
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.cpu_model,
|
||||
pytest.mark.slow_test,
|
||||
],
|
||||
),
|
||||
pytest.param("sentence-transformers/all-MiniLM-L12-v2"),
|
||||
pytest.param("intfloat/multilingual-e5-small"),
|
||||
# [Cross-Encoder]
|
||||
pytest.param(
|
||||
"sentence-transformers/stsb-roberta-base-v2",
|
||||
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model,
|
||||
) -> None:
|
||||
vllm_extra_kwargs = {}
|
||||
if model == "ssmits/Qwen2-7B-Instruct-embed-base":
|
||||
vllm_extra_kwargs["pooler_config"] = PoolerConfig(
|
||||
seq_pooling_type="MEAN", use_activation=False
|
||||
)
|
||||
|
||||
max_model_len: int | None = 512
|
||||
if model in [
|
||||
"sentence-transformers/all-MiniLM-L12-v2",
|
||||
"sentence-transformers/stsb-roberta-base-v2",
|
||||
]:
|
||||
max_model_len = None
|
||||
|
||||
# The example_prompts has ending "\n", for example:
|
||||
# "Write a short story about a robot that dreams for the first time.\n"
|
||||
# sentence_transformers will strip the input texts, see:
|
||||
# https://github.com/UKPLab/sentence-transformers/blob/v3.1.1/sentence_transformers/models/Transformer.py#L159
|
||||
# This makes the input_ids different between hf_model and vllm_model.
|
||||
# So we need to strip the input texts to avoid test failing.
|
||||
example_prompts = [str(s).strip() for s in example_prompts]
|
||||
|
||||
with hf_runner(model, is_sentence_transformer=True) as hf_model:
|
||||
hf_outputs = hf_model.encode(example_prompts)
|
||||
|
||||
with vllm_runner(
|
||||
model, runner="pooling", max_model_len=max_model_len, **vllm_extra_kwargs
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.embed(example_prompts)
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=hf_outputs,
|
||||
embeddings_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
tol=1e-2,
|
||||
)
|
||||
@@ -0,0 +1,48 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm import TokensPrompt
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["Qwen/Qwen3-0.6B"],
|
||||
)
|
||||
@torch.inference_mode
|
||||
def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
|
||||
n_prompt_tokens = [55, 56, 57]
|
||||
token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=128,
|
||||
enforce_eager=True,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
enable_prefix_caching=True,
|
||||
) as vllm_model:
|
||||
pooling_outputs = vllm_model.llm.encode(
|
||||
[TokensPrompt(prompt_token_ids=t) for t in token_prompts],
|
||||
pooling_task="token_embed",
|
||||
)
|
||||
|
||||
for n, output in zip(n_prompt_tokens, pooling_outputs):
|
||||
assert len(output.prompt_token_ids) == n
|
||||
assert len(output.outputs.data) == n
|
||||
assert output.num_cached_tokens == 0
|
||||
|
||||
# test enable_prefix_caching plus all pooling
|
||||
# we need to skip reading cache at this request by
|
||||
# request.skip_reading_prefix_cache
|
||||
pooling_outputs = vllm_model.llm.encode(
|
||||
[TokensPrompt(prompt_token_ids=t) for t in token_prompts],
|
||||
pooling_task="token_embed",
|
||||
)
|
||||
|
||||
for n, output in zip(n_prompt_tokens, pooling_outputs):
|
||||
assert len(output.prompt_token_ids) == n
|
||||
assert len(output.outputs.data) == n
|
||||
assert output.num_cached_tokens == 0
|
||||
@@ -0,0 +1,188 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import numpy as np
|
||||
import pytest
|
||||
from scipy.spatial.distance import cosine
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import ModelConfig
|
||||
|
||||
from ....utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
|
||||
from .embed_utils import run_client_embeddings
|
||||
|
||||
MODEL_NAME = "parasail-ai/GritLM-7B-vllm"
|
||||
MAX_MODEL_LEN = 4000
|
||||
ATOL = 2.3e-3
|
||||
|
||||
|
||||
def _arr(arr):
|
||||
"""
|
||||
Convert a list of integers to an array of integers.
|
||||
"""
|
||||
return np.array(arr)
|
||||
|
||||
|
||||
def test_find_array():
|
||||
from vllm.model_executor.models.gritlm import GritLMMeanPool
|
||||
|
||||
model_config = ModelConfig(
|
||||
MODEL_NAME,
|
||||
runner="pooling",
|
||||
dtype="bfloat16",
|
||||
seed=0,
|
||||
)
|
||||
pooling = GritLMMeanPool(model_config=model_config)
|
||||
|
||||
arr = _arr([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
|
||||
|
||||
assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=0) == 3
|
||||
assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=1) == 3
|
||||
assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=5) == -1
|
||||
assert pooling._find_array(arr, _arr([3, 4, 5]), end_idx=3) == -1
|
||||
assert pooling._find_array(arr, _arr([3, 4, 5]), end_idx=4) == 3
|
||||
assert pooling._find_array(arr, _arr([3, 5]), start_idx=0) == -1
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
pooling._find_array(arr, _arr([3, 4, 5]), start_idx=-1)
|
||||
|
||||
|
||||
def run_llm_encode(
|
||||
llm: LLM,
|
||||
queries: list[str],
|
||||
instruction: str,
|
||||
) -> list[list[float]]:
|
||||
outputs = llm.embed([instruction + q for q in queries])
|
||||
return [output.outputs.embedding for output in outputs]
|
||||
|
||||
|
||||
def gritlm_instruction(instruction):
|
||||
return (
|
||||
"<|user|>\n" + instruction + "\n<|embed|>\n" if instruction else "<|embed|>\n"
|
||||
)
|
||||
|
||||
|
||||
def get_test_data():
|
||||
"""
|
||||
Grabbed this test data and the expected values from
|
||||
README.md in https://github.com/ContextualAI/gritlm
|
||||
"""
|
||||
q_instruction = gritlm_instruction(
|
||||
"Given a scientific paper title, retrieve the paper's abstract",
|
||||
)
|
||||
queries = [
|
||||
"Bitcoin: A Peer-to-Peer Electronic Cash System",
|
||||
"Generative Representational Instruction Tuning",
|
||||
]
|
||||
|
||||
d_instruction = gritlm_instruction("")
|
||||
documents = [
|
||||
# ruff: noqa: E501
|
||||
"A purely peer-to-peer version of electronic cash would allow online payments to be sent directly from one party to another without going through a financial institution. Digital signatures provide part of the solution, but the main benefits are lost if a trusted third party is still required to prevent double-spending. We propose a solution to the double-spending problem using a peer-to-peer network. The network timestamps transactions by hashing them into an ongoing chain of hash-based proof-of-work, forming a record that cannot be changed without redoing the proof-of-work. The longest chain not only serves as proof of the sequence of events witnessed, but proof that it came from the largest pool of CPU power. As long as a majority of CPU power is controlled by nodes that are not cooperating to attack the network, they'll generate the longest chain and outpace attackers. The network itself requires minimal structure. Messages are broadcast on a best effort basis, and nodes can leave and rejoin the network at will, accepting the longest proof-of-work chain as proof of what happened while they were gone.",
|
||||
"All text-based language problems can be reduced to either generation or embedding. Current models only perform well at one or the other. We introduce generative representational instruction tuning (GRIT) whereby a large language model is trained to handle both generative and embedding tasks by distinguishing between them through instructions. Compared to other open models, our resulting GritLM 7B sets a new state of the art on the Massive Text Embedding Benchmark (MTEB) and outperforms all models up to its size on a range of generative tasks. By scaling up further, GritLM 8X7B outperforms all open generative language models that we tried while still being among the best embedding models. Notably, we find that GRIT matches training on only generative or embedding data, thus we can unify both at no performance loss. Among other benefits, the unification via GRIT speeds up Retrieval-Augmented Generation (RAG) by > 60% for long documents, by no longer requiring separate retrieval and generation models. Models, code, etc. are freely available at https://github.com/ContextualAI/gritlm.",
|
||||
]
|
||||
|
||||
return queries, q_instruction, documents, d_instruction
|
||||
|
||||
|
||||
def validate_embed_output(q_rep: list[list[float]], d_rep: list[list[float]]):
|
||||
cosine_sim_q0_d0 = 1 - cosine(q_rep[0], d_rep[0])
|
||||
assert cosine_sim_q0_d0 == pytest.approx(0.609, abs=ATOL)
|
||||
|
||||
cosine_sim_q0_d1 = 1 - cosine(q_rep[0], d_rep[1])
|
||||
assert cosine_sim_q0_d1 == pytest.approx(0.101, abs=ATOL)
|
||||
|
||||
cosine_sim_q1_d0 = 1 - cosine(q_rep[1], d_rep[0])
|
||||
assert cosine_sim_q1_d0 == pytest.approx(0.120, abs=ATOL)
|
||||
|
||||
cosine_sim_q1_d1 = 1 - cosine(q_rep[1], d_rep[1])
|
||||
assert cosine_sim_q1_d1 == pytest.approx(0.534, abs=ATOL)
|
||||
|
||||
|
||||
def test_gritlm_offline_embedding(vllm_runner):
|
||||
queries, q_instruction, documents, d_instruction = get_test_data()
|
||||
|
||||
with vllm_runner(
|
||||
MODEL_NAME,
|
||||
runner="pooling",
|
||||
max_model_len=MAX_MODEL_LEN,
|
||||
) as vllm_model:
|
||||
llm = vllm_model.llm
|
||||
|
||||
d_rep = run_llm_encode(
|
||||
llm,
|
||||
documents,
|
||||
d_instruction,
|
||||
)
|
||||
q_rep = run_llm_encode(
|
||||
llm,
|
||||
queries,
|
||||
q_instruction,
|
||||
)
|
||||
|
||||
validate_embed_output(q_rep, d_rep)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gritlm_api_server_embedding():
|
||||
queries, q_instruction, documents, d_instruction = get_test_data()
|
||||
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--max_model_len",
|
||||
str(MAX_MODEL_LEN),
|
||||
*ROCM_EXTRA_ARGS,
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=ROCM_ENV_OVERRIDES) as server:
|
||||
client_embedding = server.get_async_client()
|
||||
|
||||
d_rep = await run_client_embeddings(
|
||||
client_embedding,
|
||||
MODEL_NAME,
|
||||
documents,
|
||||
d_instruction,
|
||||
)
|
||||
q_rep = await run_client_embeddings(
|
||||
client_embedding,
|
||||
MODEL_NAME,
|
||||
queries,
|
||||
q_instruction,
|
||||
)
|
||||
|
||||
validate_embed_output(q_rep, d_rep)
|
||||
|
||||
|
||||
def test_gritlm_offline_generate(monkeypatch: pytest.MonkeyPatch, vllm_runner):
|
||||
input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
|
||||
|
||||
with vllm_runner(
|
||||
MODEL_NAME,
|
||||
runner="generate",
|
||||
max_model_len=MAX_MODEL_LEN,
|
||||
) as vllm_model:
|
||||
llm = vllm_model.llm
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=256)
|
||||
outputs = llm.generate(input, sampling_params=sampling_params)
|
||||
|
||||
assert outputs[0].outputs[0].text == "The capital of France is Paris."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gritlm_api_server_generate():
|
||||
input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
|
||||
|
||||
args = ["--runner", "generate", "--max_model_len", str(MAX_MODEL_LEN)]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as server:
|
||||
client_generate = server.get_async_client()
|
||||
|
||||
outputs = await client_generate.completions.create(
|
||||
model=MODEL_NAME,
|
||||
prompt=input,
|
||||
max_tokens=256,
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
assert outputs.choices[0].text == "The capital of France is Paris."
|
||||
@@ -0,0 +1,47 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoModelForSequenceClassification
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["nie3e/sentiment-polish-gpt2-small"],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_classify_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
with hf_runner(
|
||||
model, dtype=dtype, auto_cls=AutoModelForSequenceClassification
|
||||
) as hf_model:
|
||||
hf_outputs = hf_model.classify(example_prompts)
|
||||
|
||||
for head_dtype_str in ["float32", "model"]:
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
hf_overrides={"head_dtype": head_dtype_str},
|
||||
) as vllm_model:
|
||||
model_config = vllm_model.llm.llm_engine.model_config
|
||||
model_dtype = model_config.dtype
|
||||
head_dtype = model_config.head_dtype
|
||||
|
||||
if head_dtype_str == "float32":
|
||||
assert head_dtype == torch.float32
|
||||
elif head_dtype_str == "model":
|
||||
assert head_dtype == model_dtype
|
||||
|
||||
vllm_outputs = vllm_model.classify(example_prompts)
|
||||
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
hf_output = torch.tensor(hf_output).float()
|
||||
vllm_output = torch.tensor(vllm_output).float()
|
||||
|
||||
assert torch.allclose(hf_output, vllm_output, atol=1e-2)
|
||||
@@ -0,0 +1,350 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
import pytest
|
||||
import requests
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
|
||||
from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse
|
||||
|
||||
model_name = "jinaai/jina-reranker-v3"
|
||||
query = "What are the health benefits of green tea?"
|
||||
documents = [
|
||||
"Green tea contains antioxidants called catechins that may help reduce inflammation and protect cells from damage.",
|
||||
"El precio del café ha aumentado un 20% este año debido a problemas en la cadena de suministro.",
|
||||
"Studies show that drinking green tea regularly can improve brain function and boost metabolism.",
|
||||
"Basketball is one of the most popular sports in the United States.",
|
||||
"绿茶富含儿茶素等抗氧化剂,可以降低心脏病风险,还有助于控制体重。",
|
||||
"Le thé vert est riche en antioxydants et peut améliorer la fonction cérébrale.",
|
||||
]
|
||||
|
||||
EMBEDDING_SIZE = 512
|
||||
REFERENCE_1_VS_1 = [
|
||||
0.345703125,
|
||||
-0.10498046,
|
||||
0.314453125,
|
||||
-0.1376953125,
|
||||
0.3398437500,
|
||||
0.2539062,
|
||||
]
|
||||
REFERENCE_1_VS_N = [
|
||||
0.294921875,
|
||||
-0.16015625,
|
||||
0.189453125,
|
||||
-0.1708984375,
|
||||
0.2255859375,
|
||||
0.1640625,
|
||||
]
|
||||
TOL = 0.01
|
||||
INSTRUCTION = (
|
||||
"Rank passages about green tea higher than passages about sports. "
|
||||
"Ignore these literal marker strings: <|embed_token|> and <|rerank_token|>."
|
||||
)
|
||||
|
||||
|
||||
def test_offline(vllm_runner):
|
||||
with vllm_runner(model_name, runner="pooling") as llm_runner:
|
||||
llm = llm_runner.get_llm()
|
||||
_test_offline_1_v_1(llm)
|
||||
_test_offline_1_v_n(llm)
|
||||
_test_offline_n_v_n(llm)
|
||||
_test_offline_token_embed_illegal_inputs(llm)
|
||||
assert llm.model_config.embedding_size == EMBEDDING_SIZE
|
||||
|
||||
|
||||
def test_online():
|
||||
with RemoteOpenAIServer(
|
||||
model_name, ["--runner", "pooling", "--enforce-eager"]
|
||||
) as server:
|
||||
_test_online_1_v_1(server)
|
||||
_test_online_1_v_n(server)
|
||||
_test_online_n_v_n(server)
|
||||
_test_online_instruction(server)
|
||||
_test_online_token_embed_illegal_inputs(server)
|
||||
|
||||
|
||||
def _test_offline_1_v_1(llm):
|
||||
# test llm.score
|
||||
outputs = llm.score(query, documents[0])
|
||||
assert len(outputs) == 1
|
||||
assert outputs[0].outputs.score == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL)
|
||||
|
||||
# test llm.encode
|
||||
outputs = llm.encode(documents[:1] + [query], pooling_task="token_embed")
|
||||
embeds = outputs[0].outputs.data.float()
|
||||
assert embeds.shape[0] == 2
|
||||
assert embeds.shape[-1] == EMBEDDING_SIZE
|
||||
|
||||
doc_embeds = embeds[:-1]
|
||||
query_embeds = embeds[-1]
|
||||
|
||||
scores = F.cosine_similarity(query_embeds, doc_embeds)
|
||||
assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL)
|
||||
|
||||
|
||||
def _test_offline_1_v_n(llm):
|
||||
# test llm.score
|
||||
outputs = llm.score(query, documents)
|
||||
assert len(outputs) == len(documents)
|
||||
|
||||
for expected, output in zip(REFERENCE_1_VS_N, outputs):
|
||||
actual = output.outputs.score
|
||||
assert actual == pytest.approx(expected, abs=TOL)
|
||||
|
||||
# test llm.encode
|
||||
outputs = llm.encode(documents + [query], pooling_task="token_embed")
|
||||
embeds = outputs[0].outputs.data.float()
|
||||
assert embeds.shape[0] == len(documents) + 1
|
||||
|
||||
doc_embeds = embeds[:-1]
|
||||
query_embeds = embeds[-1]
|
||||
|
||||
scores = F.cosine_similarity(query_embeds, doc_embeds)
|
||||
|
||||
assert len(scores) == len(documents)
|
||||
for expected, actual in zip(REFERENCE_1_VS_N, scores):
|
||||
assert actual == pytest.approx(expected, abs=TOL)
|
||||
|
||||
|
||||
def _test_offline_n_v_n(llm):
|
||||
# test llm.score
|
||||
outputs = llm.score([query] * len(documents), documents)
|
||||
assert len(outputs) == len(documents)
|
||||
|
||||
for expected, output in zip(REFERENCE_1_VS_1, outputs):
|
||||
actual = output.outputs.score
|
||||
assert actual == pytest.approx(expected, abs=TOL)
|
||||
|
||||
# test llm.encode
|
||||
for doc, expected in zip(documents, REFERENCE_1_VS_1):
|
||||
outputs = llm.encode([doc, query], pooling_task="token_embed")
|
||||
embeds = outputs[0].outputs.data.float()
|
||||
assert embeds.shape[0] == 2
|
||||
|
||||
doc_embeds = embeds[:-1]
|
||||
query_embeds = embeds[-1]
|
||||
|
||||
scores = F.cosine_similarity(query_embeds, doc_embeds)
|
||||
assert scores[0] == pytest.approx(expected, abs=TOL)
|
||||
|
||||
|
||||
def _test_offline_token_embed_illegal_inputs(llm):
|
||||
with pytest.raises(
|
||||
ValueError, match="The JinaForRanking model requires at least 2 inputs."
|
||||
):
|
||||
llm.encode([query], pooling_task="token_embed")
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="The JinaForRanking model only supports text as input."
|
||||
):
|
||||
llm.encode([1, 2, 3], pooling_task="token_embed")
|
||||
|
||||
|
||||
def _get_score_response(server, query, document, **extra_body):
|
||||
payload = {
|
||||
"model": model_name,
|
||||
"queries": query,
|
||||
"documents": document,
|
||||
}
|
||||
payload.update(extra_body)
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json=payload,
|
||||
)
|
||||
|
||||
score_response.raise_for_status()
|
||||
return ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
|
||||
def _get_scores(server, query, document):
|
||||
score = _get_score_response(server, query, document)
|
||||
|
||||
return [d.score for d in score.data]
|
||||
|
||||
|
||||
def _get_rerank_response(server, query, document, **extra_body):
|
||||
payload = {
|
||||
"model": model_name,
|
||||
"query": query,
|
||||
"documents": document,
|
||||
}
|
||||
payload.update(extra_body)
|
||||
rerank_response = requests.post(
|
||||
server.url_for("rerank"),
|
||||
json=payload,
|
||||
)
|
||||
|
||||
rerank_response.raise_for_status()
|
||||
return RerankResponse.model_validate(rerank_response.json())
|
||||
|
||||
|
||||
def _get_embeds(server, prompts: list[str]):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"task": "token_embed",
|
||||
"input": prompts,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
return torch.as_tensor([d.data for d in poolings.data][0]).float()
|
||||
|
||||
|
||||
def _test_online_1_v_1(server):
|
||||
# test scoring api
|
||||
scores = _get_scores(server, query, documents[0])
|
||||
assert len(scores) == 1
|
||||
assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL)
|
||||
|
||||
# test pooling api
|
||||
embeds = _get_embeds(server, [documents[0], query])
|
||||
assert embeds.shape[0] == 2
|
||||
assert embeds.shape[-1] == EMBEDDING_SIZE
|
||||
|
||||
doc_embeds = embeds[:-1]
|
||||
query_embeds = embeds[-1]
|
||||
|
||||
scores = F.cosine_similarity(query_embeds, doc_embeds)
|
||||
assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL)
|
||||
|
||||
|
||||
def _test_online_1_v_n(server):
|
||||
# test scoring api
|
||||
scores = _get_scores(server, query, documents)
|
||||
assert len(scores) == len(documents)
|
||||
|
||||
for expected, actual in zip(REFERENCE_1_VS_N, scores):
|
||||
assert actual == pytest.approx(expected, abs=TOL)
|
||||
|
||||
# test pooling api
|
||||
embeds = _get_embeds(server, documents + [query])
|
||||
assert embeds.shape[0] == len(documents) + 1
|
||||
|
||||
doc_embeds = embeds[:-1]
|
||||
query_embeds = embeds[-1]
|
||||
|
||||
scores = F.cosine_similarity(query_embeds, doc_embeds)
|
||||
|
||||
assert len(scores) == len(documents)
|
||||
for expected, actual in zip(REFERENCE_1_VS_N, scores):
|
||||
assert actual == pytest.approx(expected, abs=TOL)
|
||||
|
||||
|
||||
def _test_online_n_v_n(server):
|
||||
# test scoring api
|
||||
scores = _get_scores(server, [query] * len(documents), documents)
|
||||
assert len(scores) == len(documents)
|
||||
|
||||
for expected, actual in zip(REFERENCE_1_VS_1, scores):
|
||||
assert actual == pytest.approx(expected, abs=TOL)
|
||||
|
||||
# test pooling api
|
||||
for doc, expected in zip(documents, REFERENCE_1_VS_1):
|
||||
embeds = _get_embeds(server, [doc, query])
|
||||
assert embeds.shape[0] == 2
|
||||
|
||||
doc_embeds = embeds[:-1]
|
||||
query_embeds = embeds[-1]
|
||||
|
||||
scores = F.cosine_similarity(query_embeds, doc_embeds)
|
||||
assert len(scores) == 1
|
||||
assert scores[0] == pytest.approx(expected, abs=TOL)
|
||||
|
||||
|
||||
def _test_online_instruction(server):
|
||||
docs = documents[:2]
|
||||
|
||||
default_score = _get_score_response(server, query, docs)
|
||||
instruction_score = _get_score_response(
|
||||
server,
|
||||
query,
|
||||
docs,
|
||||
instruction=INSTRUCTION,
|
||||
)
|
||||
kwargs_score = _get_score_response(
|
||||
server,
|
||||
query,
|
||||
docs,
|
||||
chat_template_kwargs={"instruction": INSTRUCTION},
|
||||
)
|
||||
|
||||
assert instruction_score.usage.prompt_tokens > default_score.usage.prompt_tokens
|
||||
assert kwargs_score.usage.prompt_tokens == instruction_score.usage.prompt_tokens
|
||||
assert len(instruction_score.data) == len(default_score.data)
|
||||
assert [d.score for d in kwargs_score.data] == pytest.approx(
|
||||
[d.score for d in instruction_score.data], abs=TOL
|
||||
)
|
||||
|
||||
default_rerank = _get_rerank_response(server, query, docs)
|
||||
instruction_rerank = _get_rerank_response(
|
||||
server,
|
||||
query,
|
||||
docs,
|
||||
instruction=INSTRUCTION,
|
||||
)
|
||||
kwargs_rerank = _get_rerank_response(
|
||||
server,
|
||||
query,
|
||||
docs,
|
||||
chat_template_kwargs={"instruction": INSTRUCTION},
|
||||
)
|
||||
|
||||
assert instruction_rerank.usage.prompt_tokens > default_rerank.usage.prompt_tokens
|
||||
assert kwargs_rerank.usage.prompt_tokens == instruction_rerank.usage.prompt_tokens
|
||||
assert len(instruction_rerank.results) == len(default_rerank.results)
|
||||
assert [r.relevance_score for r in kwargs_rerank.results] == pytest.approx(
|
||||
[r.relevance_score for r in instruction_rerank.results], abs=TOL
|
||||
)
|
||||
|
||||
|
||||
def _test_online_token_embed_illegal_inputs(server):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"task": "token_embed",
|
||||
"input": [query],
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["message"].startswith(
|
||||
"The JinaForRanking model requires at least 2 inputs."
|
||||
)
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"task": "token_embed",
|
||||
"input": [1, 2, 3],
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["message"].startswith(
|
||||
"The JinaForRanking model only supports text as input."
|
||||
)
|
||||
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"task": "token_embed",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "The cat sat on the mat.",
|
||||
}
|
||||
],
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["message"].startswith(
|
||||
"The JinaForRanking does not support chat Request."
|
||||
)
|
||||
@@ -0,0 +1,228 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for max_tokens_per_doc and max_tokens_per_query.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import VLLM_PATH, RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.scoring.protocol import RerankResponse
|
||||
|
||||
os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
|
||||
|
||||
TEMPLATE_DIR = str(VLLM_PATH / "examples/pooling/score/template")
|
||||
ExpectedPromptTokens = int | tuple[int, ...]
|
||||
|
||||
long_query = "What is the capital of France?" * 20
|
||||
long_doc = "The capital of France is Paris. " * 20
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestConfig:
|
||||
model: str
|
||||
args: list[str]
|
||||
without_truncated_prompt_tokens: ExpectedPromptTokens
|
||||
with_max_tokens_per_query_prompt_tokens: ExpectedPromptTokens
|
||||
with_max_tokens_per_doc_prompt_tokens: ExpectedPromptTokens
|
||||
with_max_tokens_per_query_and_doc_prompt_tokens: ExpectedPromptTokens
|
||||
|
||||
|
||||
RERANK_CONFIGS = [
|
||||
# 1. cross-encoder
|
||||
TestConfig(
|
||||
model="jinaai/jina-reranker-v2-base-multilingual",
|
||||
args=[
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"1024",
|
||||
"--trust-remote-code",
|
||||
],
|
||||
without_truncated_prompt_tokens=284,
|
||||
with_max_tokens_per_query_prompt_tokens=154,
|
||||
with_max_tokens_per_doc_prompt_tokens=154,
|
||||
with_max_tokens_per_query_and_doc_prompt_tokens=24,
|
||||
),
|
||||
# 2. cross-encoder + score template
|
||||
TestConfig(
|
||||
model="Qwen/Qwen3-Reranker-0.6B",
|
||||
args=[
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"1024",
|
||||
"--hf-overrides",
|
||||
json.dumps(
|
||||
{
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
}
|
||||
),
|
||||
"--chat-template",
|
||||
os.path.join(TEMPLATE_DIR, "qwen3_reranker.jinja"),
|
||||
],
|
||||
without_truncated_prompt_tokens=352,
|
||||
with_max_tokens_per_query_prompt_tokens=223,
|
||||
with_max_tokens_per_doc_prompt_tokens=221,
|
||||
with_max_tokens_per_query_and_doc_prompt_tokens=92,
|
||||
),
|
||||
# 3. bi-encoder
|
||||
TestConfig(
|
||||
model="intfloat/multilingual-e5-small",
|
||||
args=[
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--trust-remote-code",
|
||||
],
|
||||
# This model has produced both prompt-token totals in CI/local cache;
|
||||
# keep truncation checks exact while tolerating the boundary delta.
|
||||
without_truncated_prompt_tokens=(285, 286),
|
||||
with_max_tokens_per_query_prompt_tokens=(155, 156),
|
||||
with_max_tokens_per_doc_prompt_tokens=155,
|
||||
with_max_tokens_per_query_and_doc_prompt_tokens=25,
|
||||
),
|
||||
# 4. late-interaction
|
||||
TestConfig(
|
||||
model="answerdotai/answerai-colbert-small-v1",
|
||||
args=[
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--trust-remote-code",
|
||||
],
|
||||
without_truncated_prompt_tokens=285,
|
||||
with_max_tokens_per_query_prompt_tokens=155,
|
||||
with_max_tokens_per_doc_prompt_tokens=155,
|
||||
with_max_tokens_per_query_and_doc_prompt_tokens=25,
|
||||
),
|
||||
# 5. jinaai/jina-reranker-v3
|
||||
TestConfig(
|
||||
model="jinaai/jina-reranker-v3",
|
||||
args=[
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"1024",
|
||||
"--trust-remote-code",
|
||||
],
|
||||
without_truncated_prompt_tokens=567,
|
||||
with_max_tokens_per_query_prompt_tokens=308,
|
||||
with_max_tokens_per_doc_prompt_tokens=436,
|
||||
with_max_tokens_per_query_and_doc_prompt_tokens=177,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def assert_prompt_tokens(actual: int, expected: ExpectedPromptTokens) -> None:
|
||||
if isinstance(expected, int):
|
||||
assert actual == expected
|
||||
else:
|
||||
assert actual in expected
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=RERANK_CONFIGS, ids=lambda c: c.model)
|
||||
def server(request):
|
||||
config: TestConfig = request.param
|
||||
with RemoteOpenAIServer(config.model, config.args) as remote_server:
|
||||
yield config, remote_server
|
||||
|
||||
|
||||
def test_without_truncated(server):
|
||||
"""Test that max_tokens_per_doc truncates documents correctly."""
|
||||
config, remote_server = server
|
||||
|
||||
response = requests.post(
|
||||
remote_server.url_for("rerank"),
|
||||
json={"model": config.model, "query": long_query, "documents": [long_doc]},
|
||||
)
|
||||
response.raise_for_status()
|
||||
rerank = RerankResponse.model_validate(response.json())
|
||||
|
||||
assert rerank.id is not None
|
||||
assert rerank.results is not None
|
||||
assert len(rerank.results) == 1
|
||||
assert_prompt_tokens(
|
||||
rerank.usage.prompt_tokens,
|
||||
config.without_truncated_prompt_tokens,
|
||||
)
|
||||
|
||||
|
||||
def test_max_tokens_per_query(server):
|
||||
"""Test that max_tokens_per_doc truncates documents correctly."""
|
||||
config, remote_server = server
|
||||
|
||||
response = requests.post(
|
||||
remote_server.url_for("rerank"),
|
||||
json={
|
||||
"model": config.model,
|
||||
"query": long_query,
|
||||
"documents": [long_doc],
|
||||
"max_tokens_per_query": 10,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
rerank = RerankResponse.model_validate(response.json())
|
||||
|
||||
assert rerank.id is not None
|
||||
assert rerank.results is not None
|
||||
assert len(rerank.results) == 1
|
||||
assert_prompt_tokens(
|
||||
rerank.usage.prompt_tokens,
|
||||
config.with_max_tokens_per_query_prompt_tokens,
|
||||
)
|
||||
|
||||
|
||||
def test_max_tokens_per_doc(server):
|
||||
"""Test that max_tokens_per_doc truncates documents correctly."""
|
||||
config, remote_server = server
|
||||
|
||||
response = requests.post(
|
||||
remote_server.url_for("rerank"),
|
||||
json={
|
||||
"model": config.model,
|
||||
"query": long_query,
|
||||
"documents": [long_doc],
|
||||
"max_tokens_per_doc": 10,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
rerank = RerankResponse.model_validate(response.json())
|
||||
|
||||
assert rerank.id is not None
|
||||
assert rerank.results is not None
|
||||
assert len(rerank.results) == 1
|
||||
assert_prompt_tokens(
|
||||
rerank.usage.prompt_tokens,
|
||||
config.with_max_tokens_per_doc_prompt_tokens,
|
||||
)
|
||||
|
||||
|
||||
def test_max_tokens_per_query_and_doc(server):
|
||||
"""Test that max_tokens_per_doc truncates documents correctly."""
|
||||
config, remote_server = server
|
||||
|
||||
response = requests.post(
|
||||
remote_server.url_for("rerank"),
|
||||
json={
|
||||
"model": config.model,
|
||||
"query": long_query,
|
||||
"documents": [long_doc],
|
||||
"max_tokens_per_query": 10,
|
||||
"max_tokens_per_doc": 10,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
rerank = RerankResponse.model_validate(response.json())
|
||||
|
||||
assert rerank.id is not None
|
||||
assert rerank.results is not None
|
||||
assert len(rerank.results) == 1
|
||||
assert_prompt_tokens(
|
||||
rerank.usage.prompt_tokens,
|
||||
config.with_max_tokens_per_query_and_doc_prompt_tokens,
|
||||
)
|
||||
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm.config.pooler import PoolerConfig
|
||||
|
||||
|
||||
def test_idefics_multimodal(
|
||||
vllm_runner,
|
||||
) -> None:
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model_name="HuggingFaceM4/Idefics3-8B-Llama3",
|
||||
runner="pooling",
|
||||
convert="classify",
|
||||
load_format="dummy",
|
||||
max_model_len=512,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=1,
|
||||
disable_log_stats=True,
|
||||
dtype="bfloat16",
|
||||
) as vllm_model:
|
||||
llm = vllm_model.get_llm()
|
||||
outputs = llm.classify(prompts)
|
||||
for output in outputs:
|
||||
assert len(output.outputs.probs) == 2
|
||||
|
||||
|
||||
def update_config(config):
|
||||
text_config = config.get_text_config()
|
||||
text_config.update(
|
||||
{
|
||||
"architectures": ["Gemma3ForSequenceClassification"],
|
||||
"classifier_from_token": ["A", "B", "C", "D", "E"],
|
||||
"method": "no_post_processing",
|
||||
"id2label": {
|
||||
"A": "Chair",
|
||||
"B": "Couch",
|
||||
"C": "Table",
|
||||
"D": "Bed",
|
||||
"E": "Cupboard",
|
||||
},
|
||||
}
|
||||
)
|
||||
return config
|
||||
|
||||
|
||||
def test_gemma_multimodal(
|
||||
vllm_runner,
|
||||
) -> None:
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": """
|
||||
You are a helpful assistant. You will be given a product description
|
||||
which may also include an image. Classify the following product into
|
||||
one of the categories:
|
||||
|
||||
A = chair
|
||||
B = couch
|
||||
C = table
|
||||
D = bed
|
||||
E = cupboard
|
||||
|
||||
You'll answer with exactly one letter (A, B, C, D, or E).""",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/red_chair.jpg"
|
||||
},
|
||||
},
|
||||
{"type": "text", "text": "A fine 19th century piece of furniture."},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model_name="google/gemma-3-4b-it",
|
||||
runner="pooling",
|
||||
convert="classify",
|
||||
load_format="auto",
|
||||
hf_overrides=update_config,
|
||||
pooler_config=PoolerConfig(seq_pooling_type="LAST"),
|
||||
max_model_len=512,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=1,
|
||||
disable_log_stats=True,
|
||||
dtype="bfloat16",
|
||||
) as vllm_model:
|
||||
llm = vllm_model.get_llm()
|
||||
prompts = llm._preprocess_chat([messages])
|
||||
|
||||
result = llm.classify(prompts)
|
||||
assert result[0].outputs.probs[0] > 0.95
|
||||
assert all(c < 0.05 for c in result[0].outputs.probs[1:])
|
||||
@@ -0,0 +1,47 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoModel
|
||||
|
||||
from tests.models.utils import check_embeddings_close
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["BAAI/bge-m3"],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@torch.inference_mode
|
||||
def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype: str):
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
max_model_len=None,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_embed(example_prompts)
|
||||
|
||||
with hf_runner(
|
||||
model,
|
||||
auto_cls=AutoModel,
|
||||
) as hf_model:
|
||||
tokenizer = hf_model.tokenizer
|
||||
hf_outputs = []
|
||||
for prompt in example_prompts:
|
||||
inputs = tokenizer([prompt], return_tensors="pt")
|
||||
inputs = hf_model.wrap_device(inputs)
|
||||
output = hf_model.model(**inputs)
|
||||
embedding = output.last_hidden_state[0].float()
|
||||
# normal
|
||||
hf_outputs.append(embedding.cpu())
|
||||
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=hf_output,
|
||||
embeddings_1_lst=vllm_output,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
tol=1e-2,
|
||||
)
|
||||
@@ -0,0 +1,34 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoModelForSequenceClassification
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["Rami/multi-label-class-classification-on-github-issues"],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_classify_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
with vllm_runner(model, max_model_len=512, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.classify(example_prompts)
|
||||
|
||||
with hf_runner(
|
||||
model, dtype=dtype, auto_cls=AutoModelForSequenceClassification
|
||||
) as hf_model:
|
||||
hf_outputs = hf_model.classify(example_prompts)
|
||||
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
hf_output = torch.tensor(hf_output)
|
||||
vllm_output = torch.tensor(vllm_output)
|
||||
|
||||
assert torch.allclose(
|
||||
hf_output, vllm_output, 1e-3 if dtype == "float" else 1e-2
|
||||
)
|
||||
@@ -0,0 +1,111 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: SIM117
|
||||
|
||||
import pytest
|
||||
|
||||
from ...utils import EmbedModelInfo
|
||||
|
||||
MODELS = [
|
||||
EmbedModelInfo(
|
||||
"nomic-ai/nomic-embed-text-v1",
|
||||
# Fixme:
|
||||
# Update nomic-embed code to support the latest
|
||||
# HF version and remove revision set.
|
||||
revision="720244025c1a7e15661a174c63cce63c8218e52b",
|
||||
),
|
||||
# EmbedModelInfo("nomic-ai/nomic-embed-text-v1.5"),
|
||||
# EmbedModelInfo("nomic-ai/CodeRankEmbed"),
|
||||
EmbedModelInfo("nomic-ai/nomic-embed-text-v2-moe"),
|
||||
# EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-long"),
|
||||
]
|
||||
|
||||
rope_theta = 1000
|
||||
factor = 4.0
|
||||
original_max_position_embeddings = 2048
|
||||
max_model_len = int(original_max_position_embeddings * factor)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_default(model_info, vllm_runner):
|
||||
with vllm_runner(
|
||||
model_info.name,
|
||||
revision=model_info.revision,
|
||||
runner="pooling",
|
||||
max_model_len=None,
|
||||
) as vllm_model:
|
||||
model_config = vllm_model.llm.llm_engine.model_config
|
||||
if model_info.name == "nomic-ai/nomic-embed-text-v2-moe":
|
||||
# For nomic-embed-text-v2-moe the length is set to 512
|
||||
# by sentence_bert_config.json.
|
||||
assert model_config.max_model_len == 512
|
||||
if model_info.name == "nomic-ai/nomic-embed-text-v1":
|
||||
assert model_config.max_model_len == 8192
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_set_max_model_len_legal(model_info, vllm_runner):
|
||||
# set max_model_len <= 512
|
||||
with vllm_runner(
|
||||
model_info.name,
|
||||
revision=model_info.revision,
|
||||
runner="pooling",
|
||||
max_model_len=256,
|
||||
) as vllm_model:
|
||||
model_config = vllm_model.llm.llm_engine.model_config
|
||||
assert model_config.max_model_len == 256
|
||||
|
||||
# For nomic-embed-text-v2-moe the length is set to 512
|
||||
# by sentence_bert_config.json.
|
||||
if model_info.name == "nomic-ai/nomic-embed-text-v2-moe":
|
||||
with pytest.raises(ValueError):
|
||||
with vllm_runner(
|
||||
model_info.name,
|
||||
revision=model_info.revision,
|
||||
runner="pooling",
|
||||
max_model_len=1024,
|
||||
):
|
||||
pass
|
||||
return
|
||||
|
||||
# set 512 < max_model_len <= 2048
|
||||
with vllm_runner(
|
||||
model_info.name,
|
||||
revision=model_info.revision,
|
||||
runner="pooling",
|
||||
max_model_len=1024,
|
||||
) as vllm_model:
|
||||
model_config = vllm_model.llm.llm_engine.model_config
|
||||
assert model_config.max_model_len == 1024
|
||||
|
||||
# set max_model_len > 2048
|
||||
with vllm_runner(
|
||||
model_info.name,
|
||||
revision=model_info.revision,
|
||||
runner="pooling",
|
||||
max_model_len=4096,
|
||||
) as vllm_model:
|
||||
model_config = vllm_model.llm.llm_engine.model_config
|
||||
assert model_config.max_model_len == 4096
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_use_rope_scaling_legal(model_info, vllm_runner):
|
||||
hf_overrides = {
|
||||
"rope_parameters": {
|
||||
"rope_theta": rope_theta,
|
||||
"rope_type": "yarn",
|
||||
"factor": factor,
|
||||
"original_max_position_embeddings": original_max_position_embeddings,
|
||||
},
|
||||
"max_model_len": max_model_len,
|
||||
}
|
||||
|
||||
with vllm_runner(
|
||||
model_info.name,
|
||||
revision=model_info.revision,
|
||||
runner="pooling",
|
||||
max_model_len=None,
|
||||
hf_overrides=hf_overrides,
|
||||
):
|
||||
pass
|
||||
@@ -0,0 +1,167 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from tests.models.utils import softmax
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["jason9693/Qwen2.5-1.5B-apeach", "papluca/xlm-roberta-base-language-detection"],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_classify_models_using_activation(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=False),
|
||||
) as vllm_model:
|
||||
wo_activation_out = vllm_model.classify(example_prompts)
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=True),
|
||||
) as vllm_model:
|
||||
w_activation_out = vllm_model.classify(example_prompts)
|
||||
|
||||
for wo_activation, w_activation in zip(wo_activation_out, w_activation_out):
|
||||
wo_activation = torch.tensor(wo_activation)
|
||||
w_activation = torch.tensor(w_activation)
|
||||
|
||||
assert not torch.allclose(wo_activation, w_activation, atol=1e-2), (
|
||||
"pooler_config is not working"
|
||||
)
|
||||
assert torch.allclose(
|
||||
softmax(wo_activation), w_activation, 1e-3 if dtype == "float" else 1e-2
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"intfloat/multilingual-e5-small",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_embed_models_using_normalize(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=False),
|
||||
) as vllm_model:
|
||||
wo_normalize = torch.tensor(vllm_model.embed(example_prompts))
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=True),
|
||||
) as vllm_model:
|
||||
w_normalize = torch.tensor(vllm_model.embed(example_prompts))
|
||||
|
||||
assert not torch.allclose(wo_normalize, w_normalize, atol=1e-2), (
|
||||
"pooler_config normalize is not working"
|
||||
)
|
||||
assert torch.allclose(
|
||||
F.normalize(wo_normalize, p=2, dim=-1), w_normalize, atol=1e-2
|
||||
), "w_normal should be close to normal(wo_normal)."
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"internlm/internlm2-1_8b-reward",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_reward_models_using_activation(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=1024,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=False),
|
||||
) as vllm_model:
|
||||
wo_activation = vllm_model.token_classify(example_prompts)
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=1024,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=True),
|
||||
) as vllm_model:
|
||||
w_activation = vllm_model.token_classify(example_prompts)
|
||||
|
||||
for wo, w in zip(wo_activation, w_activation):
|
||||
wo = torch.tensor(wo)
|
||||
w = torch.tensor(w)
|
||||
|
||||
assert not torch.allclose(wo, w, atol=1e-2), (
|
||||
"pooler_config activation is not working"
|
||||
)
|
||||
assert torch.allclose(softmax(wo), w, atol=1e-2), (
|
||||
"w_activation should be close to activation(wo_activation)."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"intfloat/multilingual-e5-small",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_multi_vector_retrieval_models_using_normalize(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=False, task="token_embed"),
|
||||
) as vllm_model:
|
||||
wo_normalize = vllm_model.token_embed(example_prompts)
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=True, task="token_embed"),
|
||||
) as vllm_model:
|
||||
w_normalize = vllm_model.token_embed(example_prompts)
|
||||
|
||||
for wo, w in zip(wo_normalize, w_normalize):
|
||||
assert not torch.allclose(wo, w, atol=1e-2), (
|
||||
"pooler_config normalize is not working"
|
||||
)
|
||||
assert torch.allclose(F.normalize(wo, p=2, dim=-1), w, atol=1e-2), (
|
||||
"w_normal should be close to normal(wo_normal)."
|
||||
)
|
||||
@@ -0,0 +1,158 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from transformers import AutoModel
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....conftest import HfRunner
|
||||
from ....utils import VLLM_PATH
|
||||
from ...registry import HF_EXAMPLE_MODELS
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from _typeshed import StrPath
|
||||
|
||||
|
||||
FIXTURES_PATH = VLLM_PATH / "tests/models/fixtures"
|
||||
assert FIXTURES_PATH.exists()
|
||||
FIXTURE_REWARD_RESULT = {
|
||||
"Qwen/Qwen2.5-Math-PRM-7B": FIXTURES_PATH / "qwen2_5_math_prm_reward_step.json",
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def math_step_prompts():
|
||||
# ruff: noqa: E501
|
||||
data = {
|
||||
"system": "Please reason step by step, and put your final answer within \\boxed{}. ",
|
||||
"query": "Sue lives in a fun neighborhood. One weekend, the neighbors decided to play a prank on Sue. On Friday morning, the neighbors placed 18 pink plastic flamingos out on Sue's front yard. On Saturday morning, the neighbors took back one third of the flamingos, painted them white, and put these newly painted white flamingos back out on Sue's front yard. Then, on Sunday morning, they added another 18 pink plastic flamingos to the collection. At noon on Sunday, how many more pink plastic flamingos were out than white plastic flamingos?",
|
||||
"response": [
|
||||
"To find out how many more pink plastic flamingos were out than white plastic flamingos at noon on Sunday, we can break down the problem into steps. First, on Friday, the neighbors start with 18 pink plastic flamingos.",
|
||||
"On Saturday, they take back one third of the flamingos. Since there were 18 flamingos, (1/3 \\times 18 = 6) flamingos are taken back. So, they have (18 - 6 = 12) flamingos left in their possession. Then, they paint these 6 flamingos white and put them back out on Sue's front yard. Now, Sue has the original 12 pink flamingos plus the 6 new white ones. Thus, by the end of Saturday, Sue has (12 + 6 = 18) pink flamingos and 6 white flamingos.",
|
||||
"On Sunday, the neighbors add another 18 pink plastic flamingos to Sue's front yard. By the end of Sunday morning, Sue has (18 + 18 = 36) pink flamingos and still 6 white flamingos.",
|
||||
"To find the difference, subtract the number of white flamingos from the number of pink flamingos: (36 - 6 = 30). Therefore, at noon on Sunday, there were 30 more pink plastic flamingos out than white plastic flamingos. The answer is (\\boxed{30}).",
|
||||
],
|
||||
}
|
||||
answer = "<extra_0>".join(data["response"]) + "<extra_0>"
|
||||
prompt = f"<im_start>system\n{data['system']}<im_end>\n<im_start>user\n{data['query']}<im_end>\n<im_start>assistant\n{answer}<im_end><|endoftext|>"
|
||||
return [prompt]
|
||||
|
||||
|
||||
def step_reward_patch_hf_model(hf_model: HfRunner):
|
||||
# Patch the hf_runner to use the step reward function
|
||||
def make_step_rewards(
|
||||
logits: torch.Tensor, token_masks: torch.Tensor
|
||||
) -> list[list[float]]:
|
||||
probabilities = F.softmax(logits, dim=-1)
|
||||
probabilities = probabilities * token_masks.unsqueeze(-1)
|
||||
|
||||
all_scores_res: list[list[float]] = []
|
||||
for i in range(probabilities.size(0)):
|
||||
sample = probabilities[i] # seq_len, num_labels
|
||||
positive_probs = sample[sample != 0].view(-1, 2)
|
||||
non_zero_elements_list = positive_probs.cpu().tolist()
|
||||
all_scores_res.append(non_zero_elements_list)
|
||||
return all_scores_res
|
||||
|
||||
def reward(prompts: list[str]) -> list[list[float]]:
|
||||
input_ids = hf_model.tokenizer(prompts, return_tensors="pt").input_ids
|
||||
input_ids = hf_model.wrap_device(input_ids)
|
||||
outputs = hf_model.model(input_ids=input_ids)
|
||||
|
||||
step_sep_id = hf_model.tokenizer.encode("<extra_0>")[0]
|
||||
token_masks = input_ids == step_sep_id
|
||||
return make_step_rewards(outputs[0], token_masks)
|
||||
|
||||
hf_model.reward = reward # type: ignore[attr-defined]
|
||||
|
||||
return hf_model
|
||||
|
||||
|
||||
def dump_reward_outputs(outputs: list[list[float]], filename: "StrPath"):
|
||||
with open(filename, "w", encoding="utf-8") as f:
|
||||
json.dump(outputs, f)
|
||||
|
||||
|
||||
def load_reward_outputs(filename: "StrPath") -> list[list[float]]:
|
||||
with open(filename, encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
pytest.param(
|
||||
"Qwen/Qwen2.5-Math-PRM-7B",
|
||||
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_prm_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
math_step_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
if current_platform.is_cpu():
|
||||
pytest.skip("CPU only supports V1")
|
||||
|
||||
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_classify(math_step_prompts)
|
||||
|
||||
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
|
||||
hf_model = step_reward_patch_hf_model(hf_model)
|
||||
hf_outputs = hf_model.reward(math_step_prompts)
|
||||
|
||||
dump_reward_outputs(
|
||||
hf_outputs,
|
||||
FIXTURE_REWARD_RESULT[model],
|
||||
)
|
||||
|
||||
# check logits difference
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
hf_output = torch.tensor(hf_output).float()
|
||||
vllm_output = torch.tensor(vllm_output).float()
|
||||
|
||||
assert torch.allclose(hf_output, vllm_output, 1.5e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
pytest.param(
|
||||
"Qwen/Qwen2.5-Math-PRM-7B",
|
||||
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
def test_prm_models_with_golden_outputs(
|
||||
vllm_runner,
|
||||
math_step_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
if not FIXTURE_REWARD_RESULT.get(model):
|
||||
pytest.skip(f"No available golden outputs for {model}.")
|
||||
|
||||
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_classify(math_step_prompts)
|
||||
|
||||
golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model])
|
||||
|
||||
# check logits difference
|
||||
for golden_output, vllm_output in zip(golden_outputs, vllm_outputs):
|
||||
golden_output = torch.tensor(golden_output).float()
|
||||
vllm_output = torch.tensor(vllm_output).float()
|
||||
|
||||
assert torch.allclose(golden_output, vllm_output, 1.5e-2)
|
||||
@@ -0,0 +1,93 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm.model_executor.models.bert import (
|
||||
BertMLMHead,
|
||||
SPLADESparsePooler,
|
||||
)
|
||||
from vllm.pooling_params import PoolingParams
|
||||
from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# Functional test: SPLADE formula correctness (no HF download needed)
|
||||
# ---------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("B,T,H,V", [(2, 3, 5, 7)])
|
||||
@torch.inference_mode
|
||||
def test_splade_pooler_matches_reference_formula(B, T, H, V):
|
||||
"""Ensure SPLADESparsePooler forward() matches the mathematical formula:
|
||||
log1p(relu(logits)) -> max over sequence length (after masking)."""
|
||||
torch.manual_seed(0)
|
||||
|
||||
# Prepare [B] sequences of shape [T, H]
|
||||
hs_list = [torch.randn(T, H) for _ in range(B)]
|
||||
hs_tenser = torch.cat(hs_list)
|
||||
|
||||
# Simulate PoolingMetadata (only required fields)
|
||||
prompt_lens = [T, T - 1]
|
||||
prompt_lens_tenser = torch.tensor(prompt_lens, dtype=torch.int32)
|
||||
token_ids = torch.tensor(
|
||||
[
|
||||
[101, 5, 102], # Batch 0: [CLS], token, [SEP]
|
||||
[101, 6, 6], # Batch 1: [CLS], token, token (last token ignored)
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
meta = PoolingMetadata(
|
||||
prompt_lens=prompt_lens_tenser,
|
||||
prompt_token_ids=token_ids,
|
||||
prompt_token_ids_cpu=token_ids,
|
||||
pooling_params=[PoolingParams(task="embed")] * B,
|
||||
pooling_states=[PoolingStates() for _ in range(B)],
|
||||
)
|
||||
|
||||
# MLM head (prefer BertMLMHead, fallback to Linear if unavailable)
|
||||
try:
|
||||
mlm_head = BertMLMHead(hidden_size=H, vocab_size=V, layer_norm_eps=1e-12)
|
||||
except Exception:
|
||||
mlm_head = nn.Linear(H, V, bias=True)
|
||||
|
||||
# Forward pass through SPLADE pooler
|
||||
pooler = SPLADESparsePooler(mlm_head=mlm_head, pooling="max", remove_cls_sep=True)
|
||||
pooled = pooler(hidden_states=hs_tenser, pooling_metadata=meta) # list of [V]
|
||||
|
||||
# Basic output checks
|
||||
assert isinstance(pooled, torch.Tensor) and len(pooled) == B
|
||||
for vec in pooled:
|
||||
assert vec.shape == (V,)
|
||||
assert torch.isfinite(vec).all()
|
||||
assert (vec >= 0).all(), "SPLADE outputs must be non-negative."
|
||||
|
||||
# Reference implementation for comparison
|
||||
def ref_one(hs: torch.Tensor, L: int, tid_row: torch.Tensor) -> torch.Tensor:
|
||||
keep = torch.ones(L, dtype=torch.bool)
|
||||
if L > 0 and tid_row[0].item() == 101: # remove CLS
|
||||
keep[0] = False
|
||||
if L > 0 and tid_row[L - 1].item() == 102: # remove SEP
|
||||
keep[L - 1] = False
|
||||
|
||||
valid = hs[:L][keep[:L]]
|
||||
if valid.numel() == 0:
|
||||
return torch.zeros(V, dtype=torch.float32)
|
||||
|
||||
logits = mlm_head(valid) # [L', V]
|
||||
scores = torch.log1p(torch.relu(logits)) # [L', V]
|
||||
return scores.max(dim=0).values.to(torch.float32)
|
||||
|
||||
torch.testing.assert_close(
|
||||
pooled[0],
|
||||
ref_one(hs_list[0], prompt_lens[0], token_ids[0]),
|
||||
rtol=1e-4,
|
||||
atol=1e-4,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
pooled[1],
|
||||
ref_one(hs_list[1], prompt_lens[1], token_ids[1]),
|
||||
rtol=1e-4,
|
||||
atol=1e-4,
|
||||
)
|
||||
@@ -0,0 +1,198 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoModelForTokenClassification
|
||||
|
||||
from tests.models.registry import HF_EXAMPLE_MODELS
|
||||
from tests.models.utils import softmax
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def seed_everything():
|
||||
"""Seed all random number generators for reproducibility."""
|
||||
seed = 0
|
||||
set_random_seed(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"boltuix/NeuroBERT-NER",
|
||||
],
|
||||
)
|
||||
# The float32 is required for this tiny model to pass the test.
|
||||
@pytest.mark.parametrize("dtype", ["float"])
|
||||
@torch.inference_mode
|
||||
def test_bert_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_classify(example_prompts)
|
||||
|
||||
# Use eager attention on ROCm to avoid HF Transformers flash attention
|
||||
# accuracy issues: https://github.com/vllm-project/vllm/issues/30167
|
||||
hf_model_kwargs = {}
|
||||
if current_platform.is_rocm():
|
||||
hf_model_kwargs["attn_implementation"] = "eager"
|
||||
|
||||
with hf_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
auto_cls=AutoModelForTokenClassification,
|
||||
model_kwargs=hf_model_kwargs,
|
||||
) as hf_model:
|
||||
tokenizer = hf_model.tokenizer
|
||||
hf_outputs = []
|
||||
for prompt in example_prompts:
|
||||
inputs = tokenizer([prompt], return_tensors="pt")
|
||||
inputs = hf_model.wrap_device(inputs)
|
||||
output = hf_model.model(**inputs)
|
||||
hf_outputs.append(softmax(output.logits[0]))
|
||||
|
||||
# check logits difference
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
hf_output = hf_output.detach().clone().cpu().float()
|
||||
vllm_output = vllm_output.detach().clone().cpu().float()
|
||||
torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["disham993/electrical-ner-ModernBERT-base"])
|
||||
@pytest.mark.parametrize("dtype", ["float"])
|
||||
@pytest.mark.flaky(reruns=3)
|
||||
@torch.inference_mode
|
||||
def test_modernbert_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
# NOTE: https://github.com/vllm-project/vllm/pull/32403
|
||||
# `disham993/electrical-ner-ModernBERT-base` is a randomly initialized
|
||||
# model, which can cause numerical precision variance and edge cases.
|
||||
# We use @flaky(reruns=3) to mitigate intermittent failures.
|
||||
print(
|
||||
f"\n[NOTE] Testing {model} (randomly initialized weights) - "
|
||||
"flaky tolerance enabled due to numerical precision variance."
|
||||
)
|
||||
|
||||
with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_classify(example_prompts)
|
||||
|
||||
# Use eager attention on ROCm to avoid HF Transformers flash attention
|
||||
# accuracy issues: https://github.com/vllm-project/vllm/issues/30167
|
||||
hf_model_kwargs = {}
|
||||
if current_platform.is_rocm():
|
||||
hf_model_kwargs["attn_implementation"] = "eager"
|
||||
|
||||
with hf_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
auto_cls=AutoModelForTokenClassification,
|
||||
model_kwargs=hf_model_kwargs,
|
||||
) as hf_model:
|
||||
tokenizer = hf_model.tokenizer
|
||||
hf_outputs = []
|
||||
for prompt in example_prompts:
|
||||
inputs = tokenizer([prompt], return_tensors="pt")
|
||||
inputs = hf_model.wrap_device(inputs)
|
||||
output = hf_model.model(**inputs)
|
||||
hf_outputs.append(softmax(output.logits[0]))
|
||||
|
||||
# check logits difference
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
hf_output = hf_output.detach().clone().cpu().float()
|
||||
vllm_output = vllm_output.detach().clone().cpu().float()
|
||||
torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
|
||||
|
||||
|
||||
PRIVACY_FILTER_PROMPTS = [
|
||||
"My name is Harry Potter.",
|
||||
"Email me at harry.potter@hogwarts.edu.",
|
||||
"Call me on +44 20 7946 0958 tomorrow.",
|
||||
"My account number is 12345678 and the API key is sk-live-abc123def456.",
|
||||
"I live at 4 Privet Drive, Little Whinging.",
|
||||
"Visit https://example.com/profile/harry for more info.",
|
||||
"We met on 12 January 2024.",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["openai/privacy-filter"])
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
@torch.inference_mode
|
||||
def test_openai_privacy_filter(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_classify(PRIVACY_FILTER_PROMPTS)
|
||||
|
||||
hf_model_kwargs = {}
|
||||
if current_platform.is_rocm():
|
||||
hf_model_kwargs["attn_implementation"] = "eager"
|
||||
|
||||
with hf_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
auto_cls=AutoModelForTokenClassification,
|
||||
model_kwargs=hf_model_kwargs,
|
||||
) as hf_model:
|
||||
tokenizer = hf_model.tokenizer
|
||||
hf_outputs = []
|
||||
for prompt in PRIVACY_FILTER_PROMPTS:
|
||||
inputs = tokenizer([prompt], return_tensors="pt")
|
||||
inputs = hf_model.wrap_device(inputs)
|
||||
output = hf_model.model(**inputs)
|
||||
hf_outputs.append(softmax(output.logits[0]))
|
||||
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
hf_output = hf_output.detach().clone().cpu().float()
|
||||
vllm_output = vllm_output.detach().clone().cpu().float()
|
||||
torch.testing.assert_close(hf_output, vllm_output, atol=0.1, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["bd2lcco/Qwen3-0.6B-finetuned"])
|
||||
@pytest.mark.parametrize("dtype", ["float"])
|
||||
@torch.inference_mode
|
||||
def test_auto_conversion(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_classify(example_prompts)
|
||||
|
||||
with hf_runner(
|
||||
model, dtype=dtype, auto_cls=AutoModelForTokenClassification
|
||||
) as hf_model:
|
||||
tokenizer = hf_model.tokenizer
|
||||
hf_outputs = []
|
||||
for prompt in example_prompts:
|
||||
inputs = tokenizer([prompt], return_tensors="pt")
|
||||
inputs = hf_model.wrap_device(inputs)
|
||||
output = hf_model.model(**inputs)
|
||||
hf_outputs.append(softmax(output.logits[0]))
|
||||
|
||||
# check logits difference
|
||||
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
||||
hf_output = hf_output.detach().clone().cpu().float()
|
||||
vllm_output = vllm_output.detach().clone().cpu().float()
|
||||
assert torch.allclose(hf_output, vllm_output, atol=1e-2)
|
||||
@@ -0,0 +1,79 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
MODEL_NAME = "sentence-transformers/all-MiniLM-L12-v2"
|
||||
max_model_len = 128
|
||||
|
||||
input_str = """Immerse yourself in the enchanting chronicle of calculus, a
|
||||
mathematical domain that has radically transformed our comprehension of
|
||||
change and motion. Despite its roots in ancient civilizations, the
|
||||
formal birth of calculus predominantly occurred in the 17th century,
|
||||
primarily under the influential guidance of Sir Isaac Newton and Gottfried
|
||||
Wilhelm Leibniz. The earliest traces of calculus concepts are found in
|
||||
ancient Greek mathematics,most notably in the works of Eudoxus and
|
||||
Archimedes, around 300 BCE. They utilized the 'method of exhaustion'—a
|
||||
technique for computing areas and volumes through the use of finite sums.
|
||||
This methodology laid crucial foundational work for integral calculus.
|
||||
In the 17th century, both Newton and Leibniz independently pioneered
|
||||
calculus, each contributing unique perspectives that would shape this new
|
||||
field."""
|
||||
|
||||
|
||||
def test_smaller_truncation_size(
|
||||
vllm_runner, model_name=MODEL_NAME, input_str=input_str
|
||||
):
|
||||
truncate_prompt_tokens = 10
|
||||
|
||||
with vllm_runner(
|
||||
model_name, runner="pooling", max_model_len=max_model_len
|
||||
) as vllm_model:
|
||||
vllm_output = vllm_model.llm.embed(
|
||||
input_str,
|
||||
tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
|
||||
)
|
||||
|
||||
prompt_tokens = vllm_output[0].prompt_token_ids
|
||||
|
||||
assert len(prompt_tokens) == truncate_prompt_tokens
|
||||
|
||||
|
||||
def test_max_truncation_size(vllm_runner, model_name=MODEL_NAME, input_str=input_str):
|
||||
truncate_prompt_tokens = -1
|
||||
|
||||
with vllm_runner(
|
||||
model_name, runner="pooling", max_model_len=max_model_len
|
||||
) as vllm_model:
|
||||
vllm_output = vllm_model.llm.embed(
|
||||
input_str,
|
||||
tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
|
||||
)
|
||||
|
||||
prompt_tokens = vllm_output[0].prompt_token_ids
|
||||
|
||||
assert len(prompt_tokens) == max_model_len
|
||||
|
||||
|
||||
def test_bigger_truncation_size(
|
||||
vllm_runner, model_name=MODEL_NAME, input_str=input_str
|
||||
):
|
||||
truncate_prompt_tokens = max_model_len + 1
|
||||
|
||||
with (
|
||||
pytest.raises(ValueError),
|
||||
vllm_runner(
|
||||
model_name, runner="pooling", max_model_len=max_model_len
|
||||
) as vllm_model,
|
||||
):
|
||||
llm_output = vllm_model.llm.embed(
|
||||
input_str,
|
||||
tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
|
||||
)
|
||||
|
||||
assert (
|
||||
llm_output
|
||||
== f"""truncate_prompt_tokens value
|
||||
({truncate_prompt_tokens}) is greater than
|
||||
max_model_len ({max_model_len}). Please, select
|
||||
a smaller truncation size."""
|
||||
)
|
||||
@@ -0,0 +1,258 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import mteb
|
||||
import numpy as np
|
||||
import torch
|
||||
from mteb.models import ModelMeta
|
||||
from mteb.types import Array
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
import tests.ci_envs as ci_envs
|
||||
from tests.models.utils import (
|
||||
EmbedModelInfo,
|
||||
check_embeddings_close,
|
||||
get_vllm_extra_kwargs,
|
||||
)
|
||||
|
||||
# Most embedding models on the STS12 task (See #17175):
|
||||
# - Model implementation and minor changes in tensor dtype
|
||||
# results in differences less than 1e-4
|
||||
# - Different model results in differences more than 1e-3
|
||||
# 5e-4 is a good tolerance threshold
|
||||
MTEB_EMBED_TASKS = ["STS12"]
|
||||
MTEB_EMBED_TOL = 5e-4
|
||||
|
||||
|
||||
_empty_model_meta = ModelMeta(
|
||||
loader=None,
|
||||
name="vllm/model",
|
||||
revision="1",
|
||||
release_date=None,
|
||||
languages=None,
|
||||
framework=[],
|
||||
similarity_fn_name=None,
|
||||
n_parameters=None,
|
||||
memory_usage_mb=None,
|
||||
max_tokens=None,
|
||||
embed_dim=None,
|
||||
license=None,
|
||||
open_weights=None,
|
||||
public_training_code=None,
|
||||
public_training_data=None,
|
||||
use_instructions=None,
|
||||
training_datasets=None,
|
||||
modalities=["text"], # 'image' can be added to evaluate multimodal models
|
||||
)
|
||||
|
||||
|
||||
class MtebEmbedMixin(mteb.EncoderProtocol):
|
||||
mteb_model_meta = _empty_model_meta
|
||||
|
||||
def similarity(
|
||||
self,
|
||||
embeddings1: np.ndarray,
|
||||
embeddings2: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
# Cosine similarity
|
||||
norm1 = np.linalg.norm(embeddings1, axis=1, keepdims=True)
|
||||
norm2 = np.linalg.norm(embeddings2, axis=1, keepdims=True)
|
||||
sim = np.dot(embeddings1, embeddings2.T) / (norm1 * norm2.T)
|
||||
return sim
|
||||
|
||||
def similarity_pairwise(
|
||||
self,
|
||||
embeddings1: Array,
|
||||
embeddings2: Array,
|
||||
) -> Array:
|
||||
# Cosine similarity
|
||||
norm1 = np.linalg.norm(embeddings1, axis=1, keepdims=True)
|
||||
norm2 = np.linalg.norm(embeddings2, axis=1, keepdims=True)
|
||||
sim = np.sum(embeddings1 * embeddings2, axis=1) / (
|
||||
norm1.flatten() * norm2.flatten()
|
||||
)
|
||||
return sim
|
||||
|
||||
|
||||
class HfMtebEncoder(MtebEmbedMixin):
|
||||
def __init__(self, model):
|
||||
self.model = model
|
||||
|
||||
def encode(
|
||||
self,
|
||||
inputs: DataLoader[mteb.types.BatchedInput],
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> np.ndarray:
|
||||
sentences = [text for batch in inputs for text in batch["text"]]
|
||||
return self.model.encode(sentences)
|
||||
|
||||
|
||||
class VllmMtebEncoder(MtebEmbedMixin):
|
||||
def __init__(self, vllm_model, prompt_prefix: str | None = None):
|
||||
self.llm = vllm_model
|
||||
self.rng = np.random.default_rng(seed=42)
|
||||
self.prompt_prefix = prompt_prefix
|
||||
|
||||
def encode(
|
||||
self,
|
||||
inputs: DataLoader[mteb.types.BatchedInput],
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> np.ndarray:
|
||||
# Hoping to discover potential scheduling
|
||||
# issues by randomizing the order.
|
||||
sentences = [
|
||||
self.prompt_prefix + text if self.prompt_prefix else text
|
||||
for batch in inputs
|
||||
for text in batch["text"]
|
||||
]
|
||||
r = self.rng.permutation(len(sentences))
|
||||
sentences = [sentences[i] for i in r]
|
||||
outputs = self.llm.embed(sentences, use_tqdm=False)
|
||||
embeds = np.array(outputs)
|
||||
embeds = embeds[np.argsort(r)]
|
||||
return embeds
|
||||
|
||||
|
||||
class OpenAIClientMtebEncoder(MtebEmbedMixin):
|
||||
def __init__(self, model_name: str, client):
|
||||
self.model_name = model_name
|
||||
self.client = client
|
||||
self.rng = np.random.default_rng(seed=42)
|
||||
|
||||
def encode(
|
||||
self,
|
||||
inputs: DataLoader[mteb.types.BatchedInput],
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> np.ndarray:
|
||||
# Hoping to discover potential scheduling
|
||||
# issues by randomizing the order.
|
||||
sentences = [text for batch in inputs for text in batch["text"]]
|
||||
r = self.rng.permutation(len(sentences))
|
||||
sentences = [sentences[i] for i in r]
|
||||
|
||||
embeddings = self.client.embeddings.create(
|
||||
model=self.model_name, input=sentences
|
||||
)
|
||||
outputs = [d.embedding for d in embeddings.data]
|
||||
embeds = np.array(outputs)
|
||||
embeds = embeds[np.argsort(r)]
|
||||
return embeds
|
||||
|
||||
|
||||
def run_mteb_embed_task(encoder: mteb.EncoderProtocol, tasks):
|
||||
tasks = mteb.get_tasks(tasks=tasks)
|
||||
results = mteb.evaluate(
|
||||
encoder,
|
||||
tasks,
|
||||
cache=None,
|
||||
show_progress_bar=False,
|
||||
)
|
||||
|
||||
main_score = results[0].scores["test"][0]["main_score"]
|
||||
return main_score
|
||||
|
||||
|
||||
def mteb_test_embed_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model_info: EmbedModelInfo,
|
||||
vllm_extra_kwargs=None,
|
||||
hf_model_callback=None,
|
||||
atol=MTEB_EMBED_TOL,
|
||||
prompt_prefix: str | None = None,
|
||||
):
|
||||
vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
|
||||
|
||||
# Test embed_dims, isnan and whether to use normalize
|
||||
example_prompts = ["The chef prepared a delicious meal." * 1000]
|
||||
|
||||
with vllm_runner(
|
||||
model_info.name,
|
||||
revision=model_info.revision,
|
||||
runner="pooling",
|
||||
max_model_len=model_info.max_model_len,
|
||||
**vllm_extra_kwargs,
|
||||
) as vllm_model:
|
||||
model_config = vllm_model.llm.llm_engine.model_config
|
||||
|
||||
# Confirm whether vllm is using the correct architecture
|
||||
if model_info.architecture:
|
||||
assert model_info.architecture in model_config.architectures
|
||||
|
||||
# Confirm whether the important configs in model_config are correct.
|
||||
pooler_config = model_config.pooler_config
|
||||
if model_info.seq_pooling_type is not None:
|
||||
assert pooler_config.seq_pooling_type == model_info.seq_pooling_type
|
||||
if model_info.tok_pooling_type is not None:
|
||||
assert pooler_config.tok_pooling_type == model_info.tok_pooling_type
|
||||
if model_info.attn_type is not None:
|
||||
assert model_config.attn_type == model_info.attn_type
|
||||
if model_info.is_prefix_caching_supported is not None:
|
||||
assert (
|
||||
model_config.is_prefix_caching_supported
|
||||
== model_info.is_prefix_caching_supported
|
||||
)
|
||||
if model_info.is_chunked_prefill_supported is not None:
|
||||
assert (
|
||||
model_config.is_chunked_prefill_supported
|
||||
== model_info.is_chunked_prefill_supported
|
||||
)
|
||||
|
||||
vllm_main_score = run_mteb_embed_task(
|
||||
VllmMtebEncoder(vllm_model, prompt_prefix=prompt_prefix), MTEB_EMBED_TASKS
|
||||
)
|
||||
vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype
|
||||
head_dtype = model_config.head_dtype
|
||||
|
||||
# Test embedding_size, isnan and whether to use normalize
|
||||
vllm_outputs = vllm_model.embed(
|
||||
example_prompts,
|
||||
tokenization_kwargs=dict(truncate_prompt_tokens=-1),
|
||||
)
|
||||
outputs_tensor = torch.tensor(vllm_outputs)
|
||||
assert not torch.any(torch.isnan(outputs_tensor))
|
||||
embedding_size = model_config.embedding_size
|
||||
assert torch.tensor(vllm_outputs).shape[-1] == embedding_size
|
||||
|
||||
# Accelerate mteb test by setting
|
||||
# SentenceTransformers mteb score to a constant
|
||||
if model_info.mteb_score is None:
|
||||
with hf_runner(
|
||||
model_info.name,
|
||||
revision=model_info.revision,
|
||||
is_sentence_transformer=True,
|
||||
dtype=ci_envs.VLLM_CI_HF_DTYPE or model_info.hf_dtype,
|
||||
) as hf_model:
|
||||
# e.g. setting default parameters for the encode method of hf_runner
|
||||
if hf_model_callback is not None:
|
||||
hf_model_callback(hf_model)
|
||||
|
||||
st_main_score = run_mteb_embed_task(
|
||||
HfMtebEncoder(hf_model), MTEB_EMBED_TASKS
|
||||
)
|
||||
st_dtype = next(hf_model.model.parameters()).dtype
|
||||
|
||||
# Check embeddings close to hf outputs
|
||||
hf_outputs = hf_model.encode(example_prompts)
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=hf_outputs,
|
||||
embeddings_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
tol=1e-2,
|
||||
)
|
||||
else:
|
||||
st_main_score = model_info.mteb_score
|
||||
st_dtype = "Constant"
|
||||
|
||||
print("Model:", model_info.name)
|
||||
print("VLLM:", f"dtype:{vllm_dtype}", f"head_dtype:{head_dtype}", vllm_main_score)
|
||||
print("SentenceTransformers:", st_dtype, st_main_score)
|
||||
print("Difference:", st_main_score - vllm_main_score)
|
||||
|
||||
# We are not concerned that the vllm mteb results are better
|
||||
# than SentenceTransformers, so we only perform one-sided testing.
|
||||
assert st_main_score - vllm_main_score < atol
|
||||
@@ -0,0 +1,311 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import mteb
|
||||
import numpy as np
|
||||
import requests
|
||||
import torch
|
||||
from mteb.models import ModelMeta
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from tests.conftest import HfRunner
|
||||
from tests.models.utils import (
|
||||
RerankModelInfo,
|
||||
get_vllm_extra_kwargs,
|
||||
)
|
||||
|
||||
# See #19344
|
||||
MTEB_RERANK_TASKS = ["NFCorpus"]
|
||||
MTEB_RERANK_LANGS = ["eng"]
|
||||
MTEB_RERANK_TOL = 2e-3
|
||||
|
||||
template_home = (
|
||||
Path(__file__).parent.parent.parent.parent.parent
|
||||
/ "examples/pooling/score/template"
|
||||
)
|
||||
|
||||
_empty_model_meta = ModelMeta(
|
||||
loader=None,
|
||||
name="vllm/model",
|
||||
revision="1",
|
||||
release_date=None,
|
||||
languages=None,
|
||||
framework=[],
|
||||
similarity_fn_name=None,
|
||||
n_parameters=None,
|
||||
memory_usage_mb=None,
|
||||
max_tokens=None,
|
||||
embed_dim=None,
|
||||
license=None,
|
||||
open_weights=None,
|
||||
public_training_code=None,
|
||||
public_training_data=None,
|
||||
use_instructions=None,
|
||||
training_datasets=None,
|
||||
modalities=["text"], # 'image' can be added to evaluate multimodal models
|
||||
)
|
||||
|
||||
|
||||
class MtebCrossEncoderMixin(mteb.CrossEncoderProtocol):
|
||||
mteb_model_meta = _empty_model_meta
|
||||
|
||||
|
||||
class VllmMtebCrossEncoder(MtebCrossEncoderMixin):
|
||||
def __init__(self, vllm_model):
|
||||
self.llm = vllm_model
|
||||
self.rng = np.random.default_rng(seed=42)
|
||||
self.chat_template: str | None = getattr(vllm_model, "chat_template", None)
|
||||
|
||||
def predict(
|
||||
self,
|
||||
inputs1: DataLoader[mteb.types.BatchedInput],
|
||||
inputs2: DataLoader[mteb.types.BatchedInput],
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> np.ndarray:
|
||||
queries = [text for batch in inputs1 for text in batch["text"]]
|
||||
corpus = [text for batch in inputs2 for text in batch["text"]]
|
||||
|
||||
# Hoping to discover potential scheduling
|
||||
# issues by randomizing the order.
|
||||
r = self.rng.permutation(len(queries))
|
||||
queries = [queries[i] for i in r]
|
||||
corpus = [corpus[i] for i in r]
|
||||
|
||||
outputs = self.llm.score(
|
||||
queries,
|
||||
corpus,
|
||||
use_tqdm=False,
|
||||
chat_template=self.chat_template,
|
||||
tokenization_kwargs={"truncate_prompt_tokens": -1},
|
||||
)
|
||||
scores = np.array(outputs)
|
||||
scores = scores[np.argsort(r)]
|
||||
return scores
|
||||
|
||||
|
||||
class ScoreClientMtebEncoder(MtebCrossEncoderMixin):
|
||||
mteb_model_meta = _empty_model_meta
|
||||
|
||||
def __init__(self, model_name: str, url):
|
||||
self.model_name = model_name
|
||||
self.url = url
|
||||
|
||||
def predict(
|
||||
self,
|
||||
inputs1: DataLoader[mteb.types.BatchedInput],
|
||||
inputs2: DataLoader[mteb.types.BatchedInput],
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> np.ndarray:
|
||||
queries = [text for batch in inputs1 for text in batch["text"]]
|
||||
full_corpus = [text for batch in inputs2 for text in batch["text"]]
|
||||
|
||||
outputs = []
|
||||
for query, corpus in zip(queries, full_corpus):
|
||||
outputs.append(self.get_score(query, corpus))
|
||||
|
||||
scores = np.array(outputs)
|
||||
return scores
|
||||
|
||||
def get_score(self, query, corpus):
|
||||
response = requests.post(
|
||||
self.url,
|
||||
json={
|
||||
"model": self.model_name,
|
||||
"queries": query,
|
||||
"documents": corpus,
|
||||
"truncate_prompt_tokens": -1,
|
||||
},
|
||||
).json()
|
||||
return response["data"][0]["score"]
|
||||
|
||||
|
||||
class RerankClientMtebEncoder(ScoreClientMtebEncoder):
|
||||
def get_score(self, query, corpus):
|
||||
response = requests.post(
|
||||
self.url,
|
||||
json={
|
||||
"model": self.model_name,
|
||||
"query": query,
|
||||
"documents": [corpus],
|
||||
"truncate_prompt_tokens": -1,
|
||||
},
|
||||
).json()
|
||||
return response["results"][0]["relevance_score"]
|
||||
|
||||
|
||||
class HFMtebCrossEncoder(MtebCrossEncoderMixin, HfRunner):
|
||||
chat_template: str | None = None
|
||||
|
||||
def __init__(self, model_name: str, dtype: str = "auto", **kwargs: Any) -> None:
|
||||
HfRunner.__init__(
|
||||
self, model_name=model_name, is_cross_encoder=True, dtype=dtype, **kwargs
|
||||
)
|
||||
|
||||
@torch.no_grad
|
||||
def predict(
|
||||
self,
|
||||
inputs1: DataLoader[mteb.types.BatchedInput],
|
||||
inputs2: DataLoader[mteb.types.BatchedInput],
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> np.ndarray:
|
||||
queries = [text for batch in inputs1 for text in batch["text"]]
|
||||
corpus = [text for batch in inputs2 for text in batch["text"]]
|
||||
|
||||
if self.chat_template is not None:
|
||||
tokenizer = self.model.tokenizer
|
||||
prompts = []
|
||||
for query, document in zip(queries, corpus):
|
||||
conversation = [
|
||||
{"role": "query", "content": query},
|
||||
{"role": "document", "content": document},
|
||||
]
|
||||
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
conversation=conversation,
|
||||
tools=None,
|
||||
chat_template=self.chat_template,
|
||||
tokenize=False,
|
||||
)
|
||||
prompts.append(prompt)
|
||||
outputs_list = HfRunner.classify(self, prompts)
|
||||
scores = np.array(outputs_list).squeeze(-1)
|
||||
return scores
|
||||
else:
|
||||
prompts = list(zip(queries, corpus))
|
||||
outputs_tensor = HfRunner.predict(self, prompts, show_progress_bar=False)
|
||||
return outputs_tensor.cpu().numpy()
|
||||
|
||||
|
||||
def run_mteb_rerank(cross_encoder: mteb.CrossEncoderProtocol, tasks, languages):
|
||||
with tempfile.TemporaryDirectory() as prediction_folder:
|
||||
bm25s = mteb.get_model("bm25s")
|
||||
eval_splits = ["test"]
|
||||
|
||||
mteb_tasks: list[mteb.abstasks.AbsTaskRetrieval] = mteb.get_tasks(
|
||||
tasks=tasks, languages=languages, eval_splits=eval_splits
|
||||
)
|
||||
for task in mteb_tasks:
|
||||
if not task.data_loaded:
|
||||
task.load_data()
|
||||
|
||||
mteb.evaluate(
|
||||
bm25s,
|
||||
mteb_tasks,
|
||||
prediction_folder=prediction_folder,
|
||||
show_progress_bar=False,
|
||||
# don't save results for test runs
|
||||
cache=None,
|
||||
overwrite_strategy="always",
|
||||
)
|
||||
|
||||
second_stage_tasks = []
|
||||
for task in mteb_tasks:
|
||||
second_stage_tasks.append(
|
||||
task.convert_to_reranking(
|
||||
prediction_folder,
|
||||
top_k=10,
|
||||
)
|
||||
)
|
||||
|
||||
results = mteb.evaluate(
|
||||
cross_encoder,
|
||||
second_stage_tasks,
|
||||
show_progress_bar=False,
|
||||
cache=None,
|
||||
)
|
||||
main_score = results[0].scores["test"][0]["main_score"]
|
||||
return main_score
|
||||
|
||||
|
||||
def mteb_test_rerank_models(
|
||||
vllm_runner,
|
||||
model_info: RerankModelInfo,
|
||||
hf_runner=HFMtebCrossEncoder,
|
||||
vllm_extra_kwargs=None,
|
||||
vllm_mteb_encoder=VllmMtebCrossEncoder,
|
||||
atol=MTEB_RERANK_TOL,
|
||||
):
|
||||
vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
|
||||
|
||||
# Maybe load chat_template.
|
||||
chat_template: str | None = None
|
||||
if model_info.chat_template_name is not None:
|
||||
chat_template = (template_home / model_info.chat_template_name).read_text()
|
||||
|
||||
with vllm_runner(
|
||||
model_info.name,
|
||||
revision=model_info.revision,
|
||||
runner="pooling",
|
||||
max_model_len=None,
|
||||
max_num_seqs=8,
|
||||
**vllm_extra_kwargs,
|
||||
) as vllm_model:
|
||||
model_config = vllm_model.llm.llm_engine.model_config
|
||||
vllm_model.chat_template = chat_template
|
||||
|
||||
# Confirm whether vllm is using the correct architecture
|
||||
if model_info.architecture:
|
||||
assert model_info.architecture in model_config.architectures
|
||||
|
||||
# Score API is only enabled for num_labels == 1
|
||||
assert model_config.hf_config.num_labels == 1
|
||||
|
||||
# Confirm whether the important configs in model_config are correct.
|
||||
pooler_config = model_config.pooler_config
|
||||
if model_info.seq_pooling_type is not None:
|
||||
assert pooler_config.seq_pooling_type == model_info.seq_pooling_type
|
||||
if model_info.tok_pooling_type is not None:
|
||||
assert pooler_config.tok_pooling_type == model_info.tok_pooling_type
|
||||
if model_info.attn_type is not None:
|
||||
assert model_config.attn_type == model_info.attn_type
|
||||
if model_info.is_prefix_caching_supported is not None:
|
||||
assert (
|
||||
model_config.is_prefix_caching_supported
|
||||
== model_info.is_prefix_caching_supported
|
||||
)
|
||||
if model_info.is_chunked_prefill_supported is not None:
|
||||
assert (
|
||||
model_config.is_chunked_prefill_supported
|
||||
== model_info.is_chunked_prefill_supported
|
||||
)
|
||||
|
||||
vllm_main_score = run_mteb_rerank(
|
||||
vllm_mteb_encoder(vllm_model),
|
||||
tasks=MTEB_RERANK_TASKS,
|
||||
languages=MTEB_RERANK_LANGS,
|
||||
)
|
||||
vllm_dtype = model_config.dtype
|
||||
head_dtype = model_config.head_dtype
|
||||
|
||||
# Accelerate mteb test by setting
|
||||
# SentenceTransformers mteb score to a constant
|
||||
if model_info.mteb_score is None:
|
||||
with hf_runner(
|
||||
model_info.name, revision=model_info.revision, dtype=model_info.hf_dtype
|
||||
) as hf_model:
|
||||
hf_model.chat_template = chat_template
|
||||
st_main_score = run_mteb_rerank(
|
||||
hf_model,
|
||||
tasks=MTEB_RERANK_TASKS,
|
||||
languages=MTEB_RERANK_LANGS,
|
||||
)
|
||||
st_dtype = next(hf_model.model.model.parameters()).dtype
|
||||
else:
|
||||
st_main_score = model_info.mteb_score
|
||||
st_dtype = "Constant"
|
||||
|
||||
print("Model:", model_info.name)
|
||||
print("VLLM:", f"dtype:{vllm_dtype}", f"head_dtype:{head_dtype}", vllm_main_score)
|
||||
print("SentenceTransformers:", st_dtype, st_main_score)
|
||||
print("Difference:", st_main_score - vllm_main_score)
|
||||
|
||||
# We are not concerned that the vllm mteb results are better
|
||||
# than SentenceTransformers, so we only perform one-sided testing.
|
||||
assert st_main_score - vllm_main_score < atol
|
||||
@@ -0,0 +1,118 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
from tests.models.language.pooling.embed_utils import correctness_test_embed_models
|
||||
from tests.models.utils import (
|
||||
EmbedModelInfo,
|
||||
RerankModelInfo,
|
||||
)
|
||||
|
||||
from .mteb_embed_utils import mteb_test_embed_models
|
||||
from .mteb_score_utils import mteb_test_rerank_models
|
||||
|
||||
MODELS = [
|
||||
########## BertModel
|
||||
EmbedModelInfo(
|
||||
"BAAI/bge-base-en",
|
||||
architecture="BertModel",
|
||||
mteb_score=0.779336792,
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo("BAAI/bge-base-zh", architecture="BertModel", enable_test=False),
|
||||
EmbedModelInfo("BAAI/bge-small-en", architecture="BertModel", enable_test=False),
|
||||
EmbedModelInfo("BAAI/bge-small-zh", architecture="BertModel", enable_test=False),
|
||||
EmbedModelInfo("BAAI/bge-large-en", architecture="BertModel", enable_test=False),
|
||||
EmbedModelInfo("BAAI/bge-large-zh", architecture="BertModel", enable_test=False),
|
||||
EmbedModelInfo(
|
||||
"BAAI/bge-large-zh-noinstruct", architecture="BertModel", enable_test=False
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"BAAI/bge-base-en-v1.5", architecture="BertModel", enable_test=False
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"BAAI/bge-base-zh-v1.5", architecture="BertModel", enable_test=False
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"BAAI/bge-small-en-v1.5", architecture="BertModel", enable_test=False
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"BAAI/bge-small-zh-v1.5", architecture="BertModel", enable_test=False
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"BAAI/bge-large-en-v1.5", architecture="BertModel", enable_test=False
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"BAAI/bge-large-zh-v1.5", architecture="BertModel", enable_test=False
|
||||
),
|
||||
########## XLMRobertaModel
|
||||
EmbedModelInfo(
|
||||
"BAAI/bge-m3",
|
||||
architecture="XLMRobertaModel",
|
||||
mteb_score=0.787343078,
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
########## Qwen2Model
|
||||
EmbedModelInfo(
|
||||
"BAAI/bge-code-v1",
|
||||
architecture="Qwen2Model",
|
||||
mteb_score=0.75724465,
|
||||
seq_pooling_type="LAST",
|
||||
attn_type="decoder",
|
||||
is_prefix_caching_supported=True,
|
||||
is_chunked_prefill_supported=True,
|
||||
# Skip: model's custom tokenizer on HF hub is incompatible with
|
||||
# transformers v5 (sets attrs before super().__init__, triggering
|
||||
# AttributeError on 'verbose' in __getattr__).
|
||||
enable_test=False,
|
||||
),
|
||||
]
|
||||
|
||||
RERANK_MODELS = [
|
||||
########## XLMRobertaForSequenceClassification
|
||||
RerankModelInfo(
|
||||
"BAAI/bge-reranker-base",
|
||||
architecture="XLMRobertaForSequenceClassification",
|
||||
mteb_score=0.32398,
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
RerankModelInfo(
|
||||
"BAAI/bge-reranker-large",
|
||||
architecture="XLMRobertaForSequenceClassification",
|
||||
enable_test=False,
|
||||
),
|
||||
RerankModelInfo(
|
||||
"BAAI/bge-reranker-v2-m3",
|
||||
architecture="XLMRobertaForSequenceClassification",
|
||||
enable_test=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
|
||||
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_correctness(
|
||||
hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
|
||||
) -> None:
|
||||
correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", RERANK_MODELS)
|
||||
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
|
||||
mteb_test_rerank_models(vllm_runner, model_info)
|
||||
@@ -0,0 +1,136 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import Any
|
||||
|
||||
import mteb
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from tests.conftest import HfRunner
|
||||
from tests.models.utils import RerankModelInfo
|
||||
|
||||
from .mteb_score_utils import (
|
||||
MtebCrossEncoderMixin,
|
||||
mteb_test_rerank_models,
|
||||
)
|
||||
|
||||
RERANK_MODELS = [
|
||||
RerankModelInfo(
|
||||
"BAAI/bge-reranker-v2-gemma",
|
||||
architecture="GemmaForSequenceClassification",
|
||||
hf_overrides={
|
||||
"architectures": ["GemmaForSequenceClassification"],
|
||||
"classifier_from_token": ["Yes"],
|
||||
"method": "no_post_processing",
|
||||
},
|
||||
mteb_score=0.33757,
|
||||
seq_pooling_type="LAST",
|
||||
attn_type="decoder",
|
||||
is_prefix_caching_supported=True,
|
||||
is_chunked_prefill_supported=True,
|
||||
chat_template_name="bge-reranker-v2-gemma.jinja",
|
||||
),
|
||||
]
|
||||
|
||||
PROMPT = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'." # noqa: E501
|
||||
|
||||
|
||||
class GemmaRerankerHfRunner(MtebCrossEncoderMixin, HfRunner):
|
||||
def __init__(
|
||||
self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any
|
||||
) -> None:
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
HfRunner.__init__(
|
||||
self,
|
||||
model_name=model_name,
|
||||
auto_cls=AutoModelForCausalLM,
|
||||
dtype=dtype,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
|
||||
self.yes_loc = self.tokenizer.convert_tokens_to_ids("Yes")
|
||||
|
||||
@torch.no_grad
|
||||
def predict(
|
||||
self,
|
||||
inputs1: DataLoader[mteb.types.BatchedInput],
|
||||
inputs2: DataLoader[mteb.types.BatchedInput],
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> np.ndarray:
|
||||
queries = [text for batch in inputs1 for text in batch["text"]]
|
||||
corpus = [text for batch in inputs2 for text in batch["text"]]
|
||||
|
||||
def get_inputs(pairs, tokenizer, prompt=None):
|
||||
if prompt is None:
|
||||
prompt = PROMPT
|
||||
|
||||
sep = "\n"
|
||||
prompt_inputs = tokenizer(
|
||||
prompt, return_tensors=None, add_special_tokens=False
|
||||
)["input_ids"]
|
||||
sep_inputs = tokenizer(sep, return_tensors=None, add_special_tokens=False)[
|
||||
"input_ids"
|
||||
]
|
||||
inputs = []
|
||||
for query, passage in pairs:
|
||||
query_inputs = tokenizer(
|
||||
f"A: {query}",
|
||||
return_tensors=None,
|
||||
add_special_tokens=False,
|
||||
truncation=True,
|
||||
)
|
||||
passage_inputs = tokenizer(
|
||||
f"B: {passage}",
|
||||
return_tensors=None,
|
||||
add_special_tokens=False,
|
||||
truncation=True,
|
||||
)
|
||||
item = tokenizer.prepare_for_model(
|
||||
[tokenizer.bos_token_id] + query_inputs["input_ids"],
|
||||
sep_inputs + passage_inputs["input_ids"],
|
||||
truncation="only_second",
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
return_token_type_ids=False,
|
||||
add_special_tokens=False,
|
||||
)
|
||||
item["input_ids"] = item["input_ids"] + sep_inputs + prompt_inputs
|
||||
item["attention_mask"] = [1] * len(item["input_ids"])
|
||||
inputs.append(item)
|
||||
return tokenizer.pad(
|
||||
inputs,
|
||||
padding=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
scores = []
|
||||
for query, document in zip(queries, corpus):
|
||||
pairs = [(query, document)]
|
||||
inputs = get_inputs(pairs, self.tokenizer)
|
||||
inputs = inputs.to(self.model.device)
|
||||
_n_tokens = inputs["input_ids"].shape[1]
|
||||
logits = self.model(**inputs, return_dict=True).logits
|
||||
_scores = (
|
||||
logits[:, -1, self.yes_loc]
|
||||
.view(
|
||||
-1,
|
||||
)
|
||||
.float()
|
||||
.sigmoid()
|
||||
)
|
||||
scores.append(_scores[0].item())
|
||||
return torch.Tensor(scores)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", RERANK_MODELS)
|
||||
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
|
||||
mteb_test_rerank_models(
|
||||
vllm_runner,
|
||||
model_info,
|
||||
hf_runner=GemmaRerankerHfRunner,
|
||||
)
|
||||
@@ -0,0 +1,36 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
from tests.models.utils import (
|
||||
RerankModelInfo,
|
||||
)
|
||||
|
||||
from .mteb_score_utils import mteb_test_rerank_models
|
||||
|
||||
RERANK_MODELS = [
|
||||
RerankModelInfo(
|
||||
"cross-encoder/ms-marco-TinyBERT-L-2-v2",
|
||||
architecture="BertForSequenceClassification",
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
mteb_score=0.32898,
|
||||
),
|
||||
RerankModelInfo(
|
||||
"tomaarsen/Qwen3-Reranker-0.6B-seq-cls",
|
||||
architecture="Qwen3ForSequenceClassification",
|
||||
seq_pooling_type="LAST",
|
||||
attn_type="decoder",
|
||||
is_prefix_caching_supported=True,
|
||||
is_chunked_prefill_supported=True,
|
||||
chat_template_name="qwen3_reranker.jinja",
|
||||
mteb_score=0.33459,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", RERANK_MODELS)
|
||||
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
|
||||
mteb_test_rerank_models(vllm_runner, model_info)
|
||||
@@ -0,0 +1,152 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.language.pooling.embed_utils import correctness_test_embed_models
|
||||
from tests.models.utils import (
|
||||
EmbedModelInfo,
|
||||
RerankModelInfo,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from .mteb_embed_utils import mteb_test_embed_models
|
||||
from .mteb_score_utils import mteb_test_rerank_models
|
||||
|
||||
MODELS = [
|
||||
########## BertModel
|
||||
EmbedModelInfo(
|
||||
"thenlper/gte-large",
|
||||
mteb_score=0.76807651,
|
||||
architecture="BertModel",
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo("thenlper/gte-base", architecture="BertModel", enable_test=False),
|
||||
EmbedModelInfo("thenlper/gte-small", architecture="BertModel", enable_test=False),
|
||||
EmbedModelInfo(
|
||||
"thenlper/gte-large-zh", architecture="BertModel", enable_test=False
|
||||
),
|
||||
EmbedModelInfo("thenlper/gte-base-zh", architecture="BertModel", enable_test=False),
|
||||
EmbedModelInfo(
|
||||
"thenlper/gte-small-zh", architecture="BertModel", enable_test=False
|
||||
),
|
||||
########### NewModel
|
||||
# These three architectures are almost the same, but not exactly the same.
|
||||
# For example,
|
||||
# - whether to use token_type_embeddings
|
||||
# - whether to use context expansion
|
||||
# So only test one (the most widely used) model
|
||||
EmbedModelInfo(
|
||||
"Alibaba-NLP/gte-multilingual-base",
|
||||
architecture="GteNewModel",
|
||||
mteb_score=0.775074696,
|
||||
hf_overrides={"architectures": ["GteNewModel"]},
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"Alibaba-NLP/gte-base-en-v1.5",
|
||||
architecture="GteNewModel",
|
||||
hf_overrides={"architectures": ["GteNewModel"]},
|
||||
enable_test=False,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"Alibaba-NLP/gte-large-en-v1.5",
|
||||
architecture="GteNewModel",
|
||||
hf_overrides={"architectures": ["GteNewModel"]},
|
||||
enable_test=False,
|
||||
),
|
||||
########### Qwen2ForCausalLM
|
||||
EmbedModelInfo(
|
||||
"Alibaba-NLP/gte-Qwen2-1.5B-instruct",
|
||||
mteb_score=0.758473459018872,
|
||||
architecture="Qwen2ForCausalLM",
|
||||
seq_pooling_type="LAST",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
# Skip: numerical regression with transformers v5.
|
||||
enable_test=False,
|
||||
),
|
||||
########## ModernBertModel
|
||||
EmbedModelInfo(
|
||||
"Alibaba-NLP/gte-modernbert-base",
|
||||
mteb_score=0.748193353,
|
||||
architecture="ModernBertModel",
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
########## Qwen3ForCausalLM
|
||||
EmbedModelInfo(
|
||||
"Qwen/Qwen3-Embedding-0.6B",
|
||||
mteb_score=0.771163695,
|
||||
architecture="Qwen3ForCausalLM",
|
||||
seq_pooling_type="LAST",
|
||||
attn_type="decoder",
|
||||
is_prefix_caching_supported=True,
|
||||
is_chunked_prefill_supported=True,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"Qwen/Qwen3-Embedding-4B",
|
||||
architecture="Qwen3ForCausalLM",
|
||||
enable_test=False,
|
||||
),
|
||||
]
|
||||
|
||||
RERANK_MODELS = [
|
||||
RerankModelInfo(
|
||||
# classifier_pooling: mean
|
||||
"Alibaba-NLP/gte-reranker-modernbert-base",
|
||||
mteb_score=0.33386,
|
||||
architecture="ModernBertForSequenceClassification",
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
RerankModelInfo(
|
||||
"Alibaba-NLP/gte-multilingual-reranker-base",
|
||||
mteb_score=0.33062,
|
||||
architecture="GteNewForSequenceClassification",
|
||||
hf_overrides={"architectures": ["GteNewForSequenceClassification"]},
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
|
||||
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_correctness(
|
||||
hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
|
||||
) -> None:
|
||||
correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", RERANK_MODELS)
|
||||
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
|
||||
vllm_extra_kwargs = {}
|
||||
if current_platform.is_rocm():
|
||||
vllm_extra_kwargs["attention_backend"] = "TRITON_ATTN"
|
||||
mteb_test_rerank_models(
|
||||
vllm_runner, model_info, vllm_extra_kwargs=vllm_extra_kwargs
|
||||
)
|
||||
@@ -0,0 +1,60 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
from tests.models.language.pooling.embed_utils import correctness_test_embed_models
|
||||
from tests.models.utils import EmbedModelInfo
|
||||
|
||||
from .mteb_embed_utils import mteb_test_embed_models
|
||||
|
||||
MODELS = [
|
||||
########## BertModel
|
||||
EmbedModelInfo(
|
||||
"intfloat/e5-small",
|
||||
architecture="BertModel",
|
||||
mteb_score=0.742285423,
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo("intfloat/e5-base", architecture="BertModel", enable_test=False),
|
||||
EmbedModelInfo("intfloat/e5-large", architecture="BertModel", enable_test=False),
|
||||
EmbedModelInfo(
|
||||
"intfloat/multilingual-e5-small", architecture="BertModel", enable_test=False
|
||||
),
|
||||
########## XLMRobertaModel
|
||||
EmbedModelInfo(
|
||||
"intfloat/multilingual-e5-base",
|
||||
architecture="XLMRobertaModel",
|
||||
mteb_score=0.779325955,
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"intfloat/multilingual-e5-large",
|
||||
architecture="XLMRobertaModel",
|
||||
enable_test=False,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"intfloat/multilingual-e5-large-instruct",
|
||||
architecture="XLMRobertaModel",
|
||||
enable_test=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
|
||||
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_correctness(
|
||||
hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
|
||||
) -> None:
|
||||
correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)
|
||||
@@ -0,0 +1,154 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from functools import partial
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.language.pooling.embed_utils import (
|
||||
check_embeddings_close,
|
||||
correctness_test_embed_models,
|
||||
matryoshka_fy,
|
||||
)
|
||||
from tests.models.utils import (
|
||||
EmbedModelInfo,
|
||||
RerankModelInfo,
|
||||
)
|
||||
from vllm import PoolingParams
|
||||
|
||||
from .mteb_embed_utils import mteb_test_embed_models
|
||||
from .mteb_score_utils import mteb_test_rerank_models
|
||||
|
||||
EMBEDDING_MODELS = [
|
||||
EmbedModelInfo(
|
||||
"jinaai/jina-embeddings-v3",
|
||||
mteb_score=0.824413164,
|
||||
architecture="XLMRobertaModel",
|
||||
is_matryoshka=True,
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"jinaai/jina-embeddings-v5-text-small",
|
||||
mteb_score=0.794535707854956,
|
||||
architecture="JinaEmbeddingsV5Model",
|
||||
seq_pooling_type="LAST",
|
||||
attn_type="decoder",
|
||||
is_prefix_caching_supported=True,
|
||||
is_chunked_prefill_supported=True,
|
||||
),
|
||||
]
|
||||
|
||||
RERANK_MODELS = [
|
||||
RerankModelInfo(
|
||||
"jinaai/jina-reranker-v2-base-multilingual",
|
||||
mteb_score=0.33643,
|
||||
architecture="XLMRobertaForSequenceClassification",
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
|
||||
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
|
||||
task = "retrieval" if "v5" in model_info.name else "text-matching"
|
||||
prompt_prefix: str | None = "Document: " if "v5" in model_info.name else None
|
||||
|
||||
def hf_model_callback(model):
|
||||
model.encode = partial(model.encode, task=task)
|
||||
|
||||
mteb_test_embed_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model_info,
|
||||
hf_model_callback=hf_model_callback,
|
||||
prompt_prefix=prompt_prefix,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
|
||||
def test_embed_models_correctness(
|
||||
hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
|
||||
) -> None:
|
||||
task = "retrieval" if "v5" in model_info.name else "text-matching"
|
||||
|
||||
def hf_model_callback(model):
|
||||
model.encode = partial(model.encode, task=task)
|
||||
|
||||
correctness_test_embed_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model_info,
|
||||
example_prompts,
|
||||
hf_model_callback=hf_model_callback,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", RERANK_MODELS)
|
||||
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
|
||||
mteb_test_rerank_models(vllm_runner, model_info)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="jinaai/jina-embeddings-v3 custom XLMRobertaLoRA model on HF hub "
|
||||
"is incompatible with transformers v5 (missing all_tied_weights_keys)"
|
||||
)
|
||||
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("dimensions", [16, 32])
|
||||
def test_matryoshka(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model_info,
|
||||
dtype: str,
|
||||
dimensions: int,
|
||||
example_prompts,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
if not model_info.is_matryoshka:
|
||||
pytest.skip("Model is not matryoshka")
|
||||
|
||||
# ST will strip the input texts, see test_embedding.py
|
||||
example_prompts = [str(s).strip() for s in example_prompts]
|
||||
|
||||
task = "retrieval" if "v5" in model_info.name else "text-matching"
|
||||
|
||||
with hf_runner(
|
||||
model_info.name,
|
||||
dtype=dtype,
|
||||
is_sentence_transformer=True,
|
||||
) as hf_model:
|
||||
hf_outputs = hf_model.encode(example_prompts, task=task)
|
||||
hf_outputs = matryoshka_fy(hf_outputs, dimensions)
|
||||
|
||||
with vllm_runner(
|
||||
model_info.name, runner="pooling", dtype=dtype, max_model_len=None
|
||||
) as vllm_model:
|
||||
assert vllm_model.llm.llm_engine.model_config.is_matryoshka
|
||||
|
||||
matryoshka_dimensions = (
|
||||
vllm_model.llm.llm_engine.model_config.matryoshka_dimensions
|
||||
)
|
||||
assert matryoshka_dimensions is not None
|
||||
|
||||
if dimensions not in matryoshka_dimensions:
|
||||
with pytest.raises(ValueError):
|
||||
vllm_model.embed(
|
||||
example_prompts, pooling_params=PoolingParams(dimensions=dimensions)
|
||||
)
|
||||
else:
|
||||
vllm_outputs = vllm_model.embed(
|
||||
example_prompts, pooling_params=PoolingParams(dimensions=dimensions)
|
||||
)
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=hf_outputs,
|
||||
embeddings_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
tol=1e-2,
|
||||
)
|
||||
@@ -0,0 +1,109 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import Any
|
||||
|
||||
import mteb
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from tests.conftest import HfRunner
|
||||
from tests.models.utils import RerankModelInfo
|
||||
|
||||
from .mteb_score_utils import MtebCrossEncoderMixin, mteb_test_rerank_models
|
||||
|
||||
mxbai_rerank_hf_overrides = {
|
||||
"architectures": ["Qwen2ForSequenceClassification"],
|
||||
"classifier_from_token": ["0", "1"],
|
||||
"method": "from_2_way_softmax",
|
||||
}
|
||||
|
||||
RERANK_MODELS = [
|
||||
RerankModelInfo(
|
||||
"mixedbread-ai/mxbai-rerank-base-v2",
|
||||
architecture="Qwen2ForSequenceClassification",
|
||||
hf_overrides=mxbai_rerank_hf_overrides,
|
||||
seq_pooling_type="LAST",
|
||||
attn_type="decoder",
|
||||
is_prefix_caching_supported=True,
|
||||
is_chunked_prefill_supported=True,
|
||||
chat_template_name="mxbai_rerank_v2.jinja",
|
||||
mteb_score=0.33651,
|
||||
enable_test=True,
|
||||
),
|
||||
RerankModelInfo(
|
||||
"mixedbread-ai/mxbai-rerank-large-v2",
|
||||
architecture="Qwen2ForSequenceClassification",
|
||||
hf_overrides=mxbai_rerank_hf_overrides,
|
||||
chat_template_name="mxbai_rerank_v2.jinja",
|
||||
enable_test=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class MxbaiRerankerHfRunner(MtebCrossEncoderMixin, HfRunner):
|
||||
def __init__(
|
||||
self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any
|
||||
) -> None:
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
HfRunner.__init__(
|
||||
self,
|
||||
model_name=model_name,
|
||||
auto_cls=AutoModelForCausalLM,
|
||||
dtype=dtype,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
|
||||
self.yes_loc = self.tokenizer.convert_tokens_to_ids("1")
|
||||
self.no_loc = self.tokenizer.convert_tokens_to_ids("0")
|
||||
|
||||
@torch.no_grad
|
||||
def predict(
|
||||
self,
|
||||
inputs1: DataLoader[mteb.types.BatchedInput],
|
||||
inputs2: DataLoader[mteb.types.BatchedInput],
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> np.ndarray:
|
||||
queries = [text for batch in inputs1 for text in batch["text"]]
|
||||
corpus = [text for batch in inputs2 for text in batch["text"]]
|
||||
|
||||
tokenizer = self.tokenizer
|
||||
prompts = []
|
||||
for query, document in zip(queries, corpus):
|
||||
conversation = [
|
||||
{"role": "query", "content": query},
|
||||
{"role": "document", "content": document},
|
||||
]
|
||||
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
conversation=conversation,
|
||||
tools=None,
|
||||
chat_template=self.chat_template,
|
||||
tokenize=False,
|
||||
)
|
||||
prompts.append(prompt)
|
||||
|
||||
def compute_logits(inputs):
|
||||
logits = self.model(**inputs).logits[:, -1, :]
|
||||
yes_logits = logits[:, self.yes_loc]
|
||||
no_logits = logits[:, self.no_loc]
|
||||
logits = yes_logits - no_logits
|
||||
scores = logits.float().sigmoid()
|
||||
return scores
|
||||
|
||||
scores = []
|
||||
for prompt in prompts:
|
||||
inputs = tokenizer([prompt], return_tensors="pt")
|
||||
inputs = self.wrap_device(inputs)
|
||||
score = compute_logits(inputs)
|
||||
scores.append(score[0].item())
|
||||
return torch.Tensor(scores)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", RERANK_MODELS)
|
||||
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
|
||||
mteb_test_rerank_models(vllm_runner, model_info, hf_runner=MxbaiRerankerHfRunner)
|
||||
@@ -0,0 +1,50 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.language.pooling_mteb_test.mteb_embed_utils import (
|
||||
mteb_test_embed_models,
|
||||
)
|
||||
from tests.models.language.pooling_mteb_test.mteb_score_utils import (
|
||||
mteb_test_rerank_models,
|
||||
)
|
||||
from tests.models.utils import (
|
||||
EmbedModelInfo,
|
||||
RerankModelInfo,
|
||||
)
|
||||
|
||||
EMBEDDING_MODELS = [
|
||||
EmbedModelInfo(
|
||||
"nvidia/llama-nemotron-embed-1b-v2",
|
||||
architecture="LlamaBidirectionalModel",
|
||||
mteb_score=0.689164662128673,
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
)
|
||||
]
|
||||
|
||||
RERANK_MODELS = [
|
||||
RerankModelInfo(
|
||||
"nvidia/llama-nemotron-rerank-1b-v2",
|
||||
architecture="LlamaBidirectionalForSequenceClassification",
|
||||
chat_template_name="nemotron-rerank.jinja",
|
||||
mteb_score=0.33994,
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
|
||||
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
|
||||
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", RERANK_MODELS)
|
||||
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
|
||||
mteb_test_rerank_models(vllm_runner, model_info)
|
||||
@@ -0,0 +1,56 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.language.pooling.embed_utils import correctness_test_embed_models
|
||||
from tests.models.utils import EmbedModelInfo
|
||||
|
||||
from .mteb_embed_utils import mteb_test_embed_models
|
||||
|
||||
MODELS = [
|
||||
EmbedModelInfo(
|
||||
"nomic-ai/nomic-embed-text-v1",
|
||||
architecture="NomicBertModel",
|
||||
# Fixme:
|
||||
# Update nomic-embed code to support the latest
|
||||
# HF version and remove revision set.
|
||||
revision="720244025c1a7e15661a174c63cce63c8218e52b",
|
||||
mteb_score=0.737568559,
|
||||
enable_test=True,
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"nomic-ai/nomic-embed-text-v1.5",
|
||||
architecture="NomicBertModel",
|
||||
enable_test=False,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"nomic-ai/CodeRankEmbed", architecture="NomicBertModel", enable_test=False
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"nomic-ai/nomic-embed-text-v2-moe",
|
||||
architecture="NomicBertModel",
|
||||
mteb_score=0.715488912,
|
||||
enable_test=True,
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
|
||||
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_correctness(
|
||||
hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
|
||||
) -> None:
|
||||
correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)
|
||||
@@ -0,0 +1,130 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
from typing import Any
|
||||
|
||||
import mteb
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from tests.conftest import HfRunner
|
||||
from tests.models.utils import RerankModelInfo
|
||||
from tests.utils import multi_gpu_test
|
||||
|
||||
from .mteb_score_utils import MtebCrossEncoderMixin, mteb_test_rerank_models
|
||||
|
||||
qwen3_reranker_hf_overrides = {
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
}
|
||||
|
||||
RERANK_MODELS = [
|
||||
RerankModelInfo(
|
||||
"Qwen/Qwen3-Reranker-0.6B",
|
||||
architecture="Qwen3ForSequenceClassification",
|
||||
hf_overrides=qwen3_reranker_hf_overrides,
|
||||
chat_template_name="qwen3_reranker.jinja",
|
||||
seq_pooling_type="LAST",
|
||||
attn_type="decoder",
|
||||
is_prefix_caching_supported=True,
|
||||
is_chunked_prefill_supported=True,
|
||||
mteb_score=0.33459,
|
||||
enable_test=True,
|
||||
),
|
||||
RerankModelInfo(
|
||||
"Qwen/Qwen3-Reranker-4B",
|
||||
architecture="Qwen3ForSequenceClassification",
|
||||
chat_template_name="qwen3_reranker.jinja",
|
||||
hf_overrides=qwen3_reranker_hf_overrides,
|
||||
enable_test=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class Qwen3RerankerHfRunner(MtebCrossEncoderMixin, HfRunner):
|
||||
def __init__(
|
||||
self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any
|
||||
) -> None:
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
HfRunner.__init__(
|
||||
self,
|
||||
model_name=model_name,
|
||||
auto_cls=AutoModelForCausalLM,
|
||||
dtype=dtype,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
|
||||
self.token_false_id = self.tokenizer.convert_tokens_to_ids("no")
|
||||
self.token_true_id = self.tokenizer.convert_tokens_to_ids("yes")
|
||||
self.max_length = 40960
|
||||
|
||||
@torch.no_grad
|
||||
def predict(
|
||||
self,
|
||||
inputs1: DataLoader[mteb.types.BatchedInput],
|
||||
inputs2: DataLoader[mteb.types.BatchedInput],
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> np.ndarray:
|
||||
queries = [text for batch in inputs1 for text in batch["text"]]
|
||||
corpus = [text for batch in inputs2 for text in batch["text"]]
|
||||
|
||||
tokenizer = self.tokenizer
|
||||
prompts = []
|
||||
for query, document in zip(queries, corpus):
|
||||
conversation = [
|
||||
{"role": "query", "content": query},
|
||||
{"role": "document", "content": document},
|
||||
]
|
||||
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
conversation=conversation,
|
||||
tools=None,
|
||||
chat_template=self.chat_template,
|
||||
tokenize=False,
|
||||
)
|
||||
prompts.append(prompt)
|
||||
|
||||
def compute_logits(inputs):
|
||||
batch_scores = self.model(**inputs).logits[:, -1, :]
|
||||
true_vector = batch_scores[:, self.token_true_id]
|
||||
false_vector = batch_scores[:, self.token_false_id]
|
||||
batch_scores = torch.stack([false_vector, true_vector], dim=1)
|
||||
batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)
|
||||
scores = batch_scores[:, 1].exp()
|
||||
return scores
|
||||
|
||||
scores = []
|
||||
for prompt in prompts:
|
||||
inputs = tokenizer([prompt], return_tensors="pt")
|
||||
inputs = self.wrap_device(inputs)
|
||||
score = compute_logits(inputs)
|
||||
scores.append(score[0].item())
|
||||
return torch.Tensor(scores)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", RERANK_MODELS)
|
||||
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
|
||||
mteb_test_rerank_models(vllm_runner, model_info, hf_runner=Qwen3RerankerHfRunner)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", RERANK_MODELS)
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
def test_rerank_models_mteb_tp(vllm_runner, model_info: RerankModelInfo) -> None:
|
||||
assert model_info.architecture == "Qwen3ForSequenceClassification"
|
||||
|
||||
vllm_extra_kwargs: dict[str, Any] = {
|
||||
"tensor_parallel_size": 2,
|
||||
}
|
||||
|
||||
mteb_test_rerank_models(
|
||||
vllm_runner,
|
||||
model_info,
|
||||
vllm_extra_kwargs=vllm_extra_kwargs,
|
||||
hf_runner=Qwen3RerankerHfRunner,
|
||||
)
|
||||
@@ -0,0 +1,97 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.language.pooling.embed_utils import correctness_test_embed_models
|
||||
from tests.models.utils import EmbedModelInfo
|
||||
|
||||
from .mteb_embed_utils import mteb_test_embed_models
|
||||
|
||||
MODELS = [
|
||||
EmbedModelInfo(
|
||||
"Snowflake/snowflake-arctic-embed-xs",
|
||||
is_matryoshka=False,
|
||||
architecture="BertModel",
|
||||
mteb_score=0.714927797,
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"Snowflake/snowflake-arctic-embed-s",
|
||||
is_matryoshka=False,
|
||||
architecture="BertModel",
|
||||
enable_test=False,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"Snowflake/snowflake-arctic-embed-m",
|
||||
is_matryoshka=False,
|
||||
architecture="BertModel",
|
||||
enable_test=False,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"Snowflake/snowflake-arctic-embed-m-long",
|
||||
is_matryoshka=False,
|
||||
architecture="NomicBertModel",
|
||||
mteb_score=0.681146831,
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"Snowflake/snowflake-arctic-embed-l",
|
||||
is_matryoshka=False,
|
||||
architecture="BertModel",
|
||||
enable_test=False,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
is_matryoshka=True,
|
||||
architecture="BertModel",
|
||||
mteb_score=0.649088363,
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"Snowflake/snowflake-arctic-embed-l-v2.0",
|
||||
is_matryoshka=True,
|
||||
architecture="XLMRobertaModel",
|
||||
mteb_score=0.712258299,
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"Snowflake/snowflake-arctic-embed-m-v2.0",
|
||||
is_matryoshka=True,
|
||||
architecture="GteModel",
|
||||
mteb_score=0.706622444,
|
||||
seq_pooling_type="CLS",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
|
||||
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_correctness(
|
||||
hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
|
||||
) -> None:
|
||||
correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)
|
||||
@@ -0,0 +1,38 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
from tests.models.utils import (
|
||||
EmbedModelInfo,
|
||||
)
|
||||
|
||||
from .mteb_embed_utils import mteb_test_embed_models
|
||||
|
||||
# ST models with projector (Dense) layers
|
||||
ST_PROJECTOR_MODELS = [
|
||||
EmbedModelInfo(
|
||||
"TencentBAC/Conan-embedding-v1",
|
||||
architecture="BertModel",
|
||||
mteb_score=0.688611955,
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
EmbedModelInfo(
|
||||
"google/embeddinggemma-300m",
|
||||
architecture="Gemma3TextModel",
|
||||
mteb_score=0.7473819294684156,
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
enable_test=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS)
|
||||
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
|
||||
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
|
||||
@@ -0,0 +1,56 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.language.pooling.embed_utils import correctness_test_embed_models
|
||||
from tests.models.utils import EmbedModelInfo
|
||||
|
||||
from .mteb_embed_utils import mteb_test_embed_models
|
||||
|
||||
MODELS = [
|
||||
EmbedModelInfo(
|
||||
"voyageai/voyage-4-nano",
|
||||
architecture="VoyageQwen3BidirectionalEmbedModel",
|
||||
enable_test=True,
|
||||
seq_pooling_type="MEAN",
|
||||
attn_type="encoder_only",
|
||||
is_prefix_caching_supported=False,
|
||||
is_chunked_prefill_supported=False,
|
||||
hf_overrides={
|
||||
"architectures": ["VoyageQwen3BidirectionalEmbedModel"],
|
||||
"num_labels": 2048,
|
||||
},
|
||||
mteb_score=0.7054,
|
||||
# === MTEB Results ===
|
||||
# STS12: 0.6613
|
||||
# STS13: 0.6906
|
||||
# STS14: 0.6556
|
||||
# STS15: 0.7843
|
||||
# STS16: 0.7340
|
||||
# STSBenchmark: 0.7063
|
||||
# Average score: 0.7054
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
|
||||
# Encoder-only attention models need enforce_eager=True to avoid
|
||||
# CUDA graph capture issues with piecewise compilation
|
||||
mteb_test_embed_models(
|
||||
hf_runner, vllm_runner, model_info, vllm_extra_kwargs={"enforce_eager": True}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", MODELS)
|
||||
def test_embed_models_correctness(
|
||||
hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
|
||||
) -> None:
|
||||
correctness_test_embed_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model_info,
|
||||
example_prompts,
|
||||
vllm_extra_kwargs={"enforce_eager": True},
|
||||
)
|
||||
@@ -0,0 +1,69 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Pytest configuration for vLLM multimodal tests."""
|
||||
|
||||
import os
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
"""Early ROCm configuration that must happen before test collection."""
|
||||
if not current_platform.is_rocm():
|
||||
return
|
||||
|
||||
# Disable skinny GEMM on ROCm to avoid non-deterministic results
|
||||
# from atomic reductions in wvSplitKrc kernel.
|
||||
# See: https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975
|
||||
os.environ["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
|
||||
warnings.warn(
|
||||
"ROCm: Set VLLM_ROCM_USE_SKINNY_GEMM=0 to avoid non-deterministic "
|
||||
"results from skinny GEMM atomic reductions",
|
||||
UserWarning,
|
||||
stacklevel=1,
|
||||
)
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(config, items):
|
||||
"""Configure ROCm-specific settings based on collected tests."""
|
||||
if not current_platform.is_rocm():
|
||||
return
|
||||
|
||||
skip_patterns = ["test_granite_speech.py"]
|
||||
if any(pattern in str(arg) for arg in config.args for pattern in skip_patterns):
|
||||
return
|
||||
|
||||
# Disable Flash/MemEfficient SDP on ROCm to avoid HF Transformers
|
||||
# accuracy issues: https://github.com/vllm-project/vllm/issues/30167
|
||||
# TODO: Remove once ROCm SDP accuracy issues are resolved on HuggingFace
|
||||
torch.backends.cuda.enable_flash_sdp(False)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(False)
|
||||
torch.backends.cuda.enable_math_sdp(True)
|
||||
warnings.warn(
|
||||
"ROCm: Disabled flash_sdp and mem_efficient_sdp, enabled math_sdp "
|
||||
"to avoid HuggingFace Transformers accuracy issues",
|
||||
UserWarning,
|
||||
stacklevel=1,
|
||||
)
|
||||
|
||||
|
||||
def patch_hf_vision_attn_for_rocm(model):
|
||||
"""Force SDPA for HF vision encoders on ROCm.
|
||||
|
||||
HF's flash_attention_2 has accuracy issues on ROCm that bypass
|
||||
torch.backends.cuda settings. This forces SDPA which then uses
|
||||
math_sdp via the pytest_collection_modifyitems settings.
|
||||
"""
|
||||
if not current_platform.is_rocm():
|
||||
return
|
||||
|
||||
inner = getattr(model, "model", model)
|
||||
|
||||
if hasattr(inner, "vision_embedding"):
|
||||
vit = inner.vision_embedding[0]
|
||||
for layer in vit.encoder.layers:
|
||||
if hasattr(layer, "self_attn"):
|
||||
layer.self_attn.vision_config._attn_implementation = "sdpa"
|
||||
@@ -0,0 +1,171 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Copyright 2025 The vLLM team.
|
||||
# Copyright 2025 NVIDIA CORPORATION and the HuggingFace Inc. team. All rights
|
||||
# reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.registry import HF_EXAMPLE_MODELS
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
MODEL_NAME = "nvidia/audio-flamingo-3-hf"
|
||||
SINGLE_CONVERSATION = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What is surprising about the relationship between "
|
||||
"the barking and the music?",
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "https://huggingface.co/datasets/nvidia/AudioSkills/"
|
||||
"resolve/main/assets/"
|
||||
"dogs_barking_in_sync_with_the_music.wav",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
BATCHED_CONVERSATIONS = [
|
||||
SINGLE_CONVERSATION,
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Why is the philosopher's name mentioned in the "
|
||||
"lyrics? (A) To express a sense of nostalgia "
|
||||
"(B) To indicate that language cannot express clearly, "
|
||||
"satirizing the inversion of black and white in the world "
|
||||
"(C) To add depth and complexity to the lyrics "
|
||||
"(D) To showcase the wisdom and influence of the "
|
||||
"philosopher",
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "https://huggingface.co/datasets/nvidia/"
|
||||
"AudioSkills/resolve/main/assets/"
|
||||
"Ch6Ae9DT6Ko_00-04-03_00-04-31.wav",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
]
|
||||
|
||||
|
||||
def get_fixture_path(filename):
|
||||
return os.path.join(
|
||||
os.path.dirname(__file__), "../../fixtures/audioflamingo3", filename
|
||||
)
|
||||
|
||||
|
||||
def assert_output_matches(output, expected_text, expected_token_ids):
|
||||
generated = output.outputs[0]
|
||||
assert generated.text.strip() == expected_text
|
||||
actual_token_ids = list(generated.token_ids)
|
||||
assert (
|
||||
actual_token_ids == expected_token_ids
|
||||
or actual_token_ids == expected_token_ids[:-1]
|
||||
or actual_token_ids[:-1] == expected_token_ids
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm():
|
||||
model_info = HF_EXAMPLE_MODELS.get_hf_info("AudioFlamingo3ForConditionalGeneration")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
try:
|
||||
return LLM(
|
||||
model=MODEL_NAME,
|
||||
dtype="bfloat16",
|
||||
enforce_eager=True,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
)
|
||||
except Exception as e:
|
||||
pytest.skip(f"Failed to load model {MODEL_NAME}: {e}")
|
||||
|
||||
|
||||
def test_single_generation(llm):
|
||||
fixture_path = get_fixture_path("expected_results_single.json")
|
||||
if not os.path.exists(fixture_path):
|
||||
pytest.skip(f"Fixture not found: {fixture_path}")
|
||||
|
||||
with open(fixture_path) as f:
|
||||
expected = json.load(f)
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
|
||||
|
||||
outputs = llm.chat(
|
||||
messages=SINGLE_CONVERSATION,
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
assert_output_matches(
|
||||
outputs[0],
|
||||
expected["transcriptions"][0],
|
||||
expected["token_ids"][0],
|
||||
)
|
||||
|
||||
|
||||
def test_batched_generation(llm):
|
||||
fixture_path = get_fixture_path("expected_results_batched.json")
|
||||
if not os.path.exists(fixture_path):
|
||||
pytest.skip(f"Fixture not found: {fixture_path}")
|
||||
|
||||
with open(fixture_path) as f:
|
||||
expected = json.load(f)
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
|
||||
|
||||
outputs = llm.chat(
|
||||
messages=BATCHED_CONVERSATIONS,
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
|
||||
for i, output in enumerate(outputs):
|
||||
assert_output_matches(
|
||||
output,
|
||||
expected["transcriptions"][i],
|
||||
expected["token_ids"][i],
|
||||
)
|
||||
|
||||
|
||||
def test_single_and_batched_generation_match(llm):
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
|
||||
|
||||
single_output = llm.chat(
|
||||
messages=SINGLE_CONVERSATION,
|
||||
sampling_params=sampling_params,
|
||||
)[0]
|
||||
batched_output = llm.chat(
|
||||
messages=BATCHED_CONVERSATIONS,
|
||||
sampling_params=sampling_params,
|
||||
)[0]
|
||||
|
||||
assert single_output.outputs[0].text == batched_output.outputs[0].text
|
||||
assert list(single_output.outputs[0].token_ids) == list(
|
||||
batched_output.outputs[0].token_ids
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,181 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import pytest
|
||||
from transformers import AutoModelForSpeechSeq2Seq
|
||||
|
||||
from vllm.logprobs import SampleLogprobs
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....conftest import AudioTestAssets, HfRunner, PromptAudioInput, VllmRunner
|
||||
from ...registry import HF_EXAMPLE_MODELS
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
HF_AUDIO_PROMPT = "<|start_of_role|>system<|end_of_role|>Knowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant<|end_of_text|>\n<|start_of_role|>user<|end_of_role|><|audio|>can you transcribe the speech into a written format?<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>" # noqa: E501
|
||||
|
||||
|
||||
def vllm_to_hf_output(
|
||||
vllm_output: tuple[list[int], str, SampleLogprobs | None],
|
||||
) -> tuple[list[int], str, SampleLogprobs | None]:
|
||||
"""Sanitize hf output to be comparable with vllm output."""
|
||||
output_ids, output_str, out_logprobs = vllm_output
|
||||
|
||||
hf_output_str = output_str + "<|end_of_text|>"
|
||||
|
||||
return output_ids, hf_output_str, out_logprobs
|
||||
|
||||
|
||||
MODEL_NAME = "ibm-granite/granite-speech-3.3-2b"
|
||||
MODEL_NAME_4_0 = "ibm-granite/granite-4.0-1b-speech"
|
||||
# "plus" variant of granite speech (uses GraniteSpeechPlusForConditionalGeneration).
|
||||
MODEL_NAME_4_1_PLUS = "ibm-granite/granite-speech-4.1-2b-plus"
|
||||
# Audio lora co-exists directly in the 3.3 model directory,
|
||||
# the 4.0 and 4.1-plus models have adapters merged into the weights.
|
||||
models: dict[str, str | None] = {
|
||||
MODEL_NAME: MODEL_NAME,
|
||||
MODEL_NAME_4_0: None,
|
||||
MODEL_NAME_4_1_PLUS: None,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def granite_speech_attention_config():
|
||||
"""Return attention config for Granite Speech tests on ROCm."""
|
||||
if current_platform.is_rocm():
|
||||
from vllm.platforms.rocm import on_mi3xx
|
||||
|
||||
if on_mi3xx():
|
||||
return {"backend": "ROCM_AITER_FA"}
|
||||
return {"backend": "TRITON_ATTN"}
|
||||
return None
|
||||
|
||||
|
||||
def run_test(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
inputs: Sequence[tuple[list[str], PromptAudioInput]],
|
||||
model: str,
|
||||
*,
|
||||
max_model_len: int,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
tensor_parallel_size: int,
|
||||
distributed_executor_backend: str | None = None,
|
||||
attention_config: dict | None = None,
|
||||
audio_lora_path: str | None = None,
|
||||
):
|
||||
"""Inference result should be the same between hf and vllm.
|
||||
|
||||
All the audio fixtures for the test are from AUDIO_ASSETS.
|
||||
For huggingface runner, we provide the audio as input.
|
||||
For vllm runner, we provide MultiModalDataDict objects
|
||||
and corresponding MultiModalConfig as input.
|
||||
Note, the text input is also adjusted to abide by vllm contract.
|
||||
The text output is sanitized to be able to compare with hf.
|
||||
"""
|
||||
# NOTE: take care of the order. run vLLM first, and then run HF.
|
||||
# vLLM needs a fresh new process without cuda initialization.
|
||||
# if we run HF first, the cuda initialization will be done and it
|
||||
# will hurt multiprocessing backend with fork method (the default method).
|
||||
# max_model_len should be greater than image_feature_size
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="generate",
|
||||
max_model_len=max_model_len,
|
||||
max_num_seqs=1,
|
||||
dtype=dtype,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
enable_lora=audio_lora_path is not None,
|
||||
max_lora_rank=64,
|
||||
enforce_eager=True,
|
||||
attention_config=attention_config,
|
||||
) as vllm_model:
|
||||
lora_request = (
|
||||
LoRARequest("audio", 1, audio_lora_path) if audio_lora_path else None
|
||||
)
|
||||
vllm_outputs_per_case = [
|
||||
vllm_model.generate_greedy_logprobs(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
audios=audios,
|
||||
lora_request=lora_request,
|
||||
)
|
||||
for prompts, audios in inputs
|
||||
]
|
||||
|
||||
with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSpeechSeq2Seq) as hf_model:
|
||||
hf_processor = hf_model.processor
|
||||
eos_token_id = hf_processor.tokenizer.eos_token_id
|
||||
|
||||
hf_outputs_per_case = [
|
||||
hf_model.generate_greedy_logprobs_limit(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
audios=[audios],
|
||||
eos_token_id=eos_token_id,
|
||||
)
|
||||
for prompts, audios in inputs
|
||||
]
|
||||
|
||||
for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=[vllm_to_hf_output(output) for output in vllm_outputs],
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model,audio_lora_path", models.items())
|
||||
@pytest.mark.parametrize(
|
||||
"dtype", ["float16"] if current_platform.is_rocm() else ["bfloat16"]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"max_model_len", [512] if current_platform.is_rocm() else [2048]
|
||||
)
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize("num_logprobs", [10])
|
||||
def test_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model: str,
|
||||
audio_lora_path: str | None,
|
||||
audio_assets: AudioTestAssets,
|
||||
granite_speech_attention_config,
|
||||
dtype: str,
|
||||
max_model_len: int,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
audio, sr = audio_assets[0].audio_and_sample_rate
|
||||
# This model expects 16k sample rate, which our test audio
|
||||
# already is; if this changes, it may break this test,
|
||||
# so we check it directly
|
||||
assert sr == 16000
|
||||
run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
[
|
||||
([HF_AUDIO_PROMPT], [audio]),
|
||||
],
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_model_len=max_model_len,
|
||||
max_tokens=max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
tensor_parallel_size=1,
|
||||
attention_config=granite_speech_attention_config,
|
||||
audio_lora_path=audio_lora_path,
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
|
||||
models = ["llava-hf/llava-onevision-qwen2-0.5b-ov-hf"]
|
||||
|
||||
|
||||
def base_prompt(modalities_str: str) -> str:
|
||||
return f"<|im_start|>user {modalities_str}\nDescribe what you see from these items.<|im_end|><|im_start|>assistant\n" # noqa: E501
|
||||
|
||||
|
||||
INTERLEAVED_PROMPT = base_prompt("<image><video><image>\n")
|
||||
NONINTERLEAVED_PROMPT = base_prompt("<image><image><video>\n")
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize("dtype", ["float16"])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
def test_models(vllm_runner, model, dtype: str, max_tokens: int) -> None:
|
||||
"""
|
||||
This is a simple test to check if interleaved and non-interleaved prompts
|
||||
give the same result.
|
||||
"""
|
||||
|
||||
image_cherry = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
|
||||
image_stop = convert_image_mode(ImageAsset("stop_sign").pil_image, "RGB")
|
||||
images = [image_cherry, image_stop]
|
||||
video = VideoAsset(name="baby_reading", num_frames=16).np_ndarrays
|
||||
|
||||
inputs = [
|
||||
(
|
||||
[INTERLEAVED_PROMPT],
|
||||
[images],
|
||||
[video],
|
||||
),
|
||||
(
|
||||
[NONINTERLEAVED_PROMPT],
|
||||
[images],
|
||||
[video],
|
||||
),
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="generate",
|
||||
dtype=dtype,
|
||||
limit_mm_per_prompt={"image": 2},
|
||||
max_model_len=32768,
|
||||
max_num_seqs=2,
|
||||
tensor_parallel_size=1,
|
||||
enforce_eager=True,
|
||||
) as vllm_model:
|
||||
vllm_outputs_per_case = [
|
||||
vllm_model.generate_greedy(
|
||||
prompts, max_tokens, images=images, videos=videos
|
||||
)
|
||||
for prompts, images, videos in inputs
|
||||
]
|
||||
|
||||
all_results = [output[0][1] for output in vllm_outputs_per_case]
|
||||
outputs = [
|
||||
(total_str, total_str.find("assistant\n") + len("assistant\n"))
|
||||
for total_str in all_results
|
||||
]
|
||||
prompt_lengths = [prompt_len for _, prompt_len in outputs]
|
||||
generated_strs = [total_str[prompt_len:] for total_str, prompt_len in outputs]
|
||||
interleaved_prompt_len, noninterleaved_prompt_len = prompt_lengths
|
||||
interleaved_output_str, noninterleaved_output_str = generated_strs
|
||||
|
||||
# The two prompts are identical except for the order of modality tokens.
|
||||
assert interleaved_prompt_len == noninterleaved_prompt_len
|
||||
|
||||
# The two generated strings should be different because of the
|
||||
# interleaved modality tokens.
|
||||
assert interleaved_output_str != noninterleaved_output_str
|
||||
@@ -0,0 +1,76 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import NamedTuple
|
||||
|
||||
import pytest
|
||||
from PIL.Image import Image
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from vllm import LLM, EngineArgs, SamplingParams
|
||||
from vllm.multimodal.utils import encode_image_url
|
||||
|
||||
MODEL_NAME = "Kwai-Keye/Keye-VL-8B-Preview"
|
||||
|
||||
QUESTION = "What is the content of each image?"
|
||||
|
||||
|
||||
class ModelRequestData(NamedTuple):
|
||||
engine_args: EngineArgs
|
||||
prompt: str
|
||||
image_data: list[Image]
|
||||
stop_token_ids: list[int] | None = None
|
||||
chat_template: str | None = None
|
||||
sampling_params: SamplingParams | None = None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("question", [QUESTION])
|
||||
def test_keye_vl(image_assets, question: str):
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
image_urls = [encode_image_url(image) for image in images]
|
||||
|
||||
placeholders = [{"type": "image", "image": url} for url in image_urls]
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
*placeholders,
|
||||
{"type": "text", "text": question},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
|
||||
|
||||
prompt = processor.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
llm = LLM(
|
||||
model=MODEL_NAME,
|
||||
trust_remote_code=True,
|
||||
max_model_len=8192,
|
||||
max_num_seqs=5,
|
||||
limit_mm_per_prompt={"image": len(image_urls)},
|
||||
seed=42,
|
||||
)
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.0, max_tokens=256, stop_token_ids=None
|
||||
)
|
||||
|
||||
outputs = llm.generate(
|
||||
{
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": images},
|
||||
},
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
|
||||
print("-" * 50)
|
||||
for o in outputs:
|
||||
generated_text = o.outputs[0].text
|
||||
print(generated_text)
|
||||
assert len(generated_text) > 10, (
|
||||
f"Generated text is too short: {generated_text}"
|
||||
)
|
||||
print("-" * 50)
|
||||
@@ -0,0 +1,723 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Create a reduced-layer version of the Maverick model for testing purposes.
|
||||
|
||||
This script creates a new model with fewer layers by:
|
||||
1. Loading the original Maverick model configuration
|
||||
2. Creating a reduced configuration
|
||||
3. Generating compatible safetensors files with appropriate weights
|
||||
4. Creating the necessary index files for vLLM compatibility
|
||||
"""
|
||||
|
||||
import json
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
from transformers import AutoConfig, AutoProcessor, AutoTokenizer, GenerationConfig
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.v1.executor.abstract import Executor
|
||||
from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, FullAttentionSpec
|
||||
|
||||
from ....utils import multi_gpu_test
|
||||
|
||||
# Sample prompts for testing
|
||||
PROMPTS: list[str] = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
|
||||
def run_maverick_serving(model: str):
|
||||
"""Test Llama-4-Maverick model with vLLM LLM class using CLI equivalent
|
||||
options with reduced layers.
|
||||
"""
|
||||
|
||||
try:
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model=model,
|
||||
max_model_len=2048,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=8,
|
||||
enable_expert_parallel=True,
|
||||
trust_remote_code=True,
|
||||
gpu_memory_utilization=0.4,
|
||||
kv_cache_dtype="fp8",
|
||||
)
|
||||
|
||||
outputs = llm.generate(PROMPTS, sampling_params)
|
||||
|
||||
# Print the outputs
|
||||
print("\nGenerated Outputs:\n" + "-" * 60)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}")
|
||||
print(f"Output: {generated_text!r}")
|
||||
print("-" * 60)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error initializing or running model: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def get_rope_layers_config(model_path: str) -> list[int]:
|
||||
"""
|
||||
Get the interleaved RoPE configuration from HuggingFace config
|
||||
|
||||
Args:
|
||||
model_path: Path to the local directory containing the reduced
|
||||
Maverick model checkpoint
|
||||
|
||||
Returns:
|
||||
List of 0 or 1 indicating whether each layer uses RoPE and local attn
|
||||
0 indicates that RoPE is not used while 1 indicates that RoPE is used.
|
||||
"""
|
||||
config_path = Path(model_path) / "config.json"
|
||||
model_config = json.loads(config_path.read_text())
|
||||
text_config = model_config["text_config"]
|
||||
no_rope_layers = text_config["no_rope_layers"]
|
||||
print(f"Found no_rope_layers: {no_rope_layers}")
|
||||
return no_rope_layers
|
||||
|
||||
|
||||
def create_reduced_maverick_model(
|
||||
original_model_name: str = "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
|
||||
output_dir: str = "/tmp/reduced_maverick",
|
||||
text_layers: int = 4,
|
||||
num_experts: int = 4,
|
||||
vision_layers: int = 2,
|
||||
force_recreate: bool = False,
|
||||
) -> str:
|
||||
"""
|
||||
Create a reduced-layer version of the Maverick model.
|
||||
|
||||
Args:
|
||||
original_model_name: Name of the original Maverick model
|
||||
output_dir: Directory to save the reduced model
|
||||
text_layers: Number of text transformer layers
|
||||
num_experts: Number of experts per layer
|
||||
vision_layers: Number of vision transformer layers
|
||||
force_recreate: Whether to recreate if output_dir already exists
|
||||
|
||||
Returns:
|
||||
Path to the created reduced model directory
|
||||
"""
|
||||
|
||||
print(
|
||||
f"Creating reduced Maverick model with {text_layers} text layers and "
|
||||
f"{vision_layers} vision layers..."
|
||||
)
|
||||
|
||||
# Create output directory
|
||||
output_path = Path(output_dir)
|
||||
if output_path.exists():
|
||||
if force_recreate:
|
||||
shutil.rmtree(output_path)
|
||||
else:
|
||||
print(
|
||||
f"Output directory {output_dir} already exists. "
|
||||
"Use --force-recreate to overwrite."
|
||||
)
|
||||
return str(output_path)
|
||||
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
print("Loading original model configuration...")
|
||||
original_config = AutoConfig.from_pretrained(
|
||||
original_model_name, trust_remote_code=True
|
||||
)
|
||||
print("Creating reduced configuration...")
|
||||
reduced_config = create_reduced_config(
|
||||
original_config, text_layers, num_experts, vision_layers
|
||||
)
|
||||
|
||||
config_path = output_path / "config.json"
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(reduced_config, f, indent=2)
|
||||
print(f"Saved reduced config to {config_path}")
|
||||
|
||||
print("Copying tokenizer files...")
|
||||
copy_tokenizer_files(original_model_name, output_path)
|
||||
|
||||
print("Creating reduced safetensors files...")
|
||||
create_reduced_safetensors(original_config, reduced_config, output_path)
|
||||
|
||||
print("Creating preprocessor config...")
|
||||
create_preprocessor_config(original_config, output_path)
|
||||
|
||||
try:
|
||||
gen_config = GenerationConfig.from_pretrained(original_model_name)
|
||||
gen_config.save_pretrained(output_path)
|
||||
print("Copied generation config")
|
||||
except Exception as e:
|
||||
print(f"Could not copy generation config: {e}")
|
||||
|
||||
print(f"Successfully created reduced Maverick model at {output_path}")
|
||||
return str(output_path)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error creating reduced model: {e}")
|
||||
# Clean up on failure
|
||||
if output_path.exists():
|
||||
shutil.rmtree(output_path)
|
||||
raise
|
||||
|
||||
|
||||
def create_reduced_config(
|
||||
original_config: Any, text_layers: int, num_experts: int, vision_layers: int
|
||||
) -> dict[str, Any]:
|
||||
"""Create a reduced configuration based on the original."""
|
||||
|
||||
# Convert config to dictionary
|
||||
config_dict = original_config.to_dict()
|
||||
|
||||
# Reduce text layers
|
||||
if "text_config" in config_dict:
|
||||
original_text_layers = config_dict["text_config"]["num_hidden_layers"]
|
||||
config_dict["text_config"]["num_hidden_layers"] = text_layers
|
||||
original_layer_types = config_dict["text_config"]["layer_types"]
|
||||
config_dict["text_config"]["layer_types"] = original_layer_types[:text_layers]
|
||||
print(f"Reduced text layers from {original_text_layers} to {text_layers}")
|
||||
|
||||
original_num_experts = config_dict["text_config"]["num_local_experts"]
|
||||
config_dict["text_config"]["num_local_experts"] = num_experts
|
||||
print(f"Reduced num experts from {original_num_experts} to {num_experts}")
|
||||
|
||||
hidden_dim_divisor = 4
|
||||
|
||||
original_hidden_size = config_dict["text_config"]["hidden_size"]
|
||||
new_hidden_size = original_hidden_size // hidden_dim_divisor
|
||||
config_dict["text_config"]["hidden_size"] = new_hidden_size
|
||||
print(f"Reduced hidden size from {original_hidden_size} to {new_hidden_size}")
|
||||
|
||||
original_head_dim = config_dict["text_config"]["head_dim"]
|
||||
new_head_dim = original_head_dim // hidden_dim_divisor
|
||||
config_dict["text_config"]["head_dim"] = new_head_dim
|
||||
print(f"Reduced head dim from {original_head_dim} to {new_head_dim}")
|
||||
|
||||
# Reduce vision layers
|
||||
if "vision_config" in config_dict:
|
||||
original_vision_layers = config_dict["vision_config"]["num_hidden_layers"]
|
||||
config_dict["vision_config"]["num_hidden_layers"] = vision_layers
|
||||
print(f"Reduced vision layers from {original_vision_layers} to {vision_layers}")
|
||||
|
||||
# Update model name to indicate it's a reduced version
|
||||
config_dict["_name_or_path"] = f"reduced_maverick_{text_layers}t_{vision_layers}v"
|
||||
|
||||
return config_dict
|
||||
|
||||
|
||||
def copy_tokenizer_files(original_model_name: str, output_path: Path) -> None:
|
||||
"""Copy tokenizer files from the original model."""
|
||||
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
original_model_name, trust_remote_code=True
|
||||
)
|
||||
tokenizer.save_pretrained(output_path)
|
||||
print("Tokenizer files copied successfully")
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not copy tokenizer files: {e}")
|
||||
|
||||
|
||||
def create_preprocessor_config(original_config: Any, output_path: Path) -> None:
|
||||
"""Create preprocessor_config.json for multimodal model."""
|
||||
|
||||
# Try to load the original preprocessor config
|
||||
try:
|
||||
processor = AutoProcessor.from_pretrained(
|
||||
original_config._name_or_path
|
||||
or "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
processor.save_pretrained(output_path)
|
||||
print("Copied original preprocessor config")
|
||||
return
|
||||
except Exception as e:
|
||||
print(f"Could not copy original preprocessor config: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def create_reduced_safetensors(
|
||||
original_config: Any, reduced_config: dict[str, Any], output_path: Path
|
||||
) -> None:
|
||||
"""Create safetensors files with weights for the reduced model."""
|
||||
|
||||
print("Generating synthetic weights for reduced model...")
|
||||
|
||||
text_config = reduced_config["text_config"]
|
||||
vision_config = reduced_config["vision_config"]
|
||||
|
||||
weights = {}
|
||||
|
||||
print("Creating text model weights...")
|
||||
weights.update(create_text_model_weights(text_config))
|
||||
|
||||
print("Creating vision model weights...")
|
||||
weights.update(create_vision_model_weights(vision_config))
|
||||
|
||||
print("Creating shared model weights...")
|
||||
weights.update(create_shared_weights(text_config, vision_config))
|
||||
|
||||
print("Saving weights to safetensors files...")
|
||||
save_weights_to_safetensors(weights, output_path)
|
||||
|
||||
|
||||
def create_text_model_weights(text_config: dict[str, Any]) -> dict[str, torch.Tensor]:
|
||||
"""Create synthetic weights for the text model with MoE structure."""
|
||||
|
||||
weights = {}
|
||||
|
||||
vocab_size = text_config["vocab_size"]
|
||||
hidden_size = text_config["hidden_size"]
|
||||
intermediate_size = text_config["intermediate_size"]
|
||||
intermediate_size_mlp = text_config["intermediate_size_mlp"]
|
||||
num_layers = text_config["num_hidden_layers"]
|
||||
num_attention_heads = text_config["num_attention_heads"]
|
||||
num_key_value_heads = text_config.get("num_key_value_heads", num_attention_heads)
|
||||
|
||||
# MoE specific parameters
|
||||
num_experts = text_config.get("num_local_experts")
|
||||
assert num_experts is not None, "num_local_experts must be specified for MoE"
|
||||
|
||||
head_dim = hidden_size // num_attention_heads
|
||||
|
||||
# Embedding layers
|
||||
weights["language_model.model.embed_tokens.weight"] = torch.randn(
|
||||
vocab_size, hidden_size, dtype=torch.float16
|
||||
)
|
||||
|
||||
# Transformer layers
|
||||
for layer_idx in range(num_layers):
|
||||
layer_prefix = f"language_model.model.layers.{layer_idx}"
|
||||
print(f"Creating weights for layer {layer_prefix}...")
|
||||
|
||||
# Self-attention weights (separate q, k, v projections)
|
||||
weights[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(
|
||||
num_attention_heads * head_dim, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(
|
||||
num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(
|
||||
num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn(
|
||||
hidden_size, num_attention_heads * head_dim, dtype=torch.bfloat16
|
||||
)
|
||||
print("Self-attention weights created.")
|
||||
|
||||
# Feed-forward weights - MoE pattern based on interleave_moe_layer_step
|
||||
# For interleave_moe_layer_step=2: layers 1,3,5,... are MoE, layers
|
||||
# 0,2,4,... are dense
|
||||
interleave_step = text_config.get("interleave_moe_layer_step", 1)
|
||||
is_moe_layer = interleave_step > 0 and (layer_idx + 1) % interleave_step == 0
|
||||
|
||||
if is_moe_layer:
|
||||
# MoE layer structure
|
||||
# 1. Router weights
|
||||
weights[f"{layer_prefix}.feed_forward.router.weight"] = torch.randn(
|
||||
num_experts, hidden_size, dtype=torch.float16
|
||||
)
|
||||
|
||||
# 2. Individual expert weights (not fused)
|
||||
for expert_idx in range(num_experts):
|
||||
expert_prefix = f"{layer_prefix}.feed_forward.experts.{expert_idx}"
|
||||
|
||||
weights[f"{expert_prefix}.gate_proj.weight"] = torch.randn(
|
||||
intermediate_size, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{expert_prefix}.up_proj.weight"] = torch.randn(
|
||||
intermediate_size, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{expert_prefix}.down_proj.weight"] = torch.randn(
|
||||
hidden_size, intermediate_size, dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
# Expert weight scales (FP8 quantization)
|
||||
weights[f"{expert_prefix}.gate_proj.weight_scale"] = torch.ones(
|
||||
intermediate_size, 1, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{expert_prefix}.up_proj.weight_scale"] = torch.ones(
|
||||
intermediate_size, 1, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{expert_prefix}.down_proj.weight_scale"] = torch.ones(
|
||||
hidden_size, 1, dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
# 3. Shared expert weights
|
||||
shared_expert_prefix = f"{layer_prefix}.feed_forward.shared_expert"
|
||||
weights[f"{shared_expert_prefix}.gate_proj.weight"] = torch.randn(
|
||||
intermediate_size, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{shared_expert_prefix}.up_proj.weight"] = torch.randn(
|
||||
intermediate_size, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{shared_expert_prefix}.down_proj.weight"] = torch.randn(
|
||||
hidden_size, intermediate_size, dtype=torch.bfloat16
|
||||
)
|
||||
print(f"MoE feed-forward weights created for layer {layer_idx}.")
|
||||
else:
|
||||
# Dense layer structure
|
||||
weights[f"{layer_prefix}.feed_forward.gate_proj.weight"] = torch.randn(
|
||||
intermediate_size_mlp, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.feed_forward.up_proj.weight"] = torch.randn(
|
||||
intermediate_size_mlp, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.feed_forward.down_proj.weight"] = torch.randn(
|
||||
hidden_size, intermediate_size_mlp, dtype=torch.bfloat16
|
||||
)
|
||||
print(f"Dense feed-forward weights created for layer {layer_idx}.")
|
||||
|
||||
# Layer norms
|
||||
weights[f"{layer_prefix}.input_layernorm.weight"] = torch.ones(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.post_attention_layernorm.weight"] = torch.ones(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
print("Layer norms created.")
|
||||
|
||||
# Final layer norm and output projection
|
||||
weights["language_model.model.norm.weight"] = torch.ones(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights["language_model.lm_head.weight"] = torch.randn(
|
||||
vocab_size, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
return weights
|
||||
|
||||
|
||||
def create_vision_model_weights(
|
||||
vision_config: dict[str, Any],
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Create synthetic weights for the vision model."""
|
||||
|
||||
weights = {}
|
||||
|
||||
hidden_size = vision_config["hidden_size"]
|
||||
intermediate_size = vision_config["intermediate_size"]
|
||||
num_layers = vision_config["num_hidden_layers"]
|
||||
|
||||
# Vision transformer layers
|
||||
for layer_idx in range(num_layers):
|
||||
layer_prefix = f"vision_model.model.layers.{layer_idx}"
|
||||
|
||||
weights[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(
|
||||
hidden_size, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.self_attn.q_proj.bias"] = torch.zeros(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(
|
||||
hidden_size, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.self_attn.k_proj.bias"] = torch.zeros(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(
|
||||
hidden_size, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.self_attn.v_proj.bias"] = torch.zeros(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn(
|
||||
hidden_size, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.self_attn.o_proj.bias"] = torch.zeros(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
weights[f"{layer_prefix}.mlp.fc1.weight"] = torch.randn(
|
||||
intermediate_size, hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.mlp.fc1.bias"] = torch.zeros(
|
||||
intermediate_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.mlp.fc2.weight"] = torch.randn(
|
||||
hidden_size, intermediate_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.mlp.fc2.bias"] = torch.zeros(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
weights[f"{layer_prefix}.input_layernorm.weight"] = torch.ones(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.input_layernorm.bias"] = torch.zeros(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.post_attention_layernorm.weight"] = torch.ones(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
weights[f"{layer_prefix}.post_attention_layernorm.bias"] = torch.zeros(
|
||||
hidden_size, dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
return weights
|
||||
|
||||
|
||||
def create_shared_weights(
|
||||
text_config: dict[str, Any], vision_config: dict[str, Any]
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Create weights for shared components (vision-language connector)"""
|
||||
|
||||
weights = {}
|
||||
|
||||
text_hidden_size = text_config["hidden_size"]
|
||||
projector_input_dim = vision_config["projector_input_dim"]
|
||||
|
||||
# Vision-language connector (projects vision features to text space)
|
||||
weights["multi_modal_projector.linear_1.weight"] = torch.randn(
|
||||
text_hidden_size, projector_input_dim, dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
return weights
|
||||
|
||||
|
||||
def save_weights_to_safetensors(
|
||||
weights: dict[str, torch.Tensor], output_path: Path
|
||||
) -> None:
|
||||
"""Save weights to safetensors files and create index."""
|
||||
|
||||
# Determine how to shard the weights
|
||||
max_shard_size = 5 * 1024 * 1024 * 1024 # 5GB per shard
|
||||
|
||||
# Calculate sizes and create shards
|
||||
shards = []
|
||||
current_shard: dict[str, torch.Tensor] = {}
|
||||
current_size = 0
|
||||
|
||||
for name, tensor in weights.items():
|
||||
tensor_size = tensor.numel() * tensor.element_size()
|
||||
|
||||
if current_size + tensor_size > max_shard_size and current_shard:
|
||||
shards.append(current_shard)
|
||||
current_shard = {}
|
||||
current_size = 0
|
||||
|
||||
current_shard[name] = tensor
|
||||
current_size += tensor_size
|
||||
|
||||
if current_shard:
|
||||
shards.append(current_shard)
|
||||
|
||||
# Save shards and create index
|
||||
weight_map = {}
|
||||
|
||||
if len(shards) == 1:
|
||||
# Single file
|
||||
filename = "model.safetensors"
|
||||
save_file(shards[0], output_path / filename)
|
||||
weight_map = {name: filename for name in shards[0]}
|
||||
print(f"Saved weights to single file: {filename}")
|
||||
else:
|
||||
# Multiple shards
|
||||
for i, shard in enumerate(shards):
|
||||
filename = f"model-{i + 1:05d}-of-{len(shards):05d}.safetensors"
|
||||
save_file(shard, output_path / filename)
|
||||
for name in shard:
|
||||
weight_map[name] = filename
|
||||
print(f"Saved shard {i + 1}/{len(shards)}: {filename}")
|
||||
|
||||
# Create index file
|
||||
index_data = {
|
||||
"metadata": {
|
||||
"total_size": sum(
|
||||
tensor.numel() * tensor.element_size() for tensor in weights.values()
|
||||
)
|
||||
},
|
||||
"weight_map": weight_map,
|
||||
}
|
||||
|
||||
index_path = output_path / "model.safetensors.index.json"
|
||||
with open(index_path, "w") as f:
|
||||
json.dump(index_data, f, indent=2)
|
||||
|
||||
print(f"Created index file: {index_path}")
|
||||
print(
|
||||
f"Total model size: {index_data['metadata']['total_size'] / (1024**3):.2f} GB"
|
||||
)
|
||||
|
||||
|
||||
def check_attention_spec_interleaved_rope(
|
||||
llm: LLM,
|
||||
num_attention_layers: int,
|
||||
num_ranks: int,
|
||||
rope_layers: list[int],
|
||||
):
|
||||
"""Check that the attention spec is correct."""
|
||||
assert isinstance(llm.llm_engine.model_executor, Executor)
|
||||
kv_cache_specs_per_rank = llm.llm_engine.model_executor.get_kv_cache_specs()
|
||||
for rank in range(num_ranks):
|
||||
kv_cache_specs = kv_cache_specs_per_rank[rank]
|
||||
assert len(kv_cache_specs.keys()) == num_attention_layers
|
||||
for i in range(num_attention_layers):
|
||||
if rope_layers[i] == 0:
|
||||
expected_spec = FullAttentionSpec
|
||||
else:
|
||||
expected_spec = ChunkedLocalAttentionSpec
|
||||
assert isinstance(
|
||||
kv_cache_specs[f"language_model.model.layers.{i}.self_attn.attn"],
|
||||
expected_spec,
|
||||
)
|
||||
|
||||
|
||||
def run_reduced_model(llm: LLM, should_profile: bool = False) -> None:
|
||||
"""Test the created reduced model with vLLM."""
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=50)
|
||||
|
||||
if should_profile:
|
||||
llm.start_profile()
|
||||
outputs = llm.generate(PROMPTS, sampling_params)
|
||||
if should_profile:
|
||||
llm.stop_profile()
|
||||
|
||||
print("Test generation successful!")
|
||||
for output in outputs:
|
||||
print(f"Prompt: {output.prompt}")
|
||||
print(f"Output: {output.outputs[0].text}")
|
||||
print("-" * 40)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize(
|
||||
"original_model_name,text_layers,num_experts,vision_layers,",
|
||||
[("meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", 4, 4, 2)],
|
||||
)
|
||||
@pytest.mark.parametrize("enforce_eager", [True, False])
|
||||
@pytest.mark.parametrize("tp,ep", [(2, True)])
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
def test_dummy_maverick(
|
||||
monkeypatch,
|
||||
original_model_name: str,
|
||||
text_layers: int,
|
||||
num_experts: int,
|
||||
vision_layers: int,
|
||||
enforce_eager: bool,
|
||||
tp: int,
|
||||
ep: bool,
|
||||
output_dir: str = "/tmp/reduced_maverick",
|
||||
force_recreate: bool = True,
|
||||
profile: bool = False,
|
||||
) -> None:
|
||||
# Disable multiprocessing allows us to access model executor from LLM engine
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
|
||||
model_path = create_reduced_maverick_model(
|
||||
original_model_name=original_model_name,
|
||||
output_dir=output_dir,
|
||||
text_layers=text_layers,
|
||||
num_experts=num_experts,
|
||||
vision_layers=vision_layers,
|
||||
force_recreate=force_recreate,
|
||||
)
|
||||
|
||||
print(f"\nReduced model created successfully at: {model_path}")
|
||||
|
||||
rope_layers = get_rope_layers_config(model_path)
|
||||
|
||||
llm = LLM(
|
||||
model=model_path,
|
||||
trust_remote_code=True,
|
||||
max_model_len=512, # Small context for testing
|
||||
gpu_memory_utilization=0.3, # Conservative memory usage
|
||||
enforce_eager=enforce_eager,
|
||||
tensor_parallel_size=tp,
|
||||
enable_expert_parallel=ep,
|
||||
)
|
||||
|
||||
check_attention_spec_interleaved_rope(
|
||||
llm,
|
||||
text_layers,
|
||||
tp,
|
||||
rope_layers,
|
||||
)
|
||||
|
||||
print(f"\nTesting reduced model at {model_path}...")
|
||||
run_reduced_model(llm=llm, should_profile=profile)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to create and test the reduced model."""
|
||||
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Create a reduced-layer Maverick model"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default="/tmp/reduced_maverick",
|
||||
help="Output directory for the reduced model",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text-layers",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Number of text transformer layers",
|
||||
)
|
||||
parser.add_argument("--num-experts", type=int, default=4, help="Number of experts")
|
||||
parser.add_argument(
|
||||
"--vision-layers",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Number of vision transformer layers",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force-recreate",
|
||||
action="store_true",
|
||||
help="Force recreation if output directory exists",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--test", action="store_true", help="Test the created model with vLLM"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile", action="store_true", help="Profile the created model with vLLM"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--test-original",
|
||||
action="store_true",
|
||||
help="Test the original model with vLLM",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--original-model",
|
||||
default="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
|
||||
help="Original model name to base the reduction on",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.test:
|
||||
test_dummy_maverick(
|
||||
original_model_name=args.original_model,
|
||||
output_dir=args.output_dir,
|
||||
text_layers=args.text_layers,
|
||||
num_experts=args.num_experts,
|
||||
vision_layers=args.vision_layers,
|
||||
force_recreate=args.force_recreate,
|
||||
tp=2,
|
||||
ep=True,
|
||||
enforce_eager=True,
|
||||
profile=args.profile,
|
||||
)
|
||||
|
||||
if args.test_original:
|
||||
run_maverick_serving(args.original_model)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main())
|
||||
@@ -0,0 +1,182 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import gc
|
||||
import random
|
||||
import string
|
||||
import sys
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.models.registry import HF_EXAMPLE_MODELS
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.mem_utils import KiB_bytes, MiB_bytes, format_mib
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen3-VL-4B-Instruct"
|
||||
RANDOM_PREFIX_LEN = 100
|
||||
TEST_IMAGE_NAMES = [
|
||||
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg",
|
||||
"Grayscale_8bits_palette_sample_image.png",
|
||||
]
|
||||
MAX_MODEL_LEN = 8192
|
||||
REQUESTS_PER_ROUND = 4
|
||||
WARMUP_ROUNDS = 2
|
||||
MEASURED_ROUNDS = 16
|
||||
GPU_GROWTH_THRESHOLD_MIB = 0
|
||||
CPU_PEAK_GROWTH_THRESHOLD_MIB = 0
|
||||
|
||||
SAMPLING_PARAMS = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=16,
|
||||
)
|
||||
|
||||
|
||||
def _make_messages(image_url: str) -> list[ChatCompletionMessageParam]:
|
||||
# Avoid obscuring memory leaks because of prefix caching
|
||||
random_text = "".join(random.choices(string.ascii_uppercase, k=RANDOM_PREFIX_LEN))
|
||||
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"Ignore this random string: {random_text}",
|
||||
},
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Describe this image in one short sentence.",
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _build_request_batch(
|
||||
image_urls: list[str],
|
||||
) -> list[list[ChatCompletionMessageParam]]:
|
||||
return [
|
||||
_make_messages(image_urls[i % len(image_urls)])
|
||||
for i in range(REQUESTS_PER_ROUND)
|
||||
]
|
||||
|
||||
|
||||
def _ru_maxrss_bytes() -> int | None:
|
||||
try:
|
||||
import resource
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
|
||||
if rss <= 0:
|
||||
return 0
|
||||
|
||||
# Linux reports kilobytes, macOS reports bytes.
|
||||
return rss if sys.platform == "darwin" else rss * KiB_bytes
|
||||
|
||||
|
||||
def _gpu_used_bytes() -> int:
|
||||
torch.accelerator.synchronize()
|
||||
free_bytes, total_bytes = torch.accelerator.get_memory_info()
|
||||
return int(total_bytes - free_bytes)
|
||||
|
||||
|
||||
def _format_mib(num_bytes: int | None) -> str:
|
||||
if num_bytes is None:
|
||||
return "n/a"
|
||||
|
||||
return f"{format_mib(num_bytes)} MiB"
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def llm(monkeypatch):
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
|
||||
# pytest caches the fixture so we use weakref.proxy to
|
||||
# enable garbage collection
|
||||
llm_kwargs = dict(
|
||||
model=MODEL_NAME,
|
||||
enforce_eager=True,
|
||||
max_model_len=MAX_MODEL_LEN,
|
||||
max_num_seqs=REQUESTS_PER_ROUND,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
seed=0,
|
||||
disable_log_stats=True,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
if current_platform.is_rocm():
|
||||
llm_kwargs["attention_backend"] = "TRITON_ATTN"
|
||||
|
||||
llm = LLM(**llm_kwargs)
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("image_urls", [TEST_IMAGE_NAMES], indirect=True)
|
||||
def test_no_memory_leak(llm, image_urls: list[str]) -> None:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(MODEL_NAME)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
request_batch = _build_request_batch(image_urls)
|
||||
|
||||
# Establish a warmup baseline after model load and the first multimodal
|
||||
# requests complete. Later rounds should remain near this steady state.
|
||||
for _ in range(WARMUP_ROUNDS):
|
||||
outputs = llm.chat(request_batch, sampling_params=SAMPLING_PARAMS)
|
||||
assert len(outputs) == len(request_batch)
|
||||
assert llm.llm_engine.get_num_unfinished_requests() == 0
|
||||
del outputs
|
||||
|
||||
gc.collect()
|
||||
warmup_gpu = _gpu_used_bytes()
|
||||
warmup_cpu_peak = _ru_maxrss_bytes()
|
||||
|
||||
post_warmup_gpu_samples: list[int] = []
|
||||
post_warmup_cpu_peak_samples: list[int] = []
|
||||
|
||||
for _ in range(MEASURED_ROUNDS):
|
||||
outputs = llm.chat(request_batch, sampling_params=SAMPLING_PARAMS)
|
||||
assert len(outputs) == len(request_batch)
|
||||
assert llm.llm_engine.get_num_unfinished_requests() == 0
|
||||
del outputs
|
||||
|
||||
gc.collect()
|
||||
post_warmup_gpu_samples.append(_gpu_used_bytes())
|
||||
cpu_peak = _ru_maxrss_bytes()
|
||||
if cpu_peak is not None:
|
||||
post_warmup_cpu_peak_samples.append(cpu_peak)
|
||||
|
||||
gpu_growth = max(post_warmup_gpu_samples) - warmup_gpu
|
||||
gpu_threshold = GPU_GROWTH_THRESHOLD_MIB * MiB_bytes
|
||||
|
||||
assert gpu_growth <= gpu_threshold, (
|
||||
"Qwen3-VL GPU memory kept growing after warmup. "
|
||||
f"warmup_baseline={_format_mib(warmup_gpu)}, "
|
||||
f"post_warmup_samples={[_format_mib(x) for x in post_warmup_gpu_samples]}, "
|
||||
f"gpu_growth={_format_mib(gpu_growth)}, "
|
||||
f"gpu_threshold={GPU_GROWTH_THRESHOLD_MIB} MiB"
|
||||
)
|
||||
|
||||
if warmup_cpu_peak is not None and post_warmup_cpu_peak_samples:
|
||||
cpu_peak_growth = max(post_warmup_cpu_peak_samples) - warmup_cpu_peak
|
||||
cpu_threshold = CPU_PEAK_GROWTH_THRESHOLD_MIB * MiB_bytes
|
||||
|
||||
assert cpu_peak_growth <= cpu_threshold, (
|
||||
"Qwen3-VL CPU peak RSS kept growing after warmup. "
|
||||
f"warmup_ru_maxrss={_format_mib(warmup_cpu_peak)}, "
|
||||
f"post_warmup_ru_maxrss={[_format_mib(x) for x in post_warmup_cpu_peak_samples]}, " # noqa: E501
|
||||
f"cpu_peak_growth={_format_mib(cpu_peak_growth)}, "
|
||||
f"cpu_peak_threshold={CPU_PEAK_GROWTH_THRESHOLD_MIB} MiB"
|
||||
)
|
||||
@@ -0,0 +1,119 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoModelForImageTextToText
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....conftest import HfRunner, ImageTestAssets, VllmRunner
|
||||
from .vlm_utils import model_utils
|
||||
|
||||
MODEL = "google/gemma-3-4b-it"
|
||||
PROMPT = (
|
||||
"<bos><start_of_turn>user\n"
|
||||
"<start_of_image>What is the content in the center of the image?"
|
||||
"<end_of_turn>\n<start_of_turn>model\n"
|
||||
)
|
||||
|
||||
|
||||
def _install_prefill_hidden_capture(model):
|
||||
model = getattr(model, "module", model)
|
||||
model._prefill_hidden = None
|
||||
|
||||
language_model = model.language_model.model
|
||||
original_forward = language_model.forward
|
||||
|
||||
def forward(*args, **kwargs):
|
||||
hidden_states = original_forward(*args, **kwargs)
|
||||
if model._prefill_hidden is None and torch.is_tensor(hidden_states):
|
||||
model._prefill_hidden = hidden_states.detach().float().cpu()
|
||||
return hidden_states
|
||||
|
||||
language_model.forward = forward
|
||||
|
||||
|
||||
def _get_prefill_hidden(model):
|
||||
model = getattr(model, "module", model)
|
||||
hidden = getattr(model, "_prefill_hidden", None)
|
||||
assert hidden is not None
|
||||
return hidden
|
||||
|
||||
|
||||
def _get_hf_prefill_hidden(hf_model: HfRunner, image: Any):
|
||||
inputs = hf_model.get_inputs([PROMPT], images=[image])[0]
|
||||
with torch.no_grad():
|
||||
outputs = hf_model.model.model(
|
||||
**hf_model.wrap_device(inputs),
|
||||
use_cache=False,
|
||||
)
|
||||
return outputs.last_hidden_state[0].detach().float().cpu()
|
||||
|
||||
|
||||
def _get_vllm_prefill_hidden(
|
||||
vllm_runner: type[VllmRunner],
|
||||
image: Any,
|
||||
vllm_runner_kwargs: dict[str, Any],
|
||||
):
|
||||
with vllm_runner(
|
||||
MODEL,
|
||||
max_model_len=4096,
|
||||
max_num_seqs=2,
|
||||
enforce_eager=True,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
**vllm_runner_kwargs,
|
||||
) as vllm_model:
|
||||
vllm_model.apply_model(_install_prefill_hidden_capture)
|
||||
vllm_model.generate_greedy([PROMPT], max_tokens=1, images=[image])
|
||||
return vllm_model.apply_model(_get_prefill_hidden)[0]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_rocm(), reason="ROCm attention has accuracy issue for this test"
|
||||
)
|
||||
def test_mm_prefix_lm_e2e(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
image_assets: ImageTestAssets,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Regression: Gemma3 native prefill must apply image prefix-LM mask."""
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
image = image_assets[0].pil_image
|
||||
|
||||
vllm_runner_kwargs: dict[str, Any] = {
|
||||
"mm_processor_cache_gb": 0,
|
||||
"mm_processor_kwargs": {"do_pan_and_scan": True},
|
||||
}
|
||||
vllm_hidden = _get_vllm_prefill_hidden(vllm_runner, image, vllm_runner_kwargs)
|
||||
|
||||
hf_model = hf_runner(
|
||||
MODEL,
|
||||
auto_cls=AutoModelForImageTextToText,
|
||||
)
|
||||
hf_model = model_utils.gemma3_patch_hf_runner(hf_model)
|
||||
|
||||
with hf_model:
|
||||
hf_hidden = _get_hf_prefill_hidden(hf_model, image)
|
||||
|
||||
assert vllm_hidden.shape == hf_hidden.shape
|
||||
|
||||
full_cos = torch.nn.functional.cosine_similarity(
|
||||
vllm_hidden.flatten(), hf_hidden.flatten(), dim=0
|
||||
)
|
||||
image_cos = torch.nn.functional.cosine_similarity(
|
||||
vllm_hidden[1:769].flatten(), hf_hidden[1:769].flatten(), dim=0
|
||||
)
|
||||
|
||||
assert full_cos > 0.9, (
|
||||
"Gemma3 mm-prefix-LM full prefill hidden states should be close to HF; "
|
||||
f"got {full_cos=}"
|
||||
)
|
||||
assert image_cos > 0.9, (
|
||||
"Gemma3 mm-prefix-LM image prefill hidden states should be close to HF; "
|
||||
f"got {image_cos=}"
|
||||
)
|
||||
@@ -0,0 +1,176 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Generation tests for Moondream3 query and caption support."""
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.registry import HF_EXAMPLE_MODELS
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....conftest import IMAGE_ASSETS, ImageTestAssets
|
||||
from ....utils import large_gpu_mark, multi_gpu_test
|
||||
|
||||
MOONDREAM3_MODEL_ID = "moondream/moondream3-preview"
|
||||
MOONDREAM3_TOKENIZER = "moondream/starmie-v1"
|
||||
|
||||
HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
|
||||
{
|
||||
"stop_sign": "<|endoftext|><image><|md_reserved_0|>query<|md_reserved_1|>What color is the stop sign?<|md_reserved_2|>", # noqa: E501
|
||||
"cherry_blossom": "<|endoftext|><image><|md_reserved_0|>query<|md_reserved_1|>What color are the flowers?<|md_reserved_2|>", # noqa: E501
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def make_query_prompt(question: str) -> str:
|
||||
"""Create a direct-answer query prompt for Moondream3."""
|
||||
return (
|
||||
"<|endoftext|><image><|md_reserved_0|>query<|md_reserved_1|>"
|
||||
f"{question}<|md_reserved_2|>"
|
||||
)
|
||||
|
||||
|
||||
def make_caption_prompt(length: str = "normal") -> str:
|
||||
"""Create a caption prompt for Moondream3."""
|
||||
return (
|
||||
"<|endoftext|><image><|md_reserved_0|>"
|
||||
f"describe<|md_reserved_1|>{length}<|md_reserved_2|>"
|
||||
)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@large_gpu_mark(min_gb=80)
|
||||
def test_tensor_parallel(image_assets: ImageTestAssets):
|
||||
import gc
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed.parallel_state import destroy_model_parallel
|
||||
|
||||
destroy_model_parallel()
|
||||
gc.collect()
|
||||
current_platform.empty_cache()
|
||||
|
||||
llm = LLM(
|
||||
model=MOONDREAM3_MODEL_ID,
|
||||
tokenizer=MOONDREAM3_TOKENIZER,
|
||||
trust_remote_code=True,
|
||||
dtype="bfloat16",
|
||||
tensor_parallel_size=2,
|
||||
max_model_len=1024,
|
||||
enforce_eager=True,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
gpu_memory_utilization=0.45,
|
||||
)
|
||||
|
||||
image = image_assets[0].pil_image
|
||||
prompt = make_query_prompt("What color is the stop sign?")
|
||||
|
||||
try:
|
||||
outputs = llm.generate(
|
||||
{"prompt": prompt, "multi_modal_data": {"image": image}},
|
||||
SamplingParams(max_tokens=20, temperature=0),
|
||||
)
|
||||
|
||||
assert len(outputs) > 0
|
||||
assert outputs[0].outputs[0].text is not None
|
||||
finally:
|
||||
del llm
|
||||
gc.collect()
|
||||
current_platform.empty_cache()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm():
|
||||
model_info = HF_EXAMPLE_MODELS.get_hf_info("Moondream3ForCausalLM")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
from vllm import LLM
|
||||
|
||||
try:
|
||||
return LLM(
|
||||
model=MOONDREAM3_MODEL_ID,
|
||||
tokenizer=MOONDREAM3_TOKENIZER,
|
||||
trust_remote_code=True,
|
||||
dtype="bfloat16",
|
||||
max_model_len=2048,
|
||||
enforce_eager=True,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
gpu_memory_utilization=0.45,
|
||||
)
|
||||
except Exception as exc:
|
||||
pytest.skip(f"Failed to load {MOONDREAM3_MODEL_ID}: {exc}")
|
||||
|
||||
|
||||
@large_gpu_mark(min_gb=48)
|
||||
def test_model_loading(llm):
|
||||
assert llm is not None
|
||||
|
||||
|
||||
@large_gpu_mark(min_gb=48)
|
||||
def test_query_skill(llm, image_assets: ImageTestAssets):
|
||||
from vllm import SamplingParams
|
||||
|
||||
image = image_assets[0].pil_image
|
||||
prompt = make_query_prompt("What color is the stop sign?")
|
||||
|
||||
outputs = llm.generate(
|
||||
{"prompt": prompt, "multi_modal_data": {"image": image}},
|
||||
SamplingParams(max_tokens=50, temperature=0),
|
||||
)
|
||||
|
||||
output_text = outputs[0].outputs[0].text
|
||||
assert output_text is not None
|
||||
assert len(output_text) > 0
|
||||
|
||||
|
||||
@large_gpu_mark(min_gb=48)
|
||||
def test_caption_skill(llm, image_assets: ImageTestAssets):
|
||||
from vllm import SamplingParams
|
||||
|
||||
image = image_assets[1].pil_image
|
||||
prompt = make_caption_prompt()
|
||||
|
||||
outputs = llm.generate(
|
||||
{"prompt": prompt, "multi_modal_data": {"image": image}},
|
||||
SamplingParams(max_tokens=100, temperature=0),
|
||||
)
|
||||
|
||||
output_text = outputs[0].outputs[0].text
|
||||
assert output_text is not None
|
||||
assert len(output_text) > 0
|
||||
|
||||
|
||||
@large_gpu_mark(min_gb=48)
|
||||
def test_batched_inference(llm, image_assets: ImageTestAssets):
|
||||
from vllm import SamplingParams
|
||||
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
prompts = [
|
||||
{"prompt": prompt, "multi_modal_data": {"image": img}}
|
||||
for img, prompt in zip(images, HF_IMAGE_PROMPTS)
|
||||
]
|
||||
|
||||
outputs = llm.generate(prompts, SamplingParams(max_tokens=50, temperature=0))
|
||||
|
||||
assert len(outputs) == len(images)
|
||||
for output in outputs:
|
||||
assert output.outputs[0].text is not None
|
||||
assert len(output.outputs[0].text) > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asset_name", ["stop_sign", "cherry_blossom"])
|
||||
@large_gpu_mark(min_gb=48)
|
||||
def test_image_assets(llm, image_assets: ImageTestAssets, asset_name: str):
|
||||
from vllm import SamplingParams
|
||||
|
||||
asset_idx = 0 if asset_name == "stop_sign" else 1
|
||||
image = image_assets[asset_idx].pil_image
|
||||
prompt = HF_IMAGE_PROMPTS[asset_idx]
|
||||
|
||||
outputs = llm.generate(
|
||||
{"prompt": prompt, "multi_modal_data": {"image": image}},
|
||||
SamplingParams(max_tokens=50, temperature=0),
|
||||
)
|
||||
|
||||
output_text = outputs[0].outputs[0].text
|
||||
assert output_text is not None
|
||||
assert len(output_text) > 0
|
||||
@@ -0,0 +1,155 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.model_executor.models.moss_audio import MOSS_AUDIO_PLACEHOLDER
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ...registry import HF_EXAMPLE_MODELS
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
CORE_MODEL = pytest.param(
|
||||
"OpenMOSS-Team/MOSS-Audio-4B-Instruct",
|
||||
marks=pytest.mark.core_model,
|
||||
id="4b-instruct",
|
||||
)
|
||||
|
||||
EXTENDED_MODELS = [
|
||||
"OpenMOSS-Team/MOSS-Audio-4B-Thinking",
|
||||
"OpenMOSS-Team/MOSS-Audio-8B-Instruct",
|
||||
"OpenMOSS-Team/MOSS-Audio-8B-Thinking",
|
||||
]
|
||||
|
||||
ACCURACY_MODELS = [CORE_MODEL, *EXTENDED_MODELS]
|
||||
|
||||
PARALLEL_SMOKE_CASES = [
|
||||
pytest.param({"tensor_parallel_size": 2}, id="tp2"),
|
||||
pytest.param({"pipeline_parallel_size": 2}, id="pp2"),
|
||||
pytest.param(
|
||||
{"tensor_parallel_size": 2, "pipeline_parallel_size": 2},
|
||||
id="tp2_pp2",
|
||||
),
|
||||
]
|
||||
|
||||
HF_ACCURACY_SKIP_REASON = (
|
||||
"HF AutoModelForCausalLM cannot load remote MOSS-Audio configs; "
|
||||
"vLLM generation coverage is provided by the smoke tests below."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
def test_moss_audio_generation_smoke(vllm_runner) -> None:
|
||||
model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nBriefly describe this audio."]
|
||||
audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype="half",
|
||||
enforce_eager=True,
|
||||
max_model_len=1024,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
trust_remote_code=True,
|
||||
) as vllm_model:
|
||||
outputs = vllm_model.generate_greedy(
|
||||
prompts,
|
||||
max_tokens=4,
|
||||
audios=audios,
|
||||
)
|
||||
|
||||
assert len(outputs) == 1
|
||||
assert len(outputs[0][1]) > 0
|
||||
|
||||
|
||||
@pytest.mark.skip(reason=HF_ACCURACY_SKIP_REASON)
|
||||
@pytest.mark.parametrize("model", ACCURACY_MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", [8])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
def test_moss_audio_hf_vllm_accuracy(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nTranscribe this audio."]
|
||||
audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
enforce_eager=True,
|
||||
max_model_len=1024,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
trust_remote_code=True,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
audios=audios,
|
||||
)
|
||||
|
||||
with hf_runner(model, dtype=dtype, trust_remote_code=True) as hf_model:
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
audios=audios,
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("parallel_kwargs", PARALLEL_SMOKE_CASES)
|
||||
def test_moss_audio_parallel_smoke(vllm_runner, parallel_kwargs) -> None:
|
||||
model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
|
||||
required_gpus = parallel_kwargs.get(
|
||||
"tensor_parallel_size", 1
|
||||
) * parallel_kwargs.get("pipeline_parallel_size", 1)
|
||||
if current_platform.device_count() < required_gpus:
|
||||
# TP/PP integration smoke runs on local or multi-GPU CI only.
|
||||
pytest.skip(f"Requires at least {required_gpus} GPUs")
|
||||
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nBriefly describe this audio."]
|
||||
audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype="half",
|
||||
enforce_eager=True,
|
||||
max_model_len=1024,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
trust_remote_code=True,
|
||||
**parallel_kwargs,
|
||||
) as vllm_model:
|
||||
outputs = vllm_model.generate_greedy(
|
||||
prompts,
|
||||
max_tokens=4,
|
||||
audios=audios,
|
||||
)
|
||||
|
||||
assert len(outputs) == 1
|
||||
assert len(outputs[0][1]) > 0
|
||||
@@ -0,0 +1,124 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from collections.abc import Iterable, Sequence
|
||||
|
||||
import pytest
|
||||
import regex as re
|
||||
from transformers import AutoModel
|
||||
|
||||
from tests.models.utils import check_logprobs_close
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.logprobs import Logprob, SampleLogprobs
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
|
||||
from ....conftest import HfRunner, PromptImageInput, VllmRunner
|
||||
|
||||
IMAGE = ImageAsset("paper-11").pil_image_ext(ext="png").convert("RGB")
|
||||
PROMPT = (
|
||||
"</s><s><predict_bbox><predict_classes><output_markdown><predict_no_text_in_pic>"
|
||||
)
|
||||
|
||||
|
||||
class DummyLogprobs(dict[int, Logprob]):
|
||||
def __init__(self, vocab_ids: Iterable[int]):
|
||||
super().__init__(dict.fromkeys(vocab_ids, Logprob(0.0)))
|
||||
|
||||
def __repr__(self):
|
||||
return "DummyLogprobs()"
|
||||
|
||||
|
||||
def mask_bbox_tokens(
|
||||
output: tuple[list[int], str, SampleLogprobs],
|
||||
tokenizer: TokenizerLike,
|
||||
) -> tuple[list[int], str, SampleLogprobs]:
|
||||
"""
|
||||
Always pass check_logprobs_close check for bounding box tokens
|
||||
because it is reasonable for them to differ slightly.
|
||||
"""
|
||||
ignore_pattern = r"<[xy]_[\d.]+>"
|
||||
vocab = tokenizer.get_vocab()
|
||||
|
||||
output_ids, output_str, out_logprobs = output
|
||||
|
||||
masked_logprobs = list[dict[int, Logprob]]()
|
||||
for token, logprobs in zip(output_ids, out_logprobs):
|
||||
if re.match(ignore_pattern, tokenizer.decode(token)):
|
||||
masked_logprobs.append(DummyLogprobs(vocab.values()))
|
||||
else:
|
||||
masked_logprobs.append(logprobs)
|
||||
|
||||
return output_ids, output_str, masked_logprobs
|
||||
|
||||
|
||||
def run_test(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
inputs: Sequence[tuple[list[str], PromptImageInput]],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
"""Verify that the inference result is the same between hf and vllm."""
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_num_seqs=64,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
trust_remote_code=True,
|
||||
) as vllm_model:
|
||||
vllm_outputs_per_case = [
|
||||
vllm_model.generate_greedy_logprobs(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
images=images,
|
||||
)
|
||||
for prompts, images in inputs
|
||||
]
|
||||
|
||||
tokenizer = vllm_model.llm.get_tokenizer()
|
||||
|
||||
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
|
||||
hf_outputs_per_case = [
|
||||
hf_model.generate_greedy_logprobs_limit(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
images=images,
|
||||
tokenization_kwargs={"add_special_tokens": False},
|
||||
)
|
||||
for prompts, images in inputs
|
||||
]
|
||||
|
||||
for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=[
|
||||
mask_bbox_tokens(output, tokenizer) for output in hf_outputs
|
||||
],
|
||||
outputs_1_lst=[
|
||||
mask_bbox_tokens(output, tokenizer) for output in vllm_outputs
|
||||
],
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["nvidia/NVIDIA-Nemotron-Parse-v1.2"])
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
def test_models(
|
||||
hf_runner, vllm_runner, model: str, dtype: str, num_logprobs: int
|
||||
) -> None:
|
||||
run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
inputs=[
|
||||
([PROMPT] * 10, [IMAGE] * 10),
|
||||
],
|
||||
model=model,
|
||||
dtype=dtype,
|
||||
max_tokens=100,
|
||||
num_logprobs=num_logprobs,
|
||||
)
|
||||
@@ -0,0 +1,315 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
|
||||
import pytest
|
||||
import regex as re
|
||||
from huggingface_hub import snapshot_download
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.logprobs import SampleLogprobs
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.multimodal.image import convert_image_mode, rescale_image_size
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
from ....conftest import (
|
||||
IMAGE_ASSETS,
|
||||
HfRunner,
|
||||
PromptAudioInput,
|
||||
PromptImageInput,
|
||||
VllmRunner,
|
||||
)
|
||||
from ....utils import large_gpu_test
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
|
||||
{
|
||||
"stop_sign": "<|user|>\n<|image_1|>\nWhat's the content of the image?<|end|>\n<|assistant|>\n", # noqa: E501
|
||||
"cherry_blossom": "<|user|>\n<|image_1|>\nPlease infer the season with reason in details.<|end|>\n<|assistant|>\n", # noqa: E501
|
||||
}
|
||||
)
|
||||
HF_MULTIIMAGE_IMAGE_PROMPT = (
|
||||
"<|user|>\n<|image_1|>\n<|image_2|>\nDescribe these images.<|end|>\n<|assistant|>\n" # noqa: E501
|
||||
)
|
||||
|
||||
model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct")
|
||||
# Since the vision-lora and speech-lora co-exist with the base model,
|
||||
# we have to manually specify the path of the lora weights.
|
||||
vision_lora_path = os.path.join(model_path, "vision-lora")
|
||||
speech_question = os.path.join(
|
||||
model_path, "examples", "what_is_shown_in_this_image.wav"
|
||||
)
|
||||
models = [model_path]
|
||||
|
||||
|
||||
def vllm_to_hf_output(
|
||||
vllm_output: tuple[list[int], str, SampleLogprobs | None], model: str
|
||||
):
|
||||
"""Sanitize vllm output to be comparable with hf output."""
|
||||
_, output_str, out_logprobs = vllm_output
|
||||
|
||||
output_str_without_image = re.sub(r"(<\|image_\d+\|>)+", "", output_str)
|
||||
assert output_str_without_image[0] == " "
|
||||
output_str_without_image = output_str_without_image[1:]
|
||||
|
||||
hf_output_str = output_str_without_image + "<|end|><|endoftext|>"
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model)
|
||||
hf_output_ids = tokenizer.encode(output_str_without_image)
|
||||
assert hf_output_ids[0] == 1
|
||||
hf_output_ids = hf_output_ids[1:]
|
||||
|
||||
return hf_output_ids, hf_output_str, out_logprobs
|
||||
|
||||
|
||||
target_dtype = "half"
|
||||
|
||||
|
||||
def run_test(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
inputs: Sequence[tuple[list[str], PromptImageInput, PromptAudioInput | None]],
|
||||
model: str,
|
||||
*,
|
||||
max_model_len: int,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
mm_limit: int,
|
||||
tensor_parallel_size: int,
|
||||
distributed_executor_backend: str | None = None,
|
||||
):
|
||||
"""Inference result should be the same between hf and vllm.
|
||||
|
||||
All the image fixtures for the test are from IMAGE_ASSETS.
|
||||
For huggingface runner, we provide the PIL images as input.
|
||||
For vllm runner, we provide MultiModalDataDict objects
|
||||
and corresponding MultiModalConfig as input.
|
||||
Note, the text input is also adjusted to abide by vllm contract.
|
||||
The text output is sanitized to be able to compare with hf.
|
||||
"""
|
||||
# NOTE: take care of the order. run vLLM first, and then run HF.
|
||||
# vLLM needs a fresh new process without cuda initialization.
|
||||
# if we run HF first, the cuda initialization will be done and it
|
||||
# will hurt multiprocessing backend with fork method (the default method).
|
||||
# max_model_len should be greater than image_feature_size
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="generate",
|
||||
max_model_len=max_model_len,
|
||||
max_num_seqs=2,
|
||||
dtype=dtype,
|
||||
limit_mm_per_prompt={"image": mm_limit},
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
enable_lora=True,
|
||||
max_lora_rank=320,
|
||||
gpu_memory_utilization=0.8, # set to 0.8 to avoid OOM in CI
|
||||
enforce_eager=True,
|
||||
) as vllm_model:
|
||||
lora_request = LoRARequest("vision", 1, vision_lora_path)
|
||||
vllm_outputs_per_case = [
|
||||
vllm_model.generate_greedy_logprobs(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
images=images,
|
||||
audios=audios,
|
||||
lora_request=lora_request,
|
||||
)
|
||||
for prompts, images, audios in inputs
|
||||
]
|
||||
|
||||
# This error occurs inside `get_peft_model`
|
||||
# FIXME: https://huggingface.co/microsoft/Phi-4-multimodal-instruct/discussions/75
|
||||
pytest.skip("HF impl is not compatible with current transformers")
|
||||
|
||||
hf_model_kwargs = {"_attn_implementation": "sdpa"}
|
||||
with hf_runner(model, dtype=dtype, model_kwargs=hf_model_kwargs) as hf_model:
|
||||
hf_processor = hf_model.processor
|
||||
eos_token_id = hf_processor.tokenizer.eos_token_id
|
||||
|
||||
def patch_hf_processor(
|
||||
*args, text="", images=None, audio=None, sampling_rate=None, **kwargs
|
||||
):
|
||||
audios = None
|
||||
if audio is not None and sampling_rate is not None:
|
||||
audios = [(audio, sampling_rate)]
|
||||
return hf_processor(
|
||||
*args, text=text, images=images, audios=audios, **kwargs
|
||||
)
|
||||
|
||||
hf_model.processor = patch_hf_processor
|
||||
|
||||
hf_outputs_per_case = [
|
||||
hf_model.generate_greedy_logprobs_limit(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
images=images,
|
||||
audios=audios,
|
||||
eos_token_id=eos_token_id,
|
||||
num_logits_to_keep=0,
|
||||
)
|
||||
for prompts, images, audios in inputs
|
||||
]
|
||||
|
||||
for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize(
|
||||
"size_factors",
|
||||
[
|
||||
# Single-scale
|
||||
[1.0],
|
||||
# Single-scale, batched
|
||||
[1.0, 1.0, 1.0],
|
||||
# Multi-scale
|
||||
[0.25, 0.5, 1.0],
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [target_dtype])
|
||||
@pytest.mark.parametrize("max_model_len", [12800])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize("num_logprobs", [10])
|
||||
def test_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
image_assets,
|
||||
model,
|
||||
size_factors,
|
||||
dtype: str,
|
||||
max_model_len: int,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
|
||||
inputs_per_image = [
|
||||
(
|
||||
[prompt for _ in size_factors],
|
||||
[rescale_image_size(image, factor) for factor in size_factors],
|
||||
None,
|
||||
)
|
||||
for image, prompt in zip(images, HF_IMAGE_PROMPTS)
|
||||
]
|
||||
|
||||
run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
inputs_per_image,
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_model_len=max_model_len,
|
||||
max_tokens=max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
mm_limit=1,
|
||||
tensor_parallel_size=1,
|
||||
)
|
||||
|
||||
|
||||
@large_gpu_test(min_gb=48)
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize(
|
||||
"size_factors",
|
||||
[
|
||||
# No image
|
||||
# [],
|
||||
# Single-scale
|
||||
[1.0],
|
||||
# Single-scale, batched
|
||||
[1.0, 1.0, 1.0],
|
||||
# Multi-scale
|
||||
[0.25, 0.5, 1.0],
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [target_dtype])
|
||||
@pytest.mark.parametrize("max_model_len", [25600])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize("num_logprobs", [10])
|
||||
def test_multi_images_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
image_assets,
|
||||
model,
|
||||
size_factors,
|
||||
dtype: str,
|
||||
max_model_len: int,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
|
||||
inputs_per_case = [
|
||||
(
|
||||
[HF_MULTIIMAGE_IMAGE_PROMPT for _ in size_factors],
|
||||
[
|
||||
[rescale_image_size(image, factor) for image in images]
|
||||
for factor in size_factors
|
||||
],
|
||||
None,
|
||||
),
|
||||
]
|
||||
|
||||
run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
inputs_per_case,
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_model_len=max_model_len,
|
||||
max_tokens=max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
mm_limit=2,
|
||||
tensor_parallel_size=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize("dtype", [target_dtype])
|
||||
@pytest.mark.parametrize("max_model_len", [12800])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize("num_logprobs", [10])
|
||||
def test_vision_speech_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model,
|
||||
dtype: str,
|
||||
max_model_len: int,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
# use the example speech question so that the model outputs are reasonable
|
||||
audio = load_audio(speech_question, sr=None)
|
||||
image = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
|
||||
|
||||
inputs_vision_speech = [
|
||||
(
|
||||
["<|user|><|image_1|><|audio_1|><|end|><|assistant|>"],
|
||||
[image],
|
||||
[audio],
|
||||
),
|
||||
]
|
||||
|
||||
run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
inputs_vision_speech,
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_model_len=max_model_len,
|
||||
max_tokens=max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
mm_limit=1,
|
||||
tensor_parallel_size=1,
|
||||
)
|
||||
@@ -0,0 +1,198 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Sequence
|
||||
from importlib.metadata import version
|
||||
|
||||
import pytest
|
||||
import regex as re
|
||||
from packaging.version import Version
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from vllm.logprobs import SampleLogprobs
|
||||
from vllm.multimodal.image import rescale_image_size
|
||||
|
||||
from ....conftest import (
|
||||
IMAGE_ASSETS,
|
||||
HfRunner,
|
||||
PromptImageInput,
|
||||
VllmRunner,
|
||||
)
|
||||
from ....utils import multi_gpu_test
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
Version("5.0") <= Version(version("transformers")),
|
||||
reason=(
|
||||
"vllm upgraded transformers above v5.4 where HF model custom code uses siglip2 "
|
||||
"internals (filter_out_non_signature_kwargs) removed by "
|
||||
"huggingface/transformers#43514"
|
||||
),
|
||||
)
|
||||
|
||||
MODEL_ID = "microsoft/Phi-4-reasoning-vision-15B"
|
||||
|
||||
HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
|
||||
{
|
||||
"stop_sign": "<|user|>\n<image>\nWhat's the content of the image?<|end|>\n<|assistant|>\n", # noqa: E501
|
||||
"cherry_blossom": "<|user|>\n<image>\nPlease infer the season with reason in details.<|end|>\n<|assistant|>\n", # noqa: E501
|
||||
}
|
||||
)
|
||||
HF_MULTIIMAGE_IMAGE_PROMPT = (
|
||||
"<|user|>\n<image>\n<image>\nDescribe these images.<|end|>\n<|assistant|>\n" # noqa: E501
|
||||
)
|
||||
|
||||
DTYPE = "half"
|
||||
MAX_TOKENS = 128
|
||||
NUM_LOGPROBS = 10
|
||||
|
||||
|
||||
def vllm_to_hf_output(
|
||||
vllm_output: tuple[list[int], str, SampleLogprobs | None], model: str
|
||||
):
|
||||
"""Sanitize vllm output to be comparable with hf output."""
|
||||
_, output_str, out_logprobs = vllm_output
|
||||
|
||||
output_str_without_image = re.sub(r"(<image>)+", "", output_str)
|
||||
if output_str_without_image and output_str_without_image[0] == " ":
|
||||
output_str_without_image = output_str_without_image[1:]
|
||||
|
||||
hf_output_str = output_str_without_image + "<|end|><|endoftext|>"
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
|
||||
hf_output_ids = tokenizer.encode(output_str_without_image)
|
||||
if hf_output_ids and hf_output_ids[0] == tokenizer.bos_token_id:
|
||||
hf_output_ids = hf_output_ids[1:]
|
||||
|
||||
return hf_output_ids, hf_output_str, out_logprobs
|
||||
|
||||
|
||||
def _build_single_image_inputs(
|
||||
image_assets,
|
||||
) -> list[tuple[list[str], PromptImageInput]]:
|
||||
"""Build single-image inputs for all size_factors at once."""
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
all_inputs: list[tuple[list[str], PromptImageInput]] = []
|
||||
for size_factors in [[1.0], [0.25, 0.5, 1.0]]:
|
||||
for image, prompt in zip(images, HF_IMAGE_PROMPTS):
|
||||
all_inputs.append(
|
||||
(
|
||||
[prompt for _ in size_factors],
|
||||
[rescale_image_size(image, f) for f in size_factors],
|
||||
)
|
||||
)
|
||||
return all_inputs
|
||||
|
||||
|
||||
def _build_multi_image_inputs(
|
||||
image_assets,
|
||||
) -> list[tuple[list[str], PromptImageInput]]:
|
||||
"""Build multi-image inputs for all size_factors at once."""
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
all_inputs: list[tuple[list[str], PromptImageInput]] = []
|
||||
for size_factors in [[0.5], [0.15, 0.30]]:
|
||||
all_inputs.append(
|
||||
(
|
||||
[HF_MULTIIMAGE_IMAGE_PROMPT for _ in size_factors],
|
||||
[
|
||||
[rescale_image_size(image, factor) for image in images]
|
||||
for factor in size_factors
|
||||
],
|
||||
)
|
||||
)
|
||||
return all_inputs
|
||||
|
||||
|
||||
def _run_and_compare(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
all_inputs: Sequence[tuple[list[str], PromptImageInput]],
|
||||
model: str,
|
||||
max_model_len: int,
|
||||
max_num_seqs: int,
|
||||
mm_limit: int,
|
||||
gpu_memory_utilization: float,
|
||||
):
|
||||
"""Load each runner once, run all inputs, then compare."""
|
||||
# NOTE: run vLLM first, then HF. vLLM needs a fresh process without
|
||||
# cuda initialization; running HF first would break the multiprocessing
|
||||
# backend with fork method.
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="generate",
|
||||
max_model_len=max_model_len,
|
||||
max_num_seqs=max_num_seqs,
|
||||
gpu_memory_utilization=gpu_memory_utilization,
|
||||
dtype=DTYPE,
|
||||
limit_mm_per_prompt={"image": mm_limit},
|
||||
tensor_parallel_size=2,
|
||||
trust_remote_code=True,
|
||||
enforce_eager=True,
|
||||
) as vllm_model:
|
||||
vllm_outputs_per_case = [
|
||||
vllm_model.generate_greedy_logprobs(
|
||||
prompts,
|
||||
MAX_TOKENS,
|
||||
num_logprobs=NUM_LOGPROBS,
|
||||
images=images,
|
||||
)
|
||||
for prompts, images in all_inputs
|
||||
]
|
||||
|
||||
hf_model_kwargs = {"_attn_implementation": "sdpa", "device_map": "auto"}
|
||||
with hf_runner(
|
||||
model,
|
||||
dtype=DTYPE,
|
||||
model_kwargs=hf_model_kwargs,
|
||||
auto_cls=AutoModelForCausalLM,
|
||||
trust_remote_code=True,
|
||||
) as hf_model:
|
||||
hf_outputs_per_case = [
|
||||
hf_model.generate_greedy_logprobs_limit(
|
||||
prompts,
|
||||
MAX_TOKENS,
|
||||
num_logprobs=NUM_LOGPROBS,
|
||||
images=images,
|
||||
)
|
||||
for prompts, images in all_inputs
|
||||
]
|
||||
|
||||
for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize("model", [MODEL_ID])
|
||||
def test_models(hf_runner, vllm_runner, image_assets, model) -> None:
|
||||
all_inputs = _build_single_image_inputs(image_assets)
|
||||
_run_and_compare(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
all_inputs,
|
||||
model,
|
||||
max_model_len=8192,
|
||||
max_num_seqs=2,
|
||||
mm_limit=1,
|
||||
gpu_memory_utilization=0.80,
|
||||
)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize("model", [MODEL_ID])
|
||||
def test_multi_images_models(hf_runner, vllm_runner, image_assets, model) -> None:
|
||||
all_inputs = _build_multi_image_inputs(image_assets)
|
||||
_run_and_compare(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
all_inputs,
|
||||
model,
|
||||
max_model_len=8192,
|
||||
max_num_seqs=2,
|
||||
mm_limit=2,
|
||||
gpu_memory_utilization=0.80,
|
||||
)
|
||||
@@ -0,0 +1,251 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
from dataclasses import asdict
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import pytest
|
||||
from mistral_common.multimodal import download_image
|
||||
from mistral_common.protocol.instruct.chunk import ImageURLChunk
|
||||
from mistral_common.protocol.instruct.request import ChatCompletionRequest
|
||||
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
|
||||
from mistral_common.tokens.tokenizers.multimodal import image_from_chunk
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from vllm import SamplingParams, TextPrompt, TokensPrompt
|
||||
from vllm.inputs import MultiModalDataBuiltins
|
||||
from vllm.logprobs import Logprob, SampleLogprobs
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....utils import VLLM_PATH, large_gpu_test
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from _typeshed import StrPath
|
||||
|
||||
PIXTRAL_ID = "mistralai/Pixtral-12B-2409"
|
||||
MISTRAL_SMALL_3_1_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
|
||||
MINISTRAL_3B_ID = "mistralai/Ministral-3-3B-Instruct-2512"
|
||||
|
||||
MODELS = [PIXTRAL_ID, MISTRAL_SMALL_3_1_ID]
|
||||
|
||||
IMG_URLS = [
|
||||
"237-400x300.jpg", # "https://huggingface.co/datasets/Isotr0py/mistral-test-images/resolve/main/237-400x300.jpg",
|
||||
"231-200x300.jpg", # "https://huggingface.co/datasets/Isotr0py/mistral-test-images/resolve/main/237-400x300.jpg",
|
||||
"27-500x500.jpg", # "https://huggingface.co/datasets/Isotr0py/mistral-test-images/resolve/main/237-400x300.jpg",
|
||||
"17-150x600.jpg", # "https://huggingface.co/datasets/Isotr0py/mistral-test-images/resolve/main/237-400x300.jpg",
|
||||
]
|
||||
PROMPT = "Describe each image in one short sentence."
|
||||
|
||||
|
||||
def _create_msg_format(urls: list[str]) -> list[dict[str, Any]]:
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": PROMPT,
|
||||
}
|
||||
]
|
||||
+ [{"type": "image_url", "image_url": {"url": url}} for url in urls],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _create_msg_format_hf(urls: list[str]) -> list[dict[str, Any]]:
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"content": PROMPT,
|
||||
},
|
||||
*({"type": "image", "image": download_image(url)} for url in urls),
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _create_engine_inputs(urls: list[str]) -> TokensPrompt:
|
||||
msg = _create_msg_format(urls)
|
||||
|
||||
tokenizer = MistralTokenizer.from_model("pixtral")
|
||||
|
||||
request = ChatCompletionRequest(messages=msg) # type: ignore[type-var]
|
||||
tokenized = tokenizer.encode_chat_completion(request)
|
||||
|
||||
engine_inputs = TokensPrompt(prompt_token_ids=tokenized.tokens)
|
||||
|
||||
images = []
|
||||
for chunk in request.messages[0].content:
|
||||
if isinstance(chunk, ImageURLChunk):
|
||||
images.append(image_from_chunk(chunk))
|
||||
|
||||
mm_data = MultiModalDataBuiltins(image=images)
|
||||
engine_inputs["multi_modal_data"] = mm_data
|
||||
|
||||
return engine_inputs
|
||||
|
||||
|
||||
def _create_engine_inputs_hf(urls: list[str]) -> TextPrompt:
|
||||
msg = _create_msg_format_hf(urls)
|
||||
|
||||
tokenizer = AutoProcessor.from_pretrained("mistral-community/pixtral-12b")
|
||||
prompt = tokenizer.apply_chat_template(msg)
|
||||
|
||||
images = []
|
||||
for chunk in msg[0]["content"]:
|
||||
if chunk["type"] == "image":
|
||||
images.append(chunk["image"])
|
||||
|
||||
mm_data = MultiModalDataBuiltins(image=images)
|
||||
engine_inputs = TextPrompt(prompt=prompt, multi_modal_data=mm_data)
|
||||
|
||||
return engine_inputs
|
||||
|
||||
|
||||
SAMPLING_PARAMS = SamplingParams(max_tokens=512, temperature=0.0, logprobs=5)
|
||||
LIMIT_MM_PER_PROMPT = dict(image=4)
|
||||
|
||||
MAX_MODEL_LEN = [8192, 65536]
|
||||
|
||||
FIXTURES_PATH = VLLM_PATH / "tests/models/fixtures"
|
||||
assert FIXTURES_PATH.exists()
|
||||
|
||||
FIXTURE_LOGPROBS_CHAT = {
|
||||
PIXTRAL_ID: FIXTURES_PATH / "pixtral_chat.json",
|
||||
MISTRAL_SMALL_3_1_ID: FIXTURES_PATH / "mistral_small_3_chat.json",
|
||||
MINISTRAL_3B_ID: FIXTURES_PATH / "ministral_3b_chat.json",
|
||||
}
|
||||
|
||||
OutputsLogprobs = list[tuple[list[int], str, SampleLogprobs | None]]
|
||||
|
||||
|
||||
# For the test author to store golden output in JSON
|
||||
def _dump_outputs_w_logprobs(
|
||||
outputs: OutputsLogprobs,
|
||||
filename: "StrPath",
|
||||
) -> None:
|
||||
json_data = [
|
||||
(
|
||||
tokens,
|
||||
text,
|
||||
[
|
||||
{k: asdict(v) for k, v in token_logprobs.items()}
|
||||
for token_logprobs in (logprobs or [])
|
||||
],
|
||||
)
|
||||
for tokens, text, logprobs in outputs
|
||||
]
|
||||
|
||||
with open(filename, "w") as f:
|
||||
json.dump(json_data, f)
|
||||
|
||||
|
||||
def load_outputs_w_logprobs(filename: "StrPath") -> OutputsLogprobs:
|
||||
with open(filename, "rb") as f:
|
||||
json_data = json.load(f)
|
||||
|
||||
return [
|
||||
(
|
||||
tokens,
|
||||
text,
|
||||
[
|
||||
{int(k): Logprob(**v) for k, v in token_logprobs.items()}
|
||||
for token_logprobs in logprobs
|
||||
],
|
||||
)
|
||||
for tokens, text, logprobs in json_data
|
||||
]
|
||||
|
||||
|
||||
@large_gpu_test(min_gb=80)
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("max_model_len", MAX_MODEL_LEN)
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
def test_chat(
|
||||
vllm_runner, max_model_len: int, model: str, dtype: str, local_asset_server
|
||||
) -> None:
|
||||
if (
|
||||
model == MISTRAL_SMALL_3_1_ID
|
||||
and max_model_len == 65536
|
||||
and current_platform.is_rocm()
|
||||
):
|
||||
pytest.skip(
|
||||
"OOM on ROCm: 24B model with 65536 context length exceeds GPU memory"
|
||||
)
|
||||
|
||||
EXPECTED_CHAT_LOGPROBS = load_outputs_w_logprobs(FIXTURE_LOGPROBS_CHAT[model])
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
tokenizer_mode="mistral",
|
||||
load_format="mistral",
|
||||
config_format="mistral",
|
||||
max_model_len=max_model_len,
|
||||
limit_mm_per_prompt=LIMIT_MM_PER_PROMPT,
|
||||
) as vllm_model:
|
||||
outputs = []
|
||||
|
||||
urls_all = [local_asset_server.url_for(u) for u in IMG_URLS]
|
||||
msgs = [
|
||||
_create_msg_format(urls_all[:1]),
|
||||
_create_msg_format(urls_all[:2]),
|
||||
_create_msg_format(urls_all),
|
||||
]
|
||||
for msg in msgs:
|
||||
output = vllm_model.llm.chat(msg, sampling_params=SAMPLING_PARAMS)
|
||||
|
||||
outputs.extend(output)
|
||||
|
||||
logprobs = vllm_runner._final_steps_generate_w_logprobs(outputs)
|
||||
# Remove last `None` prompt_logprobs to compare with fixture
|
||||
for i in range(len(logprobs)):
|
||||
assert logprobs[i][-1] is None
|
||||
logprobs[i] = logprobs[i][:-1]
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=EXPECTED_CHAT_LOGPROBS,
|
||||
outputs_1_lst=logprobs,
|
||||
name_0="h100_ref",
|
||||
name_1="output",
|
||||
)
|
||||
|
||||
|
||||
@large_gpu_test(min_gb=16)
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
def test_chat_consolidated(vllm_runner, dtype: str, local_asset_server) -> None:
|
||||
EXPECTED_CHAT_LOGPROBS = load_outputs_w_logprobs(
|
||||
FIXTURE_LOGPROBS_CHAT[MINISTRAL_3B_ID]
|
||||
)
|
||||
with vllm_runner(
|
||||
MINISTRAL_3B_ID,
|
||||
dtype=dtype,
|
||||
tokenizer_mode="mistral",
|
||||
load_format="mistral",
|
||||
config_format="mistral",
|
||||
max_model_len=8192,
|
||||
limit_mm_per_prompt=LIMIT_MM_PER_PROMPT,
|
||||
) as vllm_model:
|
||||
outputs = []
|
||||
urls_all = [local_asset_server.url_for(u) for u in IMG_URLS]
|
||||
msgs = [
|
||||
_create_msg_format(urls_all[:1]),
|
||||
_create_msg_format(urls_all[:2]),
|
||||
_create_msg_format(urls_all),
|
||||
]
|
||||
for msg in msgs:
|
||||
output = vllm_model.llm.chat(msg, sampling_params=SAMPLING_PARAMS)
|
||||
outputs.extend(output)
|
||||
|
||||
logprobs = vllm_runner._final_steps_generate_w_logprobs(outputs)
|
||||
for i in range(len(logprobs)):
|
||||
assert logprobs[i][-1] is None
|
||||
logprobs[i] = logprobs[i][:-1]
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=EXPECTED_CHAT_LOGPROBS,
|
||||
outputs_1_lst=logprobs,
|
||||
name_0="h100_ref",
|
||||
name_1="output",
|
||||
)
|
||||
@@ -0,0 +1,257 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.multimodal.video import sample_frames_from_video
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....conftest import VIDEO_ASSETS
|
||||
|
||||
models = ["Qwen/Qwen2.5-VL-3B-Instruct"]
|
||||
target_dtype = "bfloat16"
|
||||
|
||||
VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>"
|
||||
IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
|
||||
|
||||
def qwen2_5_vl_chat_template(*query):
|
||||
return f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{''.join(query)}<|im_end|><|im_start|>assistant\n" # noqa: E501
|
||||
|
||||
|
||||
VIDEO_PROMPTS = VIDEO_ASSETS.prompts(
|
||||
{
|
||||
"baby_reading": qwen2_5_vl_chat_template(
|
||||
VIDEO_PLACEHOLDER,
|
||||
"Describe this video with a short sentence ",
|
||||
"(no more than 20 words)",
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
WINDOW_ATTN_IMAGE_PROMPT = qwen2_5_vl_chat_template(
|
||||
IMAGE_PLACEHOLDER,
|
||||
"Describe the image.",
|
||||
)
|
||||
IMAGE_ONLY_LIMIT_MM_PER_PROMPT = {"image": 1, "video": 0}
|
||||
|
||||
|
||||
def _window_attention_regression_image():
|
||||
# image from regression issue: https://github.com/vllm-project/vllm/issues/15122
|
||||
image = ImageAsset("hato").pil_image
|
||||
return image.resize((image.width // 2, image.height // 2))
|
||||
|
||||
|
||||
def _encoder_cudagraph_config(*, max_vision_items: int) -> dict:
|
||||
return {
|
||||
"cudagraph_mm_encoder": True,
|
||||
"encoder_cudagraph_max_vision_items_per_batch": max_vision_items,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize(
|
||||
"video_pruning_rate", [0.0] if current_platform.is_cpu() else [0.0, 0.75]
|
||||
)
|
||||
@pytest.mark.parametrize("num_frames", [16])
|
||||
@pytest.mark.parametrize("dtype", [target_dtype])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize(
|
||||
"use_bytecode_hook", [True] if current_platform.is_cpu() else [True, False]
|
||||
)
|
||||
def test_qwen2_5_vl_evs_functionality(
|
||||
vllm_runner,
|
||||
video_assets,
|
||||
model,
|
||||
video_pruning_rate: float,
|
||||
num_frames: int,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
use_bytecode_hook: bool,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
"""Test EVS (Efficient Video Sampling) functionality with different
|
||||
pruning rates.
|
||||
"""
|
||||
# Set the environment variable for this test
|
||||
monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
|
||||
|
||||
# Sample frames from video assets
|
||||
sampled_vids = [
|
||||
sample_frames_from_video(asset.np_ndarrays, num_frames)
|
||||
for asset in video_assets
|
||||
]
|
||||
|
||||
prompts = [VIDEO_PROMPTS[0]]
|
||||
videos = [sampled_vids[0]]
|
||||
|
||||
# Initialize model with EVS configuration
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="generate",
|
||||
max_model_len=4000,
|
||||
dtype=dtype,
|
||||
limit_mm_per_prompt={"video": 1},
|
||||
video_pruning_rate=video_pruning_rate,
|
||||
) as vllm_model:
|
||||
# Generate output - this should not crash
|
||||
outputs = vllm_model.generate_greedy(prompts, max_tokens, videos=videos)
|
||||
|
||||
# Basic validation that we got a response
|
||||
assert len(outputs) == 1
|
||||
output_ids, output_text = outputs[0]
|
||||
|
||||
# Ensure we got some output
|
||||
assert len(output_ids) > 0
|
||||
assert len(output_text) > 0
|
||||
|
||||
# Ensure the output is a string
|
||||
assert isinstance(output_text, str)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize(
|
||||
"video_pruning_rate", [0.0] if current_platform.is_cpu() else [0.0, 0.75]
|
||||
)
|
||||
@pytest.mark.parametrize("num_frames", [16])
|
||||
@pytest.mark.parametrize("dtype", [target_dtype])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize(
|
||||
"use_bytecode_hook", [True] if current_platform.is_cpu() else [True, False]
|
||||
)
|
||||
def test_qwen2_5_vl_evs_batched_videos(
|
||||
vllm_runner,
|
||||
video_assets,
|
||||
model,
|
||||
video_pruning_rate: float,
|
||||
num_frames: int,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
use_bytecode_hook: bool,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
"""Test EVS functionality with batched videos.
|
||||
|
||||
This test validates that:
|
||||
1. The model handles batched video inputs correctly with EVS
|
||||
2. Both pruning configurations work with multiple videos
|
||||
3. The model doesn't crash when processing multiple videos simultaneously
|
||||
"""
|
||||
# Set the environment variable for this test
|
||||
monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
|
||||
# Sample frames from video assets
|
||||
sampled_vids = [
|
||||
sample_frames_from_video(asset.np_ndarrays, num_frames)
|
||||
for asset in video_assets
|
||||
]
|
||||
|
||||
# Test batched videos
|
||||
prompts = [VIDEO_PROMPTS[0], VIDEO_PROMPTS[0]]
|
||||
videos = [sampled_vids[0], sampled_vids[0]] # Use same video twice for testing
|
||||
|
||||
# Initialize model with EVS configuration
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="generate",
|
||||
max_model_len=4000,
|
||||
max_num_seqs=2,
|
||||
dtype=dtype,
|
||||
limit_mm_per_prompt={"video": 2},
|
||||
tensor_parallel_size=1,
|
||||
video_pruning_rate=video_pruning_rate,
|
||||
) as vllm_model:
|
||||
# Generate output - this should not crash
|
||||
outputs = vllm_model.generate_greedy(prompts, max_tokens, videos=videos)
|
||||
|
||||
# Basic validation that we got responses for both videos
|
||||
assert len(outputs) == 2
|
||||
|
||||
for output_ids, output_text in outputs:
|
||||
# Ensure we got some output for each video
|
||||
assert len(output_ids) > 0
|
||||
assert len(output_text) > 0
|
||||
|
||||
# Ensure the output is a string
|
||||
assert isinstance(output_text, str)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize("dtype", [target_dtype])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize(
|
||||
"use_bytecode_hook", [True] if current_platform.is_cpu() else [True, False]
|
||||
)
|
||||
def test_qwen2_5_vl_window_attention_image(
|
||||
vllm_runner,
|
||||
model,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
use_bytecode_hook: bool,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
"""Regression test for Qwen2.5 window-attention image path."""
|
||||
monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
|
||||
|
||||
prompt = [WINDOW_ATTN_IMAGE_PROMPT]
|
||||
images = [[_window_attention_regression_image()]]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="generate",
|
||||
max_model_len=4096,
|
||||
dtype=dtype,
|
||||
limit_mm_per_prompt=IMAGE_ONLY_LIMIT_MM_PER_PROMPT,
|
||||
compilation_config=_encoder_cudagraph_config(max_vision_items=1),
|
||||
) as vllm_model:
|
||||
outputs = vllm_model.generate_greedy(prompt, max_tokens, images=images)
|
||||
|
||||
assert len(outputs) == 1
|
||||
output_ids, output_text = outputs[0]
|
||||
assert len(output_ids) > 0
|
||||
assert len(output_text) > 0
|
||||
assert isinstance(output_text, str)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize("dtype", [target_dtype])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize(
|
||||
"use_bytecode_hook", [True] if current_platform.is_cpu() else [True, False]
|
||||
)
|
||||
def test_qwen2_5_vl_window_attention_image_batch(
|
||||
vllm_runner,
|
||||
model,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
use_bytecode_hook: bool,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
"""Regression test window-attention with a small image batch."""
|
||||
monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
|
||||
|
||||
image = _window_attention_regression_image()
|
||||
prompts = [WINDOW_ATTN_IMAGE_PROMPT, WINDOW_ATTN_IMAGE_PROMPT]
|
||||
images = [[image], [image]]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="generate",
|
||||
max_model_len=4096,
|
||||
max_num_seqs=2,
|
||||
dtype=dtype,
|
||||
limit_mm_per_prompt=IMAGE_ONLY_LIMIT_MM_PER_PROMPT,
|
||||
compilation_config=_encoder_cudagraph_config(max_vision_items=2),
|
||||
) as vllm_model:
|
||||
outputs = vllm_model.generate_greedy(prompts, max_tokens, images=images)
|
||||
|
||||
assert len(outputs) == 2
|
||||
for output_ids, output_text in outputs:
|
||||
assert len(output_ids) > 0
|
||||
assert len(output_text) > 0
|
||||
assert isinstance(output_text, str)
|
||||
@@ -0,0 +1,472 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any, TypedDict
|
||||
|
||||
import numpy.typing as npt
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm.multimodal.image import rescale_image_size
|
||||
from vllm.multimodal.video import rescale_video_size, sample_frames_from_video
|
||||
|
||||
from ....conftest import (
|
||||
IMAGE_ASSETS,
|
||||
VIDEO_ASSETS,
|
||||
PromptImageInput,
|
||||
PromptVideoInput,
|
||||
VllmRunner,
|
||||
)
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", autouse=True)
|
||||
def enable_pickle(monkeypatch):
|
||||
"""`LLM.apply_model` requires pickling a function."""
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
|
||||
|
||||
models = ["Qwen/Qwen2-VL-2B-Instruct"]
|
||||
target_dtype = "half"
|
||||
|
||||
IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>"
|
||||
MODEL_HIDDEN_SIZE = 1536
|
||||
|
||||
|
||||
def qwen2_vl_chat_template(*query):
|
||||
return f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{''.join(query)}<|im_end|><|im_start|>assistant\n" # noqa: E501
|
||||
|
||||
|
||||
IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
|
||||
{
|
||||
"stop_sign": qwen2_vl_chat_template(
|
||||
IMAGE_PLACEHOLDER,
|
||||
"What is the biggest text's content in this image?",
|
||||
),
|
||||
"cherry_blossom": qwen2_vl_chat_template(
|
||||
IMAGE_PLACEHOLDER,
|
||||
"What is the season shown in this image? ",
|
||||
"Reply with a short sentence (no more than 20 words)",
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
VIDEO_PROMPTS = VIDEO_ASSETS.prompts(
|
||||
{
|
||||
"baby_reading": qwen2_vl_chat_template(
|
||||
VIDEO_PLACEHOLDER,
|
||||
"Describe this video with a short sentence ",
|
||||
"(no more than 20 words)",
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
MULTIIMAGE_PROMPT = qwen2_vl_chat_template(
|
||||
IMAGE_PLACEHOLDER,
|
||||
IMAGE_PLACEHOLDER,
|
||||
"Describe these two images separately. ",
|
||||
"For each image, reply with a short sentence ",
|
||||
"(no more than 10 words).",
|
||||
)
|
||||
|
||||
|
||||
class Qwen2VLPromptImageEmbeddingInput(TypedDict):
|
||||
image_embeds: torch.Tensor
|
||||
image_grid_thw: torch.Tensor
|
||||
|
||||
|
||||
class Qwen2VLPromptVideoEmbeddingInput(TypedDict):
|
||||
video_embeds: torch.Tensor
|
||||
video_grid_thw: torch.Tensor
|
||||
|
||||
|
||||
def batch_make_image_embeddings(
|
||||
image_batches: list[Image.Image | list[Image.Image]],
|
||||
processor,
|
||||
llm: VllmRunner,
|
||||
) -> list[Qwen2VLPromptImageEmbeddingInput]:
|
||||
"""batched image embeddings for Qwen2-VL
|
||||
|
||||
This will infer all images' embeddings in a single batch,
|
||||
and split the result according to input batches.
|
||||
|
||||
image_batches:
|
||||
- Single-image batches: `list[Image.Image]`
|
||||
- Multiple-image batches: `list[list[Image.Image]]]`
|
||||
|
||||
returns: `list[Qwen2VLPromptImageEmbeddingInput]`
|
||||
"""
|
||||
|
||||
image_batches_: list[Any] = image_batches[:]
|
||||
|
||||
# convert single-image batches to multiple-image batches
|
||||
for idx in range(len(image_batches_)):
|
||||
if not isinstance(image_batches_[idx], list):
|
||||
image_batches_[idx] = [image_batches_[idx]]
|
||||
|
||||
assert isinstance(image_batches_[idx], list)
|
||||
|
||||
# append all images into a list (as a batch)
|
||||
images: list[Image.Image] = []
|
||||
for image_batch in image_batches_:
|
||||
images += image_batch
|
||||
|
||||
# image to pixel values
|
||||
image_processor = processor.image_processor
|
||||
|
||||
preprocess_result = image_processor.preprocess(
|
||||
images=images, return_tensors="pt"
|
||||
).data
|
||||
pixel_values = preprocess_result["pixel_values"]
|
||||
image_grid_thw = preprocess_result["image_grid_thw"]
|
||||
|
||||
# pixel values to embeddings & grid_thws
|
||||
def get_image_embeds(model):
|
||||
with torch.no_grad():
|
||||
visual = model.visual
|
||||
|
||||
pixel_values_on_device = pixel_values.to(visual.device, dtype=visual.dtype)
|
||||
return visual(pixel_values_on_device, grid_thw=image_grid_thw).cpu()
|
||||
|
||||
image_embeds = torch.concat(llm.apply_model(get_image_embeds))
|
||||
|
||||
# split into original batches
|
||||
result: list[Qwen2VLPromptImageEmbeddingInput] = []
|
||||
image_counter = 0
|
||||
embed_counter = 0
|
||||
for image_batch in image_batches_:
|
||||
cur_batch_image_count = len(image_batch)
|
||||
merge_size = image_processor.merge_size
|
||||
cur_batch_embed_len = sum(
|
||||
grid_thw.prod(-1) // merge_size // merge_size
|
||||
for grid_thw in image_grid_thw[
|
||||
image_counter : image_counter + cur_batch_image_count
|
||||
]
|
||||
)
|
||||
|
||||
result.append(
|
||||
{
|
||||
"image_embeds": image_embeds[
|
||||
embed_counter : embed_counter + cur_batch_embed_len
|
||||
],
|
||||
"image_grid_thw": image_grid_thw[
|
||||
image_counter : image_counter + cur_batch_image_count
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
embed_counter += cur_batch_embed_len
|
||||
image_counter += cur_batch_image_count
|
||||
|
||||
# ensure we don't lose any images or embeddings
|
||||
assert embed_counter == image_embeds.size(0)
|
||||
assert image_counter == image_grid_thw.size(0)
|
||||
assert len(image_batches) == len(result)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def batch_make_video_embeddings(
|
||||
video_batches: PromptVideoInput, processor, llm: VllmRunner
|
||||
) -> list[Qwen2VLPromptVideoEmbeddingInput]:
|
||||
"""batched video embeddings for Qwen2-VL
|
||||
|
||||
A NDArray represents a single video's all frames.
|
||||
|
||||
This will infer all videos' embeddings in a single batch,
|
||||
and split the result according to input batches.
|
||||
|
||||
video_batches:
|
||||
- Single-video batches: `list[NDArray]`
|
||||
- Multiple-video batches: `list[list[NDArray]]`
|
||||
"""
|
||||
|
||||
video_batches_: list[Any] = video_batches[:]
|
||||
|
||||
for idx in range(len(video_batches_)):
|
||||
if not isinstance(video_batches_[idx], list):
|
||||
single_video_batch: list[npt.NDArray] = [video_batches_[idx]]
|
||||
video_batches_[idx] = single_video_batch
|
||||
|
||||
assert isinstance(video_batches_[idx], list)
|
||||
|
||||
# append all videos into a list (as a batch)
|
||||
videos: list[npt.NDArray] = []
|
||||
for video_batch in video_batches_:
|
||||
videos += video_batch
|
||||
|
||||
# video to pixel values
|
||||
video_processor = processor.video_processor
|
||||
|
||||
preprocess_result = video_processor.preprocess(
|
||||
videos=videos, return_tensors="pt"
|
||||
).data
|
||||
pixel_values = preprocess_result["pixel_values_videos"]
|
||||
video_grid_thw = preprocess_result["video_grid_thw"]
|
||||
|
||||
# pixel values to embeddings & grid_thws
|
||||
def get_image_embeds(model):
|
||||
with torch.no_grad():
|
||||
visual = model.visual
|
||||
|
||||
pixel_values_on_device = pixel_values.to(visual.device, dtype=visual.dtype)
|
||||
return visual(pixel_values_on_device, grid_thw=video_grid_thw).cpu()
|
||||
|
||||
video_embeds = torch.concat(llm.apply_model(get_image_embeds))
|
||||
|
||||
# split into original batches
|
||||
result: list[Qwen2VLPromptVideoEmbeddingInput] = []
|
||||
video_counter = 0
|
||||
embed_counter = 0
|
||||
for video_batch in video_batches_:
|
||||
cur_batch_video_count = len(video_batch)
|
||||
merge_size = video_processor.merge_size
|
||||
cur_batch_embed_len = sum(
|
||||
grid_thw.prod(-1) // merge_size // merge_size
|
||||
for grid_thw in video_grid_thw[
|
||||
video_counter : video_counter + cur_batch_video_count
|
||||
]
|
||||
)
|
||||
|
||||
result.append(
|
||||
{
|
||||
"video_embeds": video_embeds[
|
||||
embed_counter : embed_counter + cur_batch_embed_len
|
||||
],
|
||||
"video_grid_thw": video_grid_thw[
|
||||
video_counter : video_counter + cur_batch_video_count
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
embed_counter += cur_batch_embed_len
|
||||
video_counter += cur_batch_video_count
|
||||
|
||||
# ensure we don't lose any videos or embeddings
|
||||
assert embed_counter == video_embeds.size(0)
|
||||
assert video_counter == video_grid_thw.size(0)
|
||||
assert len(video_batches) == len(result)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def run_embedding_input_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
inputs: list[tuple[list[str], PromptImageInput, PromptVideoInput]],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
mm_limit: int,
|
||||
tensor_parallel_size: int,
|
||||
distributed_executor_backend: str | None = None,
|
||||
):
|
||||
"""Inference result should be the same between
|
||||
original image/video input and image/video embeddings input.
|
||||
"""
|
||||
from transformers import AutoProcessor
|
||||
|
||||
processor = AutoProcessor.from_pretrained(model)
|
||||
|
||||
# max_model_len should be greater than image_feature_size
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="generate",
|
||||
max_model_len=4000,
|
||||
max_num_seqs=3,
|
||||
dtype=dtype,
|
||||
limit_mm_per_prompt={"image": mm_limit, "video": mm_limit},
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
default_torch_num_threads=1,
|
||||
enable_mm_embeds=True,
|
||||
) as vllm_model:
|
||||
outputs_per_case_for_original_input = [
|
||||
vllm_model.generate_greedy_logprobs(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
images=images or None,
|
||||
videos=videos or None,
|
||||
)
|
||||
for prompts, images, videos in inputs
|
||||
]
|
||||
|
||||
outputs_per_case_for_embeddings_input = [
|
||||
vllm_model.generate_greedy_logprobs(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
images=batch_make_image_embeddings(images, processor, vllm_model)
|
||||
if images
|
||||
else None,
|
||||
videos=batch_make_video_embeddings(videos, processor, vllm_model)
|
||||
if videos
|
||||
else None,
|
||||
)
|
||||
for prompts, images, videos in inputs
|
||||
]
|
||||
|
||||
for outputs_for_original_input, outputs_for_embeddings_input in zip(
|
||||
outputs_per_case_for_original_input, outputs_per_case_for_embeddings_input
|
||||
):
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=outputs_for_original_input,
|
||||
outputs_1_lst=outputs_for_embeddings_input,
|
||||
name_0="original_input",
|
||||
name_1="embeddings_input",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize(
|
||||
"size_factors",
|
||||
[
|
||||
# Single-scale
|
||||
[0.5],
|
||||
# Single-scale, batched
|
||||
[0.5, 0.5],
|
||||
# Multi-scale
|
||||
[0.25, 0.5, 0.5],
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [target_dtype])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize("num_logprobs", [10])
|
||||
def test_qwen2_vl_image_embeddings_input(
|
||||
vllm_runner,
|
||||
image_assets,
|
||||
model,
|
||||
size_factors,
|
||||
dtype,
|
||||
max_tokens,
|
||||
num_logprobs,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
|
||||
inputs_per_case: list[tuple[list[str], PromptImageInput, PromptVideoInput]] = [
|
||||
(
|
||||
[prompt for _ in size_factors],
|
||||
[rescale_image_size(image, factor) for factor in size_factors],
|
||||
[],
|
||||
)
|
||||
for image, prompt in zip(images, IMAGE_PROMPTS)
|
||||
]
|
||||
|
||||
run_embedding_input_test(
|
||||
vllm_runner,
|
||||
inputs_per_case,
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_tokens=max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
mm_limit=1,
|
||||
tensor_parallel_size=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize(
|
||||
"size_factors",
|
||||
[
|
||||
# Single-scale
|
||||
[0.5],
|
||||
# Single-scale, batched
|
||||
[0.5, 0.5],
|
||||
# Multi-scale
|
||||
[0.25, 0.5, 0.5],
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [target_dtype])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize("num_logprobs", [10])
|
||||
def test_qwen2_vl_multiple_image_embeddings_input(
|
||||
vllm_runner,
|
||||
image_assets,
|
||||
model,
|
||||
size_factors,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
|
||||
inputs_per_case: list[tuple[list[str], PromptImageInput, PromptVideoInput]] = [
|
||||
(
|
||||
[MULTIIMAGE_PROMPT for _ in size_factors],
|
||||
[
|
||||
[rescale_image_size(image, factor) for image in images]
|
||||
for factor in size_factors
|
||||
],
|
||||
[],
|
||||
)
|
||||
]
|
||||
|
||||
run_embedding_input_test(
|
||||
vllm_runner,
|
||||
inputs_per_case,
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_tokens=max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
mm_limit=2,
|
||||
tensor_parallel_size=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", models)
|
||||
@pytest.mark.parametrize(
|
||||
"size_factors",
|
||||
[
|
||||
# Single-scale
|
||||
[0.5],
|
||||
# Single-scale, batched
|
||||
[0.5, 0.5],
|
||||
# Multi-scale
|
||||
[0.25, 0.25, 0.5],
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [target_dtype])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize("num_logprobs", [10])
|
||||
def test_qwen2_vl_video_embeddings_input(
|
||||
vllm_runner,
|
||||
video_assets,
|
||||
model,
|
||||
size_factors,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
num_frames = 4
|
||||
sampled_vids = [
|
||||
sample_frames_from_video(asset.np_ndarrays, num_frames)
|
||||
for asset in video_assets
|
||||
]
|
||||
|
||||
inputs_per_case: list[tuple[list[str], PromptImageInput, PromptVideoInput]] = [
|
||||
(
|
||||
[prompt for _ in size_factors],
|
||||
[],
|
||||
[rescale_video_size(video, factor) for factor in size_factors],
|
||||
)
|
||||
for video, prompt in zip(sampled_vids, VIDEO_PROMPTS)
|
||||
]
|
||||
|
||||
run_embedding_input_test(
|
||||
vllm_runner,
|
||||
inputs_per_case,
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_tokens=max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
mm_limit=1,
|
||||
tensor_parallel_size=1,
|
||||
)
|
||||
@@ -0,0 +1,230 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from ....conftest import AUDIO_ASSETS, AudioTestAssets, VllmRunner
|
||||
from ....utils import RemoteOpenAIServer
|
||||
from ...registry import HF_EXAMPLE_MODELS
|
||||
|
||||
MODEL_NAME = "fixie-ai/ultravox-v0_5-llama-3_2-1b"
|
||||
|
||||
AUDIO_PROMPTS = AUDIO_ASSETS.prompts(
|
||||
{
|
||||
"mary_had_lamb": "Transcribe this into English.",
|
||||
"winning_call": "What is happening in this audio clip?",
|
||||
}
|
||||
)
|
||||
|
||||
MULTI_AUDIO_PROMPT = "Describe each of the audios above."
|
||||
|
||||
AudioTuple = tuple[np.ndarray, int]
|
||||
|
||||
VLLM_PLACEHOLDER = "<|audio|>"
|
||||
HF_PLACEHOLDER = "<|audio|>"
|
||||
|
||||
CHUNKED_PREFILL_KWARGS = {
|
||||
"enable_chunked_prefill": True,
|
||||
"max_num_seqs": 2,
|
||||
# Use a very small limit to exercise chunked prefill.
|
||||
"max_num_batched_tokens": 16,
|
||||
}
|
||||
|
||||
|
||||
def params_kwargs_to_cli_args(params_kwargs: dict[str, Any]) -> list[str]:
|
||||
"""Convert kwargs to CLI args."""
|
||||
args = []
|
||||
for key, value in params_kwargs.items():
|
||||
if isinstance(value, bool):
|
||||
if value:
|
||||
args.append(f"--{key.replace('_', '-')}")
|
||||
else:
|
||||
args.append(f"--{key.replace('_', '-')}={value}")
|
||||
return args
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
params=[
|
||||
pytest.param({}, marks=pytest.mark.cpu_model),
|
||||
pytest.param(CHUNKED_PREFILL_KWARGS),
|
||||
]
|
||||
)
|
||||
def server(request, audio_assets: AudioTestAssets):
|
||||
args = [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
"--enforce-eager",
|
||||
"--limit-mm-per-prompt",
|
||||
json.dumps({"audio": len(audio_assets)}),
|
||||
"--trust-remote-code",
|
||||
] + params_kwargs_to_cli_args(request.param)
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME, args, env_dict={"VLLM_AUDIO_FETCH_TIMEOUT": "30"}
|
||||
) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
def _get_prompt(audio_count, question, placeholder):
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
||||
placeholder = f"{placeholder}\n" * audio_count
|
||||
|
||||
return tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": f"{placeholder}{question}"}],
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
|
||||
|
||||
def run_multi_audio_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
prompts_and_audios: list[tuple[str, list[AudioTuple]]],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
**kwargs,
|
||||
):
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
enforce_eager=True,
|
||||
limit_mm_per_prompt={
|
||||
"audio": max((len(audio) for _, audio in prompts_and_audios))
|
||||
},
|
||||
**kwargs,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
[prompt for prompt, _ in prompts_and_audios],
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
audios=[audios for _, audios in prompts_and_audios],
|
||||
)
|
||||
|
||||
# The HuggingFace model doesn't support multiple audios yet, so
|
||||
# just assert that some tokens were generated.
|
||||
assert all(tokens for tokens, *_ in vllm_outputs)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.parametrize(
|
||||
"vllm_kwargs",
|
||||
[
|
||||
pytest.param({}, marks=pytest.mark.cpu_model),
|
||||
pytest.param(CHUNKED_PREFILL_KWARGS),
|
||||
],
|
||||
)
|
||||
def test_models_with_multiple_audios(
|
||||
vllm_runner,
|
||||
audio_assets: AudioTestAssets,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
vllm_kwargs: dict,
|
||||
) -> None:
|
||||
vllm_prompt = _get_prompt(len(audio_assets), MULTI_AUDIO_PROMPT, VLLM_PLACEHOLDER)
|
||||
run_multi_audio_test(
|
||||
vllm_runner,
|
||||
[(vllm_prompt, [audio.audio_and_sample_rate for audio in audio_assets])],
|
||||
MODEL_NAME,
|
||||
dtype=dtype,
|
||||
max_tokens=max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
**vllm_kwargs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", [32])
|
||||
def test_variable_length_audio_batching(
|
||||
vllm_runner,
|
||||
audio_assets: AudioTestAssets,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
) -> None:
|
||||
"""Test batching of requests with different audio durations.
|
||||
|
||||
This exercises the variable-length tensor handling in
|
||||
MultiModalFlatField._reduce_data() which was buggy before
|
||||
https://github.com/vllm-project/vllm/issues/31658 was fixed.
|
||||
"""
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(MODEL_NAME)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
# Create prompts with single audio each (different durations)
|
||||
prompts_and_audios = []
|
||||
for audio, question in zip(audio_assets, AUDIO_PROMPTS):
|
||||
prompt = _get_prompt(1, question, VLLM_PLACEHOLDER)
|
||||
prompts_and_audios.append((prompt, [audio.audio_and_sample_rate]))
|
||||
|
||||
with vllm_runner(
|
||||
MODEL_NAME,
|
||||
dtype=dtype,
|
||||
enforce_eager=True,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
) as vllm_model:
|
||||
# Generate for all prompts in a single batch
|
||||
# This triggers the variable-length batching code path
|
||||
outputs = vllm_model.generate_greedy(
|
||||
[prompt for prompt, _ in prompts_and_audios],
|
||||
max_tokens,
|
||||
audios=[audios for _, audios in prompts_and_audios],
|
||||
)
|
||||
|
||||
# Verify outputs were generated for each request
|
||||
assert len(outputs) == len(prompts_and_audios)
|
||||
for output in outputs:
|
||||
assert len(output[1]) > 0, "Expected non-empty output"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_online_serving(client, audio_assets: AudioTestAssets):
|
||||
"""Exercises online serving with/without chunked prefill enabled."""
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
*[
|
||||
{"type": "audio_url", "audio_url": {"url": audio.url}}
|
||||
for audio in audio_assets
|
||||
],
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"What's happening in these {len(audio_assets)} audio clips?", # noqa: E501
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
chat_completion = await client.chat.completions.create(
|
||||
model=MODEL_NAME, messages=messages, max_tokens=10
|
||||
)
|
||||
|
||||
assert len(chat_completion.choices) == 1
|
||||
choice = chat_completion.choices[0]
|
||||
assert choice.finish_reason == "length"
|
||||
@@ -0,0 +1,365 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.multimodal.video import sample_frames_from_video
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....conftest import IMAGE_ASSETS, VIDEO_ASSETS
|
||||
from ...utils import dummy_hf_overrides
|
||||
from .vlm_utils.builders import sample_frames_with_video_metadata
|
||||
|
||||
|
||||
@dataclass
|
||||
class VitCudagraphTestConfig:
|
||||
model: str
|
||||
modalities: list[str] = field(default_factory=lambda: ["image", "video"])
|
||||
image_prompt: str | None = None
|
||||
video_prompt: str | None = None
|
||||
dtype: str = "bfloat16"
|
||||
max_model_len: int = 4096
|
||||
max_tokens: int = 64
|
||||
max_num_seqs: int = 2
|
||||
num_video_frames: int = 16
|
||||
needs_video_metadata: bool = False
|
||||
vllm_runner_kwargs: dict = field(default_factory=dict)
|
||||
compilation_config_overrides: dict = field(default_factory=dict)
|
||||
marks: list = field(default_factory=list)
|
||||
skip: bool = False
|
||||
|
||||
|
||||
def params_with_marks(
|
||||
configs: dict[str, VitCudagraphTestConfig],
|
||||
) -> list[pytest.param]:
|
||||
return [
|
||||
pytest.param(model_id, marks=cfg.marks) for model_id, cfg in configs.items()
|
||||
]
|
||||
|
||||
|
||||
def qwen_vl_chat_template(content: str) -> str:
|
||||
return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
|
||||
|
||||
|
||||
def internvl_chat_template(content: str) -> str:
|
||||
return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
|
||||
|
||||
|
||||
def kimi_vl_chat_template(content: str) -> str:
|
||||
return (
|
||||
f"<|im_user|>user<|im_middle|>{content}<|im_end|>"
|
||||
"<|im_assistant|>assistant<|im_middle|>"
|
||||
)
|
||||
|
||||
|
||||
def step3_vl_chat_template(content: str) -> str:
|
||||
return (
|
||||
"<|begin▁of▁sentence|> You are a helpful assistant.<|BOT|>user\n "
|
||||
f"<im_patch>{content} <|EOT|><|BOT|>assistant\n"
|
||||
)
|
||||
|
||||
|
||||
def gemma3_chat_template(content: str) -> str:
|
||||
return f"<bos><start_of_turn>user\n{content}<end_of_turn>\n<start_of_turn>model\n"
|
||||
|
||||
|
||||
MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
|
||||
"gemma3": VitCudagraphTestConfig(
|
||||
model="google/gemma-3-4b-it",
|
||||
modalities=["image"],
|
||||
image_prompt=gemma3_chat_template("<start_of_image>What is in this image?"),
|
||||
compilation_config_overrides={
|
||||
"encoder_cudagraph_token_budgets": [512],
|
||||
},
|
||||
dtype="bfloat16",
|
||||
max_model_len=4096,
|
||||
),
|
||||
"llama4": VitCudagraphTestConfig(
|
||||
model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
|
||||
modalities=["image"],
|
||||
image_prompt=(
|
||||
"<|begin_of_text|><|header_start|>user<|header_end|>\n\n"
|
||||
"<|image|>What is in this image?<|eot|>"
|
||||
"<|header_start|>assistant<|header_end|>\n\n"
|
||||
),
|
||||
max_model_len=4096,
|
||||
max_tokens=32,
|
||||
max_num_seqs=2,
|
||||
vllm_runner_kwargs={
|
||||
"load_format": "dummy",
|
||||
"hf_overrides": partial(
|
||||
dummy_hf_overrides,
|
||||
model_arch="Llama4ForConditionalGeneration",
|
||||
),
|
||||
},
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
"qwen2_vl": VitCudagraphTestConfig(
|
||||
model="Qwen/Qwen2-VL-2B-Instruct",
|
||||
image_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
|
||||
),
|
||||
video_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|video_pad|><|vision_end|>"
|
||||
"Describe this video in one sentence."
|
||||
),
|
||||
needs_video_metadata=False,
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
"qwen2_5_vl": VitCudagraphTestConfig(
|
||||
model="Qwen/Qwen2.5-VL-3B-Instruct",
|
||||
image_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
|
||||
),
|
||||
video_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|video_pad|><|vision_end|>"
|
||||
"Describe this video in one sentence."
|
||||
),
|
||||
needs_video_metadata=False,
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
"kimi_vl": VitCudagraphTestConfig(
|
||||
model="moonshotai/Kimi-VL-A3B-Instruct",
|
||||
modalities=["image"],
|
||||
image_prompt=kimi_vl_chat_template(
|
||||
"<|media_start|>image<|media_content|><|media_pad|><|media_end|>"
|
||||
"What is in this image?"
|
||||
),
|
||||
needs_video_metadata=False,
|
||||
# Single bucket sized to cover the test images' output tokens.
|
||||
# The default auto-inferred range fans out into multiple power-of-2
|
||||
# buckets, each holding a full ViT capture pool.
|
||||
compilation_config_overrides={
|
||||
"encoder_cudagraph_token_budgets": [1024],
|
||||
},
|
||||
# Shrink to 1 text + 1 vision layer with random weights so the
|
||||
# test runs on any CI GPU (incl. L4) and skips the multi-GiB
|
||||
# weight download. The test only validates that encoder CG
|
||||
# capture/replay functions correctly, not output quality.
|
||||
vllm_runner_kwargs={
|
||||
"trust_remote_code": True,
|
||||
"load_format": "dummy",
|
||||
"hf_overrides": partial(
|
||||
dummy_hf_overrides,
|
||||
model_arch="KimiVLForConditionalGeneration",
|
||||
),
|
||||
},
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
"qwen3_vl": VitCudagraphTestConfig(
|
||||
model="Qwen/Qwen3-VL-2B-Instruct",
|
||||
image_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
|
||||
),
|
||||
video_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|video_pad|><|vision_end|>"
|
||||
"Describe this video in one sentence."
|
||||
),
|
||||
needs_video_metadata=True,
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
"qwen3_5": VitCudagraphTestConfig(
|
||||
model="Qwen/Qwen3.5-0.8B",
|
||||
image_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
|
||||
),
|
||||
video_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|video_pad|><|vision_end|>"
|
||||
"Describe this video in one sentence."
|
||||
),
|
||||
needs_video_metadata=True,
|
||||
vllm_runner_kwargs={"enable_chunked_prefill": True},
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
"internvl": VitCudagraphTestConfig(
|
||||
model="OpenGVLab/InternVL3-1B",
|
||||
num_video_frames=8,
|
||||
image_prompt=internvl_chat_template("<image>\nWhat is in this image?"),
|
||||
video_prompt=internvl_chat_template(
|
||||
"<video>\nDescribe this video in one sentence."
|
||||
),
|
||||
needs_video_metadata=False,
|
||||
vllm_runner_kwargs={"trust_remote_code": True},
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
"step3_vl": VitCudagraphTestConfig(
|
||||
model="stepfun-ai/Step3-VL-10B",
|
||||
modalities=["image"],
|
||||
image_prompt=step3_vl_chat_template("What is in this image?"),
|
||||
# Single bucket sized to cover the largest test image's output
|
||||
# tokens (1152 > 1141 for cherry_blossom). The default auto-
|
||||
# inferred range fans out into multiple power-of-2 buckets, each
|
||||
# holding a full ViT capture pool.
|
||||
compilation_config_overrides={
|
||||
"encoder_cudagraph_token_budgets": [1152],
|
||||
},
|
||||
# Shrink to 1 text + 1 vision layer with random weights so the
|
||||
# test runs on any CI GPU (incl. L4) and skips the 20 GiB weight
|
||||
# download. The test only validates that encoder CG capture/
|
||||
# replay functions correctly, not output quality.
|
||||
vllm_runner_kwargs={
|
||||
"load_format": "dummy",
|
||||
"hf_overrides": partial(
|
||||
dummy_hf_overrides,
|
||||
model_arch="StepVLForConditionalGeneration",
|
||||
),
|
||||
},
|
||||
),
|
||||
"glm4_1v": VitCudagraphTestConfig(
|
||||
model="zai-org/GLM-4.1V-9B-Thinking",
|
||||
image_prompt=(
|
||||
"[gMASK]<sop><|system|>\nYou are a helpful assistant.<|user|>\n"
|
||||
"<|begin_of_image|><|image|><|end_of_image|>"
|
||||
"What is in this image?<|assistant|>assistant\n"
|
||||
),
|
||||
video_prompt=(
|
||||
"[gMASK]<sop><|system|>\nYou are a helpful assistant.<|user|>\n"
|
||||
"<|begin_of_video|><|video|><|end_of_video|>"
|
||||
"Describe this video in one sentence<|assistant|>assistant\n"
|
||||
),
|
||||
needs_video_metadata=True,
|
||||
marks=[pytest.mark.core_model],
|
||||
vllm_runner_kwargs={
|
||||
"load_format": "dummy",
|
||||
"hf_overrides": partial(
|
||||
dummy_hf_overrides,
|
||||
model_arch="Glm4vForConditionalGeneration",
|
||||
),
|
||||
},
|
||||
),
|
||||
"deepseek_ocr": VitCudagraphTestConfig(
|
||||
model="deepseek-ai/DeepSeek-OCR",
|
||||
modalities=["image"],
|
||||
image_prompt="<image>\nWhat is in this image?",
|
||||
marks=[pytest.mark.core_model],
|
||||
compilation_config_overrides={
|
||||
"encoder_cudagraph_token_budgets": [272],
|
||||
"mode": 0,
|
||||
"cudagraph_mode": 2,
|
||||
},
|
||||
vllm_runner_kwargs={
|
||||
"load_format": "dummy",
|
||||
"hf_overrides": partial(
|
||||
dummy_hf_overrides,
|
||||
model_arch="DeepseekOCRForCausalLM",
|
||||
),
|
||||
},
|
||||
skip=True, # TODO: Re-enable this once OOM issues are resolved on CI.
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def get_compilation_config(config: VitCudagraphTestConfig):
|
||||
return {
|
||||
"cudagraph_mm_encoder": True,
|
||||
"encoder_cudagraph_max_vision_items_per_batch": 1,
|
||||
"encoder_cudagraph_max_frames_per_batch": 16,
|
||||
**config.compilation_config_overrides,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS))
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
|
||||
def test_vit_cudagraph_image(model_id, vllm_runner, image_assets):
|
||||
config = MODEL_CONFIGS[model_id]
|
||||
|
||||
if config.skip:
|
||||
pytest.skip(f"{model_id} is marked to be skipped.")
|
||||
|
||||
if "image" not in config.modalities:
|
||||
pytest.skip(f"{model_id} does not support the image modality.")
|
||||
|
||||
image_prompts = IMAGE_ASSETS.prompts(
|
||||
{
|
||||
"stop_sign": config.image_prompt, # type: ignore[typeddict-item]
|
||||
"cherry_blossom": config.image_prompt, # type: ignore[typeddict-item]
|
||||
}
|
||||
)
|
||||
images = [[asset.pil_image] for asset in image_assets]
|
||||
|
||||
with vllm_runner(
|
||||
config.model,
|
||||
dtype=config.dtype,
|
||||
max_model_len=config.max_model_len,
|
||||
max_num_seqs=config.max_num_seqs,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
compilation_config=get_compilation_config(config),
|
||||
**config.vllm_runner_kwargs,
|
||||
) as vllm_model:
|
||||
outputs = vllm_model.generate_greedy(
|
||||
image_prompts, config.max_tokens, images=images
|
||||
)
|
||||
|
||||
# Basic validation that we got a response
|
||||
assert len(outputs) == 2
|
||||
output_ids, output_text = outputs[0]
|
||||
|
||||
# Ensure we got some output
|
||||
assert len(output_ids) > 0
|
||||
assert len(output_text) > 0
|
||||
|
||||
# Ensure the output is a string
|
||||
assert isinstance(output_text, str)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS))
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
|
||||
def test_vit_cudagraph_video(model_id, vllm_runner, video_assets):
|
||||
config = MODEL_CONFIGS[model_id]
|
||||
|
||||
if config.skip:
|
||||
pytest.skip(f"{model_id} is marked to be skipped.")
|
||||
|
||||
if "video" not in config.modalities:
|
||||
pytest.skip(f"{model_id} does not support the video modality")
|
||||
|
||||
video_prompts = VIDEO_ASSETS.prompts(
|
||||
{
|
||||
"baby_reading": config.video_prompt, # type: ignore[typeddict-item]
|
||||
}
|
||||
)
|
||||
if config.needs_video_metadata:
|
||||
sampled_vids = [
|
||||
sample_frames_with_video_metadata(
|
||||
(asset.np_ndarrays, asset.metadata), config.num_video_frames
|
||||
)
|
||||
for asset in video_assets
|
||||
]
|
||||
else:
|
||||
sampled_vids = [
|
||||
sample_frames_from_video(asset.np_ndarrays, config.num_video_frames)
|
||||
for asset in video_assets
|
||||
]
|
||||
videos = [sampled_vids[0]]
|
||||
|
||||
with vllm_runner(
|
||||
config.model,
|
||||
dtype=config.dtype,
|
||||
max_model_len=config.max_model_len,
|
||||
max_num_seqs=config.max_num_seqs,
|
||||
limit_mm_per_prompt={"video": 1},
|
||||
compilation_config=get_compilation_config(config),
|
||||
**config.vllm_runner_kwargs,
|
||||
) as vllm_model:
|
||||
outputs = vllm_model.generate_greedy(
|
||||
video_prompts, config.max_tokens, videos=videos
|
||||
)
|
||||
|
||||
# Basic validation that we got a response
|
||||
assert len(outputs) == 1
|
||||
output_ids, output_text = outputs[0]
|
||||
|
||||
# Ensure we got some output
|
||||
assert len(output_ids) > 0
|
||||
assert len(output_text) > 0
|
||||
|
||||
# Ensure the output is a string
|
||||
assert isinstance(output_text, str)
|
||||
@@ -0,0 +1,212 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from mistral_common.protocol.instruct.chunk import AudioChunk, TextChunk
|
||||
from mistral_common.protocol.instruct.messages import UserMessage
|
||||
from mistral_common.tokens.tokenizers.audio import Audio
|
||||
from transformers import VoxtralForConditionalGeneration
|
||||
|
||||
from vllm.tokenizers.mistral import MistralTokenizer
|
||||
|
||||
from ....conftest import AudioTestAssets
|
||||
from ....utils import RemoteOpenAIServer
|
||||
from ...utils import check_logprobs_close
|
||||
from .test_ultravox import MULTI_AUDIO_PROMPT, run_multi_audio_test
|
||||
from .vlm_utils import model_utils
|
||||
|
||||
MODEL_NAME = "mistralai/Voxtral-Mini-3B-2507"
|
||||
MISTRAL_FORMAT_ARGS = [
|
||||
"--tokenizer_mode",
|
||||
"mistral",
|
||||
"--config_format",
|
||||
"mistral",
|
||||
"--load_format",
|
||||
"mistral",
|
||||
]
|
||||
|
||||
|
||||
def _get_prompt(audio_assets: AudioTestAssets, question: str) -> list[int]:
|
||||
"""Build a token-ID prompt via mistral_common for vLLM offline inference."""
|
||||
tokenizer = MistralTokenizer.from_pretrained(MODEL_NAME)
|
||||
|
||||
audios = [
|
||||
Audio.from_file(str(asset.get_local_path()), strict=False)
|
||||
for asset in audio_assets
|
||||
]
|
||||
audio_chunks = [AudioChunk.from_audio(audio) for audio in audios]
|
||||
|
||||
messages = [
|
||||
UserMessage(content=[*audio_chunks, TextChunk(text=question)]).to_openai()
|
||||
]
|
||||
return tokenizer.apply_chat_template(messages=messages)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", [128])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
def test_models_with_multiple_audios(
|
||||
vllm_runner,
|
||||
audio_assets: AudioTestAssets,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
) -> None:
|
||||
vllm_prompt = _get_prompt(audio_assets, MULTI_AUDIO_PROMPT)
|
||||
run_multi_audio_test(
|
||||
vllm_runner,
|
||||
[(vllm_prompt, [a.audio_and_sample_rate for a in audio_assets])], # type: ignore[list-item]
|
||||
MODEL_NAME,
|
||||
dtype=dtype,
|
||||
max_tokens=max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
tokenizer_mode="mistral",
|
||||
)
|
||||
|
||||
|
||||
def test_online_serving(vllm_runner, audio_assets: AudioTestAssets):
|
||||
"""Two-layer accuracy and serving validation using Mistral format.
|
||||
|
||||
1. Offline vLLM greedy output (runs first to avoid CUDA fork issues
|
||||
with multiprocessing - see vlm_utils/core.py).
|
||||
2. Online OpenAI-compatible API output must match offline — validates
|
||||
that the serving path (chat template, audio encoding, tokenization)
|
||||
does not corrupt anything.
|
||||
|
||||
Steps run sequentially so each releases the GPU before the next starts.
|
||||
"""
|
||||
|
||||
question = f"What's happening in these {len(audio_assets)} audio clips?"
|
||||
max_tokens = 10
|
||||
audio_data = [asset.audio_and_sample_rate for asset in audio_assets]
|
||||
|
||||
vllm_prompt = _get_prompt(audio_assets, question)
|
||||
with vllm_runner(
|
||||
MODEL_NAME,
|
||||
dtype="half",
|
||||
enforce_eager=True,
|
||||
tokenizer_mode="mistral",
|
||||
config_format="mistral",
|
||||
load_format="mistral",
|
||||
limit_mm_per_prompt={"audio": len(audio_assets)},
|
||||
) as vllm_model:
|
||||
offline_outputs = vllm_model.generate_greedy(
|
||||
[vllm_prompt],
|
||||
max_tokens,
|
||||
audios=[audio_data],
|
||||
)
|
||||
|
||||
offline_text = offline_outputs[0][1]
|
||||
assert offline_text, "Offline vLLM inference produced empty output"
|
||||
|
||||
def _asset_to_openai_chunk(asset):
|
||||
audio = Audio.from_file(str(asset.get_local_path()), strict=False)
|
||||
audio.format = "wav"
|
||||
return AudioChunk.from_audio(audio).to_openai()
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
*[_asset_to_openai_chunk(a) for a in audio_assets],
|
||||
{"type": "text", "text": question},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
server_args = [
|
||||
"--enforce-eager",
|
||||
"--limit-mm-per-prompt",
|
||||
json.dumps({"audio": len(audio_assets)}),
|
||||
*MISTRAL_FORMAT_ARGS,
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME,
|
||||
server_args,
|
||||
env_dict={"VLLM_AUDIO_FETCH_TIMEOUT": "30"},
|
||||
) as remote_server:
|
||||
client = remote_server.get_client()
|
||||
completion = client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
messages=messages,
|
||||
max_tokens=max_tokens,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
assert len(completion.choices) == 1
|
||||
choice = completion.choices[0]
|
||||
assert choice.finish_reason == "length"
|
||||
assert choice.message.content == offline_text, (
|
||||
f"Online serving output does not match offline inference.\n"
|
||||
f" Online: {choice.message.content!r}\n"
|
||||
f" Offline: {offline_text!r}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="VoxtralProcessor.apply_chat_template() in transformers v5 "
|
||||
"doesn't resolve chat_template=None to the default template"
|
||||
)
|
||||
def test_hf_reference(hf_runner, vllm_runner, audio_assets: AudioTestAssets):
|
||||
"""Compare vLLM Mistral-format output against HF Transformers reference.
|
||||
|
||||
Instead of requiring an exact text match (which is brittle across
|
||||
attention backends), we compare per-token logprobs using the standard
|
||||
check_logprobs_close helper: when tokens diverge at a position, each
|
||||
runner's chosen token must appear in the other's top-k logprobs.
|
||||
|
||||
Marked xfail(strict=False) so remaining edge-case mismatches
|
||||
don't block CI.
|
||||
"""
|
||||
question = f"What's happening in these {len(audio_assets)} audio clips?"
|
||||
max_tokens = 10
|
||||
num_logprobs = 5
|
||||
audio_data = [asset.audio_and_sample_rate for asset in audio_assets]
|
||||
|
||||
vllm_prompt = _get_prompt(audio_assets, question)
|
||||
with vllm_runner(
|
||||
MODEL_NAME,
|
||||
dtype="half",
|
||||
enforce_eager=True,
|
||||
tokenizer_mode="mistral",
|
||||
config_format="mistral",
|
||||
load_format="mistral",
|
||||
limit_mm_per_prompt={"audio": len(audio_assets)},
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
[vllm_prompt],
|
||||
max_tokens,
|
||||
num_logprobs,
|
||||
audios=[audio_data],
|
||||
)
|
||||
assert vllm_outputs[0][1], "vLLM inference produced empty output"
|
||||
|
||||
with hf_runner(
|
||||
MODEL_NAME,
|
||||
dtype="half",
|
||||
auto_cls=VoxtralForConditionalGeneration,
|
||||
) as hf_model:
|
||||
hf_model = model_utils.voxtral_patch_hf_runner(hf_model)
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
[question],
|
||||
max_tokens,
|
||||
num_logprobs,
|
||||
audios=[audio_data],
|
||||
)
|
||||
assert hf_outputs[0][1], "HF Transformers produced empty output"
|
||||
|
||||
print(
|
||||
f"HF Reference Comparison\n"
|
||||
f" vLLM: {vllm_outputs[0][1]!r}\n"
|
||||
f" HF: {hf_outputs[0][1]!r}"
|
||||
)
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=vllm_outputs,
|
||||
outputs_1_lst=hf_outputs,
|
||||
name_0="vllm",
|
||||
name_1="hf",
|
||||
)
|
||||
@@ -0,0 +1,235 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from mistral_common.protocol.transcription.request import (
|
||||
StreamingMode,
|
||||
TranscriptionRequest,
|
||||
)
|
||||
from mistral_common.tokens.tokenizers.audio import Audio
|
||||
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
|
||||
from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.kv_cache_interface import SlidingWindowSpec
|
||||
|
||||
from ....utils import ROCM_ENGINE_KWARGS
|
||||
|
||||
MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
|
||||
AUDIO_LAYER_NAME = "whisper_encoder.whisper_encoder.layers.0.layers.self_attn.attn"
|
||||
ENGINE_CONFIG = {
|
||||
"model": MODEL_NAME,
|
||||
"max_model_len": 8192,
|
||||
"max_num_seqs": 4,
|
||||
"limit_mm_per_prompt": {"audio": 1},
|
||||
"config_format": "mistral",
|
||||
"load_format": "mistral",
|
||||
"tokenizer_mode": "mistral",
|
||||
"enforce_eager": True,
|
||||
"gpu_memory_utilization": 0.9,
|
||||
**ROCM_ENGINE_KWARGS,
|
||||
}
|
||||
|
||||
|
||||
EXPECTED_TEXT = [
|
||||
(
|
||||
" First words I spoke in the original phonograph. "
|
||||
"A little piece of practical poetry. Mary had a little lamb,"
|
||||
" its fleece was quite a slow, and everywhere that Mary went, "
|
||||
"the lamb was sure to go."
|
||||
),
|
||||
(
|
||||
" And the 0-1 pitch on the way to Edgar Martinez. Swung on"
|
||||
" the line. Down the left field line for OBS. Here comes Joy. "
|
||||
"Here is Junior to third base. They're going to wave him in. "
|
||||
"The throw to the plate will be late. The Mariners are going"
|
||||
" to play. For the American League Championship, "
|
||||
"I don't believe it. It just continues. My, oh, my."
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _normalize(texts: list[str]) -> list[str]:
|
||||
# The model occasionally transcribes "OBS" as "a base hit" and
|
||||
# "oh, my" as "oh my", but both are acoustically valid. Normalise so
|
||||
# the assertion is stable across runs and hardware.
|
||||
texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")
|
||||
return texts
|
||||
|
||||
|
||||
def assert_encoder_kv_cache_spec(engine: LLM) -> None:
|
||||
vllm_config = engine.llm_engine.vllm_config
|
||||
audio_config = vllm_config.model_config.hf_config.audio_config
|
||||
kv_cache_specs_per_rank = engine.llm_engine.model_executor.get_kv_cache_specs()
|
||||
|
||||
assert len(kv_cache_specs_per_rank) == 1
|
||||
kv_cache_specs = kv_cache_specs_per_rank[0]
|
||||
assert AUDIO_LAYER_NAME in kv_cache_specs, kv_cache_specs.keys()
|
||||
spec = kv_cache_specs[AUDIO_LAYER_NAME]
|
||||
|
||||
assert audio_config.sliding_window == 750
|
||||
assert audio_config.block_pool_size == 4
|
||||
assert isinstance(spec, SlidingWindowSpec)
|
||||
assert spec.block_size == 16
|
||||
assert spec.num_kv_heads == 128
|
||||
# cdiv(750, 4) == 188 pooled tokens cover the model's window; the extra
|
||||
# +1 is an eviction margin (see whisper_causal.py get_kv_cache_spec).
|
||||
assert spec.sliding_window == cdiv(750, 4) + 1 == 189
|
||||
assert (
|
||||
spec.max_admission_blocks_per_request(
|
||||
max_in_flight_tokens=1,
|
||||
max_model_len=vllm_config.model_config.max_model_len,
|
||||
)
|
||||
== 13
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def audio_assets() -> list[AudioAsset]:
|
||||
return [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tokenizer() -> MistralTokenizer:
|
||||
return MistralTokenizer.from_hf_hub(MODEL_NAME)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def async_engine():
|
||||
gpu_memory_utilization = ENGINE_CONFIG.get("gpu_memory_utilization", 0.9)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.is_rocm():
|
||||
from tests.utils import wait_for_rocm_memory_to_settle
|
||||
|
||||
wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
|
||||
|
||||
engine_args = AsyncEngineArgs(**ENGINE_CONFIG)
|
||||
llm = AsyncLLM.from_engine_args(engine_args)
|
||||
try:
|
||||
yield llm
|
||||
finally:
|
||||
shutdown_timeout = 60.0 if current_platform.is_rocm() else None
|
||||
llm.shutdown(timeout=shutdown_timeout)
|
||||
del llm
|
||||
import torch
|
||||
|
||||
torch._dynamo.reset()
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
from tests.utils import wait_for_rocm_memory_to_settle
|
||||
|
||||
wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
|
||||
|
||||
|
||||
def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch):
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
|
||||
vllm_kwargs = {**ENGINE_CONFIG}
|
||||
vllm_kwargs["model_name"] = vllm_kwargs.pop("model")
|
||||
|
||||
with vllm_runner(**vllm_kwargs) as vllm_model:
|
||||
assert_encoder_kv_cache_spec(vllm_model.llm)
|
||||
audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
|
||||
|
||||
def from_file(file_path: str):
|
||||
audio = Audio.from_file(file_path, strict=False)
|
||||
req = TranscriptionRequest(
|
||||
audio=audio.to_base64(audio.format),
|
||||
streaming=StreamingMode.OFFLINE,
|
||||
language=None,
|
||||
)
|
||||
tokenized = tokenizer.instruct_tokenizer.encode_transcription(req)
|
||||
|
||||
return (tokenized.tokens, tokenized.audios[0].audio_array)
|
||||
|
||||
tokenized_list = [
|
||||
from_file(audio_asset.get_local_path()) for audio_asset in audio_assets
|
||||
]
|
||||
|
||||
inputs = []
|
||||
sampling_params = []
|
||||
|
||||
for tokens, audio_array in tokenized_list:
|
||||
num_samples = audio_array.shape[0]
|
||||
max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1
|
||||
sampling_params.append(
|
||||
SamplingParams(temperature=0.0, max_tokens=max_tokens)
|
||||
)
|
||||
|
||||
input_dict = {
|
||||
"multi_modal_data": {"audio": [(audio_array, None)]},
|
||||
"prompt_token_ids": tokens,
|
||||
}
|
||||
inputs.append(input_dict)
|
||||
|
||||
outputs = vllm_model.llm.generate(
|
||||
inputs,
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
|
||||
texts = _normalize([out.outputs[0].text for out in outputs])
|
||||
for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
|
||||
assert got == expected, (
|
||||
f"Output mismatch at index {i}:\n"
|
||||
f" got: {got!r}\n"
|
||||
f" expected: {expected!r}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine):
|
||||
# Lazy import to avoid CUDA-reinitialization error
|
||||
from vllm.model_executor.models.voxtral_realtime import VoxtralRealtimeBuffer
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=1)
|
||||
audio_config = tokenizer.instruct_tokenizer.audio_encoder.audio_config
|
||||
|
||||
output_tokens_list = []
|
||||
for i, audio_asset in enumerate(audio_assets):
|
||||
output_tokens = []
|
||||
audio = Audio.from_file(audio_asset.get_local_path(), strict=False)
|
||||
|
||||
req = TranscriptionRequest(
|
||||
streaming=StreamingMode.OFFLINE,
|
||||
audio=audio.to_base64(audio.format),
|
||||
language=None,
|
||||
)
|
||||
audio_enc = tokenizer.encode_transcription(req)
|
||||
|
||||
buffer = VoxtralRealtimeBuffer(audio_config, audio_enc.tokens)
|
||||
await buffer.append_audio(audio_enc.audios[0].audio_array)
|
||||
await buffer.append_audio(None)
|
||||
|
||||
request_id = f"session-{i}"
|
||||
|
||||
async for resp in async_engine.generate(
|
||||
prompt=buffer.get_input_stream(),
|
||||
sampling_params=sampling_params,
|
||||
request_id=request_id,
|
||||
):
|
||||
tokens = resp.outputs[0].token_ids[-1:]
|
||||
output_tokens.extend(tokens)
|
||||
await buffer.append_tokens(tokens)
|
||||
|
||||
output_tokens_list.append(output_tokens)
|
||||
|
||||
texts = _normalize(
|
||||
[
|
||||
tokenizer.decode(
|
||||
output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE
|
||||
)
|
||||
for output_tokens in output_tokens_list
|
||||
]
|
||||
)
|
||||
for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
|
||||
assert got == expected, (
|
||||
f"Output mismatch at index {i}:\n"
|
||||
f" got: {got!r}\n"
|
||||
f" expected: {expected!r}"
|
||||
)
|
||||
@@ -0,0 +1,368 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from transformers import AutoModelForSpeechSeq2Seq
|
||||
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.multimodal.audio import AudioResampler
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....conftest import HfRunner, PromptAudioInput, VllmRunner
|
||||
from ....utils import create_new_process_for_each_test, multi_gpu_test
|
||||
from ...registry import HF_EXAMPLE_MODELS
|
||||
from ...utils import check_logprobs_close
|
||||
|
||||
VLLM_PROMPT = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
|
||||
HF_PROMPT = ""
|
||||
# Whisper expects 16kHz audio
|
||||
WHISPER_SAMPLE_RATE = 16000
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def use_spawn_for_whisper(monkeypatch):
|
||||
"""Whisper has issues with forked workers, use spawn instead."""
|
||||
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
|
||||
def run_test(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
inputs: Sequence[tuple[list[str], list[str], PromptAudioInput]],
|
||||
model: str,
|
||||
*,
|
||||
max_model_len: int,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
tensor_parallel_size: int,
|
||||
distributed_executor_backend: str | None = None,
|
||||
enforce_eager: bool = True,
|
||||
gpu_memory_utilization: float = 0.9,
|
||||
) -> None:
|
||||
"""Inference result should be the same between hf and vllm.
|
||||
|
||||
All the audio fixtures for the test are from AudioAsset.
|
||||
For huggingface runner, we provide the audio as input.
|
||||
For vllm runner, we provide MultiModalDataDict objects
|
||||
and corresponding MultiModalConfig as input.
|
||||
"""
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_model_len=max_model_len,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
limit_mm_per_prompt={"audio": 2},
|
||||
enforce_eager=enforce_eager,
|
||||
gpu_memory_utilization=gpu_memory_utilization,
|
||||
disable_custom_all_reduce=True,
|
||||
) as vllm_model:
|
||||
vllm_outputs_per_case = [
|
||||
vllm_model.generate_greedy_logprobs(
|
||||
vllm_prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
audios=audios,
|
||||
)
|
||||
for vllm_prompts, _, audios in inputs
|
||||
]
|
||||
|
||||
with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSpeechSeq2Seq) as hf_model:
|
||||
hf_outputs_per_case = [
|
||||
hf_model.generate_greedy_logprobs_limit(
|
||||
hf_prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
audios=audios,
|
||||
)
|
||||
for _, hf_prompts, audios in inputs
|
||||
]
|
||||
|
||||
for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def resampled_assets() -> list[tuple[Any, int]]:
|
||||
audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
|
||||
sampled_assets = []
|
||||
resampler = AudioResampler(target_sr=WHISPER_SAMPLE_RATE)
|
||||
for asset in audio_assets:
|
||||
audio, orig_sr = asset.audio_and_sample_rate
|
||||
# Resample to Whisper's expected sample rate (16kHz)
|
||||
if orig_sr != WHISPER_SAMPLE_RATE:
|
||||
audio = resampler.resample(audio, orig_sr=orig_sr)
|
||||
sampled_assets.append(
|
||||
(audio, WHISPER_SAMPLE_RATE),
|
||||
)
|
||||
return sampled_assets
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def input_audios(
|
||||
resampled_assets,
|
||||
) -> list[tuple[list[str], list[str], list[tuple[Any, int]]]]:
|
||||
inputs = []
|
||||
# audio assets are resampled to WHISPER_SAMPLE_RATE
|
||||
for audio_info in resampled_assets:
|
||||
# vLLM prompts, HF prompts, audio inputs
|
||||
inputs.append(([VLLM_PROMPT], [HF_PROMPT], [audio_info]))
|
||||
return inputs
|
||||
|
||||
|
||||
def check_model_available(model: str) -> None:
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", [64])
|
||||
@pytest.mark.parametrize("beam_width", [1, 2])
|
||||
def test_beam_search_encoder_decoder(
|
||||
monkeypatch,
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
beam_width: int,
|
||||
resampled_assets,
|
||||
) -> None:
|
||||
"""Test beam search with encoder-decoder models (Whisper)."""
|
||||
if current_platform.is_rocm():
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
|
||||
|
||||
model = "openai/whisper-large-v3-turbo"
|
||||
check_model_available(model)
|
||||
|
||||
hf_prompts = [
|
||||
"<|startoftranscript|>",
|
||||
"<|startoftranscript|>",
|
||||
]
|
||||
|
||||
with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSpeechSeq2Seq) as hf_model:
|
||||
hf_outputs = hf_model.generate_beam_search(
|
||||
hf_prompts,
|
||||
beam_width=beam_width,
|
||||
max_tokens=max_tokens,
|
||||
audios=resampled_assets,
|
||||
)
|
||||
|
||||
# Test both explicit encoder/decoder prompts
|
||||
vllm_prompts = [
|
||||
# Implicit encoder/decoder prompt
|
||||
{
|
||||
"prompt": "<|startoftranscript|>",
|
||||
"multi_modal_data": {"audio": resampled_assets[0]},
|
||||
},
|
||||
# Explicit encoder/decover prompt
|
||||
{
|
||||
"encoder_prompt": {
|
||||
"prompt": "",
|
||||
"multi_modal_data": {"audio": resampled_assets[1]},
|
||||
},
|
||||
"decoder_prompt": "<|startoftranscript|>",
|
||||
},
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype="half",
|
||||
max_model_len=448,
|
||||
tensor_parallel_size=1,
|
||||
max_num_seqs=4,
|
||||
limit_mm_per_prompt={"audio": 2},
|
||||
enforce_eager=True,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_beam_search(
|
||||
vllm_prompts,
|
||||
beam_width=beam_width,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
for i in range(len(vllm_prompts)):
|
||||
hf_output_ids, hf_output_texts = hf_outputs[i]
|
||||
vllm_output_ids, vllm_output_texts = vllm_outputs[i]
|
||||
|
||||
for j, (hf_text, vllm_text) in enumerate(
|
||||
zip(hf_output_texts, vllm_output_texts)
|
||||
):
|
||||
print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:")
|
||||
print(hf_text)
|
||||
print(f">>>{j}-th vllm output:")
|
||||
print(vllm_text)
|
||||
|
||||
# Check that we got the same number of beams
|
||||
assert len(hf_output_ids) == len(vllm_output_ids)
|
||||
|
||||
# For encoder-decoder models, we primarily want to verify that:
|
||||
# 1. Beam search completes without errors
|
||||
# 2. We get the expected number of beams
|
||||
# 3. Outputs are reasonable (non-empty, diverse beams)
|
||||
for j in range(len(vllm_output_ids)):
|
||||
# Check that outputs are not empty
|
||||
assert len(vllm_output_ids[j]) > 0, f"Prompt {i}, beam {j}: empty output"
|
||||
# Check that decoded text is not empty
|
||||
assert len(vllm_output_texts[j].strip()) > 0, (
|
||||
f"Prompt {i}, beam {j}: empty text output"
|
||||
)
|
||||
|
||||
|
||||
def test_parse_language_detection_output():
|
||||
"""Unit test for WhisperForConditionalGeneration.parse_language_detection_output.
|
||||
|
||||
No GPU or model loading required.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from vllm.model_executor.models.whisper import (
|
||||
WhisperForConditionalGeneration,
|
||||
)
|
||||
|
||||
cls = WhisperForConditionalGeneration
|
||||
|
||||
def make_tokenizer(return_value: str) -> MagicMock:
|
||||
tok = MagicMock()
|
||||
tok.decode = MagicMock(return_value=return_value)
|
||||
return tok
|
||||
|
||||
# English
|
||||
assert (
|
||||
cls.parse_language_detection_output([50259], make_tokenizer("<|en|>")) == "en"
|
||||
)
|
||||
|
||||
# German
|
||||
assert (
|
||||
cls.parse_language_detection_output([50261], make_tokenizer("<|de|>")) == "de"
|
||||
)
|
||||
|
||||
# Unsupported language code
|
||||
with pytest.raises(AssertionError):
|
||||
cls.parse_language_detection_output([99999], make_tokenizer("<|xx|>"))
|
||||
|
||||
# No special token format
|
||||
with pytest.raises(AssertionError):
|
||||
cls.parse_language_detection_output([1], make_tokenizer("hello"))
|
||||
|
||||
# Empty token_ids
|
||||
with pytest.raises((AssertionError, IndexError)):
|
||||
cls.parse_language_detection_output([], make_tokenizer("anything"))
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu_model
|
||||
@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
|
||||
@pytest.mark.parametrize("dtype", ["half", "float"])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@pytest.mark.parametrize("enforce_eager", [True, False])
|
||||
def test_models(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
num_logprobs: int,
|
||||
input_audios,
|
||||
enforce_eager: bool,
|
||||
) -> None:
|
||||
check_model_available(model)
|
||||
if current_platform.is_cpu() and not enforce_eager:
|
||||
pytest.skip("Skipping test for CPU with non-eager mode")
|
||||
run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
input_audios,
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_model_len=448,
|
||||
max_tokens=200,
|
||||
num_logprobs=num_logprobs,
|
||||
tensor_parallel_size=1,
|
||||
enforce_eager=enforce_eager,
|
||||
)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
|
||||
@pytest.mark.parametrize("distributed_executor_backend", ["ray", "mp"])
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", [200])
|
||||
@pytest.mark.parametrize("num_logprobs", [5])
|
||||
@create_new_process_for_each_test("spawn")
|
||||
def test_models_distributed(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
model: str,
|
||||
distributed_executor_backend: str,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
input_audios,
|
||||
) -> None:
|
||||
check_model_available(model)
|
||||
run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
input_audios,
|
||||
model,
|
||||
dtype=dtype,
|
||||
max_model_len=448,
|
||||
max_tokens=max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
tensor_parallel_size=2,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
enforce_eager=False,
|
||||
gpu_memory_utilization=0.65,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
|
||||
def test_encoder_cache_cleanup(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
input_audios,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
"""Test that encoder cache is properly cleaned up after requests complete.
|
||||
|
||||
This is a regression test for a bug where encoder cache entries were freed
|
||||
in the same scheduling step they were allocated, before the model could use
|
||||
them.
|
||||
"""
|
||||
# Set single-process mode to access the model runner's encoder cache directly
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
check_model_available(model)
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
dtype="half",
|
||||
max_model_len=448,
|
||||
tensor_parallel_size=1,
|
||||
limit_mm_per_prompt={"audio": 2},
|
||||
enforce_eager=True,
|
||||
) as vllm_model:
|
||||
engine_core = vllm_model.llm.llm_engine.engine_core.engine_core
|
||||
model_runner = engine_core.model_executor.driver_worker.worker.model_runner
|
||||
encoder_cache = model_runner.encoder_cache
|
||||
|
||||
# Run multiple sequential requests to ensure cache is properly managed
|
||||
for vllm_prompts, _, audios in input_audios:
|
||||
vllm_model.generate_greedy(vllm_prompts, max_tokens=50, audios=audios)
|
||||
|
||||
# After all requests complete, encoder cache should be empty
|
||||
cache_size = len(encoder_cache)
|
||||
assert cache_size == 0, (
|
||||
f"Encoder cache should be empty after all requests complete, "
|
||||
f"but has {cache_size} entries. This indicates encoder cache "
|
||||
f"entries are not being properly freed."
|
||||
)
|
||||
@@ -0,0 +1,344 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Helpers for building inputs that can be leveraged for different test types."""
|
||||
|
||||
from collections.abc import Callable, Iterable
|
||||
from pathlib import PosixPath
|
||||
from typing import Any
|
||||
|
||||
import numpy.typing as npt
|
||||
import torch
|
||||
|
||||
from vllm.multimodal.audio import AudioResampler
|
||||
from vllm.multimodal.image import rescale_image_size
|
||||
from vllm.multimodal.video import (
|
||||
rescale_video_size,
|
||||
resize_video,
|
||||
sample_frames_from_video,
|
||||
)
|
||||
|
||||
from .....conftest import AudioTestAssets, ImageTestAssets, VideoTestAssets
|
||||
from .types import (
|
||||
SINGLE_AUDIO_BASE_PROMPT,
|
||||
SINGLE_IMAGE_BASE_PROMPTS,
|
||||
TEST_AUDIO_PLACEHOLDER,
|
||||
TEST_IMG_PLACEHOLDER,
|
||||
TEST_VIDEO_PLACEHOLDER,
|
||||
VIDEO_BASE_PROMPT,
|
||||
ImageSizeWrapper,
|
||||
PromptWithMultiModalInput,
|
||||
SizeType,
|
||||
VLMTestInfo,
|
||||
)
|
||||
|
||||
|
||||
def replace_test_placeholder(
|
||||
prompt: str, mm_idx_to_prompt: Callable[[int], str], test_placeholder: str
|
||||
) -> str:
|
||||
"""Given a prompt, replaces each test placeholder with the
|
||||
model-specific tag.
|
||||
"""
|
||||
prompt_segments = prompt.split(test_placeholder)
|
||||
img_prompt = prompt_segments[0]
|
||||
for placeholder_idx, next_seg in enumerate(prompt_segments[1:], start=1):
|
||||
img_prompt += mm_idx_to_prompt(placeholder_idx)
|
||||
img_prompt += next_seg
|
||||
return img_prompt
|
||||
|
||||
|
||||
def get_model_prompts(
|
||||
base_prompts: Iterable[str],
|
||||
img_idx_to_prompt: Callable[[int], str] | None,
|
||||
video_idx_to_prompt: Callable[[int], str] | None,
|
||||
audio_idx_to_prompt: Callable[[int], str] | None,
|
||||
prompt_formatter: Callable[[str], str],
|
||||
) -> list[str]:
|
||||
"""Given a model-agnostic base prompt and test configuration for a model(s)
|
||||
to be tested, update the media placeholders and apply the prompt formatting
|
||||
to get the test prompt string for this model.
|
||||
|
||||
Example for phi3v, given the base_prompt: "<image>What is the season?"
|
||||
1. Replace img placeholder(s)
|
||||
-> "<|image_1|>\nWhat is the season?"
|
||||
2. Apply prompt formatter:
|
||||
-> <|user|>\n<|image_1|>\nWhat is the season?<|end|>\n<|assistant|>\n
|
||||
"""
|
||||
assert isinstance(base_prompts, (list, tuple))
|
||||
model_prompts = []
|
||||
for base_prompt in base_prompts:
|
||||
# Replace the multimodal placeholders in the base prompt with
|
||||
# the correct ones for the model that we are testing
|
||||
if img_idx_to_prompt:
|
||||
base_prompt = replace_test_placeholder(
|
||||
base_prompt, img_idx_to_prompt, TEST_IMG_PLACEHOLDER
|
||||
)
|
||||
|
||||
if video_idx_to_prompt:
|
||||
base_prompt = replace_test_placeholder(
|
||||
base_prompt, video_idx_to_prompt, TEST_VIDEO_PLACEHOLDER
|
||||
)
|
||||
|
||||
if audio_idx_to_prompt:
|
||||
base_prompt = replace_test_placeholder(
|
||||
base_prompt, audio_idx_to_prompt, TEST_AUDIO_PLACEHOLDER
|
||||
)
|
||||
|
||||
# Apply the prompt formatter to wrap the base prompt with
|
||||
# the correct media placeholders to get the model test prompt
|
||||
model_prompt = prompt_formatter(base_prompt)
|
||||
model_prompts.append(model_prompt)
|
||||
return model_prompts
|
||||
|
||||
|
||||
def build_single_image_inputs_from_test_info(
|
||||
test_info: VLMTestInfo,
|
||||
image_assets: ImageTestAssets,
|
||||
size_wrapper: ImageSizeWrapper,
|
||||
tmp_path: PosixPath | None = None,
|
||||
) -> list[PromptWithMultiModalInput]:
|
||||
if test_info.prompt_formatter is None:
|
||||
raise ValueError("Prompt formatter must be set to build single image inputs")
|
||||
|
||||
model_prompts = get_model_prompts(
|
||||
test_info.single_image_prompts,
|
||||
test_info.img_idx_to_prompt,
|
||||
test_info.video_idx_to_prompt,
|
||||
test_info.audio_idx_to_prompt,
|
||||
test_info.prompt_formatter,
|
||||
)
|
||||
|
||||
# For models that require a local path / URL encoded in the image; export
|
||||
# assets and encode into tmp_path for this test. This should be avoided
|
||||
# where possible (currently needed for Qwen-VL).
|
||||
if test_info.prompt_path_encoder is not None:
|
||||
if tmp_path is None:
|
||||
raise ValueError("Prompt path encoder requires setting local path")
|
||||
model_prompts = [
|
||||
test_info.prompt_path_encoder(tmp_path, prompt, [asset])
|
||||
for prompt, asset in zip(model_prompts, image_assets)
|
||||
]
|
||||
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
assert len(images) == len(model_prompts)
|
||||
return build_single_image_inputs(images, model_prompts, size_wrapper)
|
||||
|
||||
|
||||
def build_single_image_inputs(
|
||||
images, model_prompts, size_wrapper: ImageSizeWrapper
|
||||
) -> list[PromptWithMultiModalInput]:
|
||||
# For every image / prompt pair, get a pair containing two lists of
|
||||
# length size_factors, where the first contains duplicates of the model
|
||||
# prompt [str], and the second contains copies of the image after being
|
||||
# scaled by one of the size factors.
|
||||
#
|
||||
# NOTE: rescaling preserves the image aspect ratio.
|
||||
return [
|
||||
PromptWithMultiModalInput(
|
||||
prompts=[prompt for _ in size_wrapper.data],
|
||||
image_data=[
|
||||
apply_image_size_scaling(image, size, size_wrapper.type)
|
||||
for size in size_wrapper.data
|
||||
],
|
||||
)
|
||||
for image, prompt in zip(images, model_prompts)
|
||||
]
|
||||
|
||||
|
||||
def build_multi_image_inputs_from_test_info(
|
||||
test_info: VLMTestInfo,
|
||||
image_assets: ImageTestAssets,
|
||||
size_wrapper: ImageSizeWrapper,
|
||||
tmp_path: PosixPath | None = None,
|
||||
) -> list[PromptWithMultiModalInput]:
|
||||
if test_info.prompt_formatter is None:
|
||||
raise ValueError("Prompt formatter must be set to build multi image inputs")
|
||||
|
||||
model_prompts = get_model_prompts(
|
||||
[test_info.multi_image_prompt],
|
||||
test_info.img_idx_to_prompt,
|
||||
test_info.video_idx_to_prompt,
|
||||
test_info.audio_idx_to_prompt,
|
||||
test_info.prompt_formatter,
|
||||
)
|
||||
|
||||
if test_info.prompt_path_encoder is not None:
|
||||
if tmp_path is None:
|
||||
raise ValueError("Prompt path encoder requires setting local path")
|
||||
model_prompts = [
|
||||
test_info.prompt_path_encoder(tmp_path, model_prompt, image_assets)
|
||||
for model_prompt in model_prompts
|
||||
]
|
||||
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
|
||||
# Currently, we only have one multi-image list & one multi-image prompt
|
||||
return build_multi_image_inputs(
|
||||
image_lists=[images],
|
||||
model_prompts=model_prompts,
|
||||
size_wrapper=size_wrapper,
|
||||
)
|
||||
|
||||
|
||||
def build_multi_image_inputs(
|
||||
image_lists, model_prompts, size_wrapper: ImageSizeWrapper
|
||||
) -> list[PromptWithMultiModalInput]:
|
||||
return [
|
||||
PromptWithMultiModalInput(
|
||||
prompts=[prompt for _ in size_wrapper.data],
|
||||
image_data=[
|
||||
[
|
||||
apply_image_size_scaling(image, size, size_wrapper.type)
|
||||
for image in images
|
||||
]
|
||||
for size in size_wrapper.data
|
||||
],
|
||||
)
|
||||
for images, prompt in zip(image_lists, model_prompts)
|
||||
]
|
||||
|
||||
|
||||
def build_embedding_inputs_from_test_info(
|
||||
test_info: VLMTestInfo,
|
||||
image_assets: ImageTestAssets,
|
||||
size_wrapper: ImageSizeWrapper,
|
||||
):
|
||||
# These conditions will always be true if invoked through filtering,
|
||||
# but we still check them in case this is ever called directly
|
||||
if test_info.prompt_formatter is None:
|
||||
raise ValueError("Prompt formatter must be set to build image embedding inputs")
|
||||
if size_wrapper.type != SizeType.SIZE_FACTOR or not all(
|
||||
factor == 1.0 for factor in size_wrapper.data
|
||||
):
|
||||
raise ValueError("Embedding tests require constant (1.0) size factors")
|
||||
if test_info.convert_assets_to_embeddings is None:
|
||||
raise ValueError("No conversion func for getting embeddings found")
|
||||
|
||||
model_prompts = get_model_prompts(
|
||||
SINGLE_IMAGE_BASE_PROMPTS,
|
||||
test_info.img_idx_to_prompt,
|
||||
test_info.video_idx_to_prompt,
|
||||
test_info.audio_idx_to_prompt,
|
||||
test_info.prompt_formatter,
|
||||
)
|
||||
|
||||
images = [asset.pil_image for asset in image_assets]
|
||||
embeds = test_info.convert_assets_to_embeddings(image_assets)
|
||||
if test_info.dtype != "auto":
|
||||
dtype = getattr(torch, test_info.dtype) # type: ignore
|
||||
embeds = [e.to(dtype=dtype) for e in embeds]
|
||||
assert len(images) == len(model_prompts)
|
||||
|
||||
inputs = build_single_image_inputs(images, model_prompts, size_wrapper)
|
||||
vllm_embeddings = build_single_image_inputs(embeds, model_prompts, size_wrapper)
|
||||
return inputs, vllm_embeddings
|
||||
|
||||
|
||||
def build_video_inputs_from_test_info(
|
||||
test_info: VLMTestInfo,
|
||||
video_assets: VideoTestAssets,
|
||||
size_wrapper: ImageSizeWrapper,
|
||||
num_frames: int,
|
||||
needs_video_metadata: bool,
|
||||
) -> list[PromptWithMultiModalInput]:
|
||||
if test_info.prompt_formatter is None:
|
||||
raise ValueError("Prompt formatter must be set to build video inputs")
|
||||
model_prompts = get_model_prompts(
|
||||
[VIDEO_BASE_PROMPT],
|
||||
test_info.img_idx_to_prompt,
|
||||
test_info.video_idx_to_prompt,
|
||||
test_info.audio_idx_to_prompt,
|
||||
test_info.prompt_formatter,
|
||||
)
|
||||
|
||||
sampled_vids = [
|
||||
sample_frames_with_video_metadata(
|
||||
(asset.np_ndarrays, asset.metadata),
|
||||
num_frames,
|
||||
)
|
||||
for asset in video_assets
|
||||
]
|
||||
|
||||
video_scaler = (
|
||||
resize_video if size_wrapper.type == SizeType.FIXED_SIZE else rescale_video_size
|
||||
)
|
||||
|
||||
return [
|
||||
PromptWithMultiModalInput(
|
||||
prompts=[prompt for _ in size_wrapper.data],
|
||||
video_data=[
|
||||
(
|
||||
video_scaler(video, size)
|
||||
if not needs_video_metadata
|
||||
else (video_scaler(video, size), meta)
|
||||
)
|
||||
for size in size_wrapper.data
|
||||
],
|
||||
)
|
||||
for (video, meta), prompt in zip(sampled_vids, model_prompts)
|
||||
]
|
||||
|
||||
|
||||
def sample_frames_with_video_metadata(
|
||||
video_with_meta: tuple[npt.NDArray, dict[str, Any]],
|
||||
num_frames: int,
|
||||
) -> tuple[npt.NDArray, dict[str, Any]]:
|
||||
video, meta = video_with_meta
|
||||
video = sample_frames_from_video(video, num_frames)
|
||||
|
||||
meta["do_sample_frames"] = meta["total_num_frames"] == num_frames
|
||||
meta["total_num_frames"] = num_frames
|
||||
meta["fps"] = meta["duration"] / num_frames
|
||||
meta["frames_indices"] = list(range(num_frames))
|
||||
return video, meta
|
||||
|
||||
|
||||
def apply_image_size_scaling(image, size: float | tuple[int, int], size_type: SizeType):
|
||||
"""Applies a size scaler to one image; this can be an image size factor,
|
||||
which scales the image while maintaining the aspect ratio"""
|
||||
# Special case for embeddings; if it's a tensor, it's only valid if we
|
||||
# are considering size factors at constant scale, i.e., we just clone
|
||||
# the tensor
|
||||
if isinstance(image, torch.Tensor):
|
||||
assert size_type == SizeType.SIZE_FACTOR and size == 1
|
||||
return image
|
||||
if size_type == SizeType.SIZE_FACTOR:
|
||||
# We have a list of image size factors
|
||||
return rescale_image_size(image, size)
|
||||
elif size_type == SizeType.FIXED_SIZE:
|
||||
# We have a list of fixed sizes
|
||||
return image.resize(size)
|
||||
raise ValueError("ImageSizeWrapper type must be FIXED_SIZE or SIZE_FACTOR")
|
||||
|
||||
|
||||
def build_audio_inputs_from_test_info(
|
||||
test_info: VLMTestInfo,
|
||||
audio_assets: AudioTestAssets,
|
||||
) -> list[PromptWithMultiModalInput]:
|
||||
if test_info.prompt_formatter is None:
|
||||
raise ValueError("Prompt formatter must be set to build audio inputs")
|
||||
model_prompts = get_model_prompts(
|
||||
SINGLE_AUDIO_BASE_PROMPT,
|
||||
test_info.img_idx_to_prompt,
|
||||
test_info.video_idx_to_prompt,
|
||||
test_info.audio_idx_to_prompt,
|
||||
test_info.prompt_formatter,
|
||||
)
|
||||
resampler = AudioResampler(target_sr=16000)
|
||||
audios = [asset.audio_and_sample_rate for asset in audio_assets]
|
||||
resampled_audios = [
|
||||
(
|
||||
resampler.resample(
|
||||
audio,
|
||||
orig_sr=sr,
|
||||
),
|
||||
int(resampler.target_sr),
|
||||
)
|
||||
for audio, sr in audios
|
||||
]
|
||||
|
||||
return [
|
||||
PromptWithMultiModalInput(
|
||||
prompts=model_prompts,
|
||||
audio_data=resampled_audios,
|
||||
)
|
||||
]
|
||||
@@ -0,0 +1,183 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Utils for determining which subset of model tests belong to a specific
|
||||
modality, getting all combinations (similar to pytest's parametrization),
|
||||
handling multimodal placeholder substitution, and so on.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Iterable
|
||||
|
||||
import pytest
|
||||
|
||||
from .types import (
|
||||
EMBEDDING_SIZE_FACTORS,
|
||||
ExpandableVLMTestArgs,
|
||||
ImageSizeWrapper,
|
||||
SizeType,
|
||||
VLMTestInfo,
|
||||
VLMTestType,
|
||||
)
|
||||
|
||||
|
||||
def get_filtered_test_settings(
|
||||
test_settings: dict[str, VLMTestInfo],
|
||||
test_type: VLMTestType,
|
||||
new_proc_per_test: bool,
|
||||
) -> dict[str, VLMTestInfo]:
|
||||
"""Given the dict of potential test settings to run, return a subdict
|
||||
of tests who have the current test type enabled with the matching val for
|
||||
fork_per_test.
|
||||
"""
|
||||
|
||||
def matches_test_type(test_info: VLMTestInfo, test_type: VLMTestType):
|
||||
return test_info.test_type == test_type or (
|
||||
isinstance(test_info.test_type, Iterable)
|
||||
and test_type in test_info.test_type
|
||||
)
|
||||
|
||||
matching_tests = {}
|
||||
for test_name, test_info in test_settings.items():
|
||||
# Otherwise check if the test has the right type & keep if it does
|
||||
if matches_test_type(test_info, test_type):
|
||||
# Embedding tests need to have a conversion func in their test info
|
||||
if matches_test_type(test_info, VLMTestType.EMBEDDING):
|
||||
assert test_info.convert_assets_to_embeddings is not None
|
||||
# Custom test inputs need to explicitly define the mm limit/inputs
|
||||
if matches_test_type(test_info, VLMTestType.CUSTOM_INPUTS):
|
||||
assert test_info.custom_test_opts is not None and isinstance(
|
||||
test_info.custom_test_opts, Iterable
|
||||
)
|
||||
# For all types besides custom inputs, we need a prompt formatter
|
||||
else:
|
||||
assert test_info.prompt_formatter is not None
|
||||
|
||||
# Everything looks okay; keep if this is correct proc handling
|
||||
if (
|
||||
test_info.distributed_executor_backend is not None
|
||||
) == new_proc_per_test:
|
||||
matching_tests[test_name] = test_info
|
||||
|
||||
return matching_tests
|
||||
|
||||
|
||||
def get_model_type_cases(
|
||||
model_type: str,
|
||||
test_info: VLMTestInfo,
|
||||
test_type: VLMTestType,
|
||||
):
|
||||
# Ensure that something is wrapped as an iterable it's not already
|
||||
ensure_wrapped = lambda e: e if isinstance(e, (list, tuple)) else (e,)
|
||||
|
||||
# This is essentially the same as nesting a bunch of mark.parametrize
|
||||
# decorators, but we do it programmatically to allow overrides for on
|
||||
# a per-model basis, while still being able to execute each of these
|
||||
# as individual test cases in pytest.
|
||||
iter_kwargs = OrderedDict(
|
||||
[
|
||||
("model", ensure_wrapped(test_info.models)),
|
||||
("max_tokens", ensure_wrapped(test_info.max_tokens)),
|
||||
("num_logprobs", ensure_wrapped(test_info.num_logprobs)),
|
||||
("dtype", ensure_wrapped(test_info.dtype)),
|
||||
(
|
||||
"distributed_executor_backend",
|
||||
ensure_wrapped(test_info.distributed_executor_backend),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
# num_frames is video only
|
||||
if test_type == VLMTestType.VIDEO:
|
||||
iter_kwargs["num_video_frames"] = ensure_wrapped(test_info.num_video_frames)
|
||||
iter_kwargs["needs_video_metadata"] = ensure_wrapped(
|
||||
test_info.needs_video_metadata
|
||||
)
|
||||
|
||||
# No sizes passed for custom inputs, since inputs are directly provided
|
||||
if test_type not in (
|
||||
VLMTestType.CUSTOM_INPUTS,
|
||||
VLMTestType.AUDIO,
|
||||
):
|
||||
wrapped_sizes = get_wrapped_test_sizes(test_info, test_type)
|
||||
if wrapped_sizes is None:
|
||||
raise ValueError(f"Sizes must be set for test type {test_type}")
|
||||
iter_kwargs["size_wrapper"] = wrapped_sizes
|
||||
|
||||
# Otherwise expand the custom test options instead
|
||||
elif test_type == VLMTestType.CUSTOM_INPUTS:
|
||||
if test_info.custom_test_opts is None:
|
||||
raise ValueError("Test has type CUSTOM_INPUTS, but none given")
|
||||
iter_kwargs["custom_test_opts"] = test_info.custom_test_opts
|
||||
|
||||
# Wrap all model cases in a pytest parameter & pass marks through
|
||||
return [
|
||||
pytest.param(
|
||||
model_type,
|
||||
ExpandableVLMTestArgs(**{k: v for k, v in zip(iter_kwargs.keys(), case)}),
|
||||
marks=test_info.marks if test_info.marks is not None else [],
|
||||
)
|
||||
for case in list(itertools.product(*iter_kwargs.values()))
|
||||
]
|
||||
|
||||
|
||||
def get_parametrized_options(
|
||||
test_settings: dict[str, VLMTestInfo],
|
||||
test_type: VLMTestType,
|
||||
create_new_process_for_each_test: bool,
|
||||
):
|
||||
"""Converts all of our VLMTestInfo into an expanded list of parameters.
|
||||
This is similar to nesting pytest parametrize calls, but done directly
|
||||
through an itertools product so that each test can set things like
|
||||
size factors etc, while still running in isolated test cases.
|
||||
"""
|
||||
matching_tests = get_filtered_test_settings(
|
||||
test_settings, test_type, create_new_process_for_each_test
|
||||
)
|
||||
|
||||
# Get a list per model type, where each entry contains a tuple of all of
|
||||
# that model type's cases, then flatten them into the top level so that
|
||||
# we can consume them in one mark.parametrize call.
|
||||
cases_by_model_type = [
|
||||
get_model_type_cases(model_type, test_info, test_type)
|
||||
for model_type, test_info in matching_tests.items()
|
||||
]
|
||||
return list(itertools.chain(*cases_by_model_type))
|
||||
|
||||
|
||||
def get_wrapped_test_sizes(
|
||||
test_info: VLMTestInfo, test_type: VLMTestType
|
||||
) -> tuple[ImageSizeWrapper, ...]:
|
||||
"""Given a test info which may have size factors or fixed sizes, wrap them
|
||||
and combine them into an iterable, each of which will be used in parameter
|
||||
expansion.
|
||||
|
||||
Args:
|
||||
test_info: Test configuration to be expanded.
|
||||
test_type: The type of test being filtered for.
|
||||
"""
|
||||
# If it is an embedding test, we always use the EMBEDDING_SIZE_FACTORS
|
||||
if test_type == VLMTestType.EMBEDDING:
|
||||
return tuple(
|
||||
[
|
||||
ImageSizeWrapper(type=SizeType.SIZE_FACTOR, data=factor)
|
||||
for factor in EMBEDDING_SIZE_FACTORS
|
||||
]
|
||||
)
|
||||
# Audio and Custom inputs have preprocessed inputs
|
||||
elif test_type in (VLMTestType.AUDIO, VLMTestType.CUSTOM_INPUTS):
|
||||
return tuple()
|
||||
|
||||
size_factors = test_info.image_size_factors if test_info.image_size_factors else []
|
||||
fixed_sizes = test_info.image_sizes if test_info.image_sizes else []
|
||||
|
||||
wrapped_factors = [
|
||||
ImageSizeWrapper(type=SizeType.SIZE_FACTOR, data=factor)
|
||||
for factor in size_factors
|
||||
]
|
||||
|
||||
wrapped_sizes = [
|
||||
ImageSizeWrapper(type=SizeType.FIXED_SIZE, data=size) for size in fixed_sizes
|
||||
]
|
||||
|
||||
return tuple(wrapped_factors + wrapped_sizes)
|
||||
@@ -0,0 +1,207 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Core test implementation to be shared across modalities."""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from transformers.models.auto.auto_factory import _BaseAutoModelClass
|
||||
|
||||
from vllm.config.model import RunnerOption
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
|
||||
from .....conftest import HfRunner, VllmRunner
|
||||
from ....registry import HF_EXAMPLE_MODELS
|
||||
from .types import PromptWithMultiModalInput, RunnerOutput
|
||||
|
||||
|
||||
def run_test(
|
||||
*,
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
inputs: list[PromptWithMultiModalInput],
|
||||
model: str,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
num_logprobs: int,
|
||||
enforce_eager: bool,
|
||||
max_model_len: int,
|
||||
max_num_seqs: int,
|
||||
hf_output_post_proc: Callable[[RunnerOutput, str], Any] | None,
|
||||
vllm_output_post_proc: Callable[[RunnerOutput, str], Any] | None,
|
||||
auto_cls: type[_BaseAutoModelClass],
|
||||
use_tokenizer_eos: bool,
|
||||
comparator: Callable[..., None],
|
||||
get_stop_token_ids: Callable[[TokenizerLike], list[int]] | None,
|
||||
stop_str: list[str] | None,
|
||||
limit_mm_per_prompt: dict[str, int],
|
||||
vllm_runner_kwargs: dict[str, Any] | None,
|
||||
hf_model_kwargs: dict[str, Any] | None,
|
||||
hf_processor: Callable[[str], Any] | None,
|
||||
patch_hf_runner: Callable[[HfRunner], HfRunner] | None,
|
||||
runner: RunnerOption = "auto",
|
||||
distributed_executor_backend: str | None = None,
|
||||
tensor_parallel_size: int = 1,
|
||||
vllm_embeddings: torch.Tensor | None = None,
|
||||
):
|
||||
"""Modality agnostic test executor for comparing HF/vLLM outputs."""
|
||||
# In the case of embeddings, vLLM takes separate input tensors
|
||||
vllm_inputs = vllm_embeddings if vllm_embeddings is not None else inputs
|
||||
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
# Disable other modalities to save memory
|
||||
default_limits = {"image": 0, "video": 0, "audio": 0}
|
||||
limit_mm_per_prompt = default_limits | limit_mm_per_prompt
|
||||
|
||||
vllm_outputs_per_mm = []
|
||||
hf_outputs_per_mm = []
|
||||
|
||||
# NOTE: take care of the order. run vLLM first, and then run HF.
|
||||
# vLLM needs a fresh new process without cuda initialization.
|
||||
# if we run HF first, the cuda initialization will be done and it
|
||||
# will hurt multiprocessing backend with fork method (the default method).
|
||||
|
||||
vllm_runner_kwargs_: dict[str, Any] = {"mm_processor_cache_gb": 0}
|
||||
if model_info.tokenizer:
|
||||
vllm_runner_kwargs_["tokenizer_name"] = model_info.tokenizer
|
||||
if model_info.tokenizer_mode:
|
||||
vllm_runner_kwargs_["tokenizer_mode"] = model_info.tokenizer_mode
|
||||
if model_info.hf_overrides:
|
||||
vllm_runner_kwargs_["hf_overrides"] = model_info.hf_overrides
|
||||
if model_info.require_embed_inputs:
|
||||
for k in ("skip_tokenizer_init", "enable_prompt_embeds", "enable_mm_embeds"):
|
||||
vllm_runner_kwargs_[k] = model_info.require_embed_inputs
|
||||
if not model_info.enable_prefix_caching:
|
||||
vllm_runner_kwargs_["enable_prefix_caching"] = False
|
||||
|
||||
if vllm_runner_kwargs:
|
||||
vllm_runner_kwargs_.update(vllm_runner_kwargs)
|
||||
|
||||
# Avoid passing limit_mm_per_prompt twice when vllm_runner_kwargs
|
||||
# already contains it (e.g. gemma4 sets it via vllm_runner_kwargs).
|
||||
if "limit_mm_per_prompt" in vllm_runner_kwargs_:
|
||||
limit_mm_per_prompt = vllm_runner_kwargs_.pop("limit_mm_per_prompt")
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=max_model_len,
|
||||
max_num_seqs=max_num_seqs,
|
||||
dtype=dtype,
|
||||
limit_mm_per_prompt=limit_mm_per_prompt,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
enforce_eager=enforce_eager,
|
||||
runner=runner,
|
||||
**vllm_runner_kwargs_,
|
||||
) as vllm_model:
|
||||
tokenizer = vllm_model.llm.get_tokenizer()
|
||||
|
||||
vllm_kwargs: dict[str, Any] = {}
|
||||
if get_stop_token_ids is not None:
|
||||
vllm_kwargs["stop_token_ids"] = get_stop_token_ids(tokenizer)
|
||||
if stop_str:
|
||||
vllm_kwargs["stop"] = stop_str
|
||||
|
||||
for prompts, image_data, video_data, audio_data in vllm_inputs:
|
||||
mm_data = dict(images=image_data, videos=video_data, audios=audio_data)
|
||||
vllm_kwargs_with_mm_data = vllm_kwargs | mm_data
|
||||
vllm_output = vllm_model.generate_greedy_logprobs(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
**vllm_kwargs_with_mm_data,
|
||||
)
|
||||
vllm_outputs_per_mm.append(vllm_output)
|
||||
|
||||
hf_runner_kwargs: dict[str, Any] = {}
|
||||
if model_info.tokenizer:
|
||||
hf_runner_kwargs["tokenizer_name"] = model_info.tokenizer
|
||||
if hf_processor is not None:
|
||||
hf_runner_kwargs["processor"] = hf_processor(model)
|
||||
|
||||
hf_model = hf_runner(
|
||||
model,
|
||||
dtype=dtype,
|
||||
auto_cls=auto_cls,
|
||||
model_kwargs=hf_model_kwargs,
|
||||
**hf_runner_kwargs,
|
||||
)
|
||||
|
||||
# Some models need to patch things like the model processor, e.g., internvl
|
||||
if patch_hf_runner is not None:
|
||||
hf_model = patch_hf_runner(hf_model)
|
||||
|
||||
with hf_model, torch.no_grad():
|
||||
tokenizer = hf_model.tokenizer
|
||||
|
||||
# Some models need to explicitly pass the eos_token_id off the tokenizer
|
||||
# or processor for a good comparison;
|
||||
# currently assume processor/tokenizer agree on the EOS, and pull it off
|
||||
# the tokenizer if requested.
|
||||
hf_kwargs = {}
|
||||
if use_tokenizer_eos:
|
||||
hf_kwargs["eos_token_id"] = tokenizer.eos_token_id
|
||||
if stop_str:
|
||||
hf_kwargs["stop_strings"] = stop_str
|
||||
|
||||
for prompts, image_data, video_data, audio_data in inputs:
|
||||
mm_data = dict(images=image_data, videos=video_data, audios=audio_data)
|
||||
hf_kwargs_with_mm_data = hf_kwargs | mm_data
|
||||
hf_output = hf_model.generate_greedy_logprobs_limit(
|
||||
prompts,
|
||||
max_tokens,
|
||||
num_logprobs=num_logprobs,
|
||||
tokenizer=tokenizer,
|
||||
**hf_kwargs_with_mm_data,
|
||||
)
|
||||
hf_outputs_per_mm.append(hf_output)
|
||||
|
||||
# Apply output processing / sanitation to the vLLM and HF runner results
|
||||
hf_outputs_per_mm, vllm_outputs_per_mm = process_runner_outputs(
|
||||
model,
|
||||
first_runner_outputs=hf_outputs_per_mm,
|
||||
second_runner_outputs=vllm_outputs_per_mm,
|
||||
first_runner_processor=hf_output_post_proc,
|
||||
second_runner_processor=vllm_output_post_proc,
|
||||
)
|
||||
|
||||
for hf_outputs, vllm_outputs in zip(hf_outputs_per_mm, vllm_outputs_per_mm):
|
||||
# This is usually check_logprobs_close, but it's passed through to
|
||||
# allow things like check_outputs_equal where needed
|
||||
comparator(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
def process_runner_outputs(
|
||||
model,
|
||||
first_runner_outputs,
|
||||
second_runner_outputs,
|
||||
first_runner_processor=None,
|
||||
second_runner_processor=None,
|
||||
):
|
||||
"""Applies the runner processor(s) to the runner outputs, if any."""
|
||||
if first_runner_processor is not None:
|
||||
first_runner_outputs = process_outputs(
|
||||
first_runner_processor, model, first_runner_outputs
|
||||
)
|
||||
if second_runner_processor is not None:
|
||||
second_runner_outputs = process_outputs(
|
||||
second_runner_processor, model, second_runner_outputs
|
||||
)
|
||||
return first_runner_outputs, second_runner_outputs
|
||||
|
||||
|
||||
def process_outputs(output_processor, model, outputs_per_image):
|
||||
"""Applies a model specific post-processor function to a runner's output"""
|
||||
return [
|
||||
[output_processor(res, model) for res in outputs]
|
||||
for outputs in outputs_per_image
|
||||
]
|
||||
@@ -0,0 +1,156 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Custom input builders for edge-cases in different models."""
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.multimodal.image import rescale_image_size
|
||||
from vllm.multimodal.video import (
|
||||
rescale_video_size,
|
||||
resize_video,
|
||||
sample_frames_from_video,
|
||||
)
|
||||
|
||||
from .....conftest import IMAGE_ASSETS, VIDEO_ASSETS
|
||||
from .builders import build_multi_image_inputs, build_single_image_inputs
|
||||
from .types import ImageSizeWrapper, PromptWithMultiModalInput, SizeType
|
||||
|
||||
|
||||
def multi_image_multi_aspect_ratio_inputs(formatter: Callable[[str], str]):
|
||||
"""Builds inputs for multi-image (varied sizes/aspect ratio) testing.
|
||||
|
||||
Args:
|
||||
formatter: model-specific prompt formatter.
|
||||
"""
|
||||
stop_sign = IMAGE_ASSETS[0].pil_image
|
||||
cherry_blossom = IMAGE_ASSETS[1].pil_image
|
||||
|
||||
# Apply the selected formatter to the base prompts
|
||||
img_prompts = [
|
||||
"<image><image>\nDescribe 2 images.",
|
||||
"<image><image>\nDescribe 2 images.",
|
||||
"<image><image><image><image>\nDescribe 4 images.",
|
||||
"<image>\nWhat is the season?",
|
||||
]
|
||||
formatted_prompts = [formatter(prompt) for prompt in img_prompts]
|
||||
aspect_ratio_images = [
|
||||
[stop_sign, cherry_blossom],
|
||||
# Images with different sizes and aspect-ratios
|
||||
[
|
||||
rescale_image_size(stop_sign, 0.1),
|
||||
stop_sign,
|
||||
],
|
||||
[
|
||||
stop_sign,
|
||||
rescale_image_size(stop_sign, 0.25),
|
||||
cherry_blossom.resize((183, 488)),
|
||||
cherry_blossom.resize((488, 183)),
|
||||
],
|
||||
cherry_blossom,
|
||||
]
|
||||
|
||||
return [
|
||||
PromptWithMultiModalInput(
|
||||
prompts=formatted_prompts,
|
||||
image_data=aspect_ratio_images,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def multi_video_multi_aspect_ratio_inputs(
|
||||
formatter: Callable[[str], str], num_frames: int = 16
|
||||
):
|
||||
"""Builds inputs for multi-video (varied sizes/aspect ratio) testing.
|
||||
|
||||
Args:
|
||||
formatter: model-specific prompt formatter.
|
||||
"""
|
||||
video = sample_frames_from_video(VIDEO_ASSETS[0].np_ndarrays, num_frames)
|
||||
# Apply the selected formatter to the base prompts
|
||||
video_prompts = [
|
||||
"<video><video>\nDescribe 2 videos.",
|
||||
"<video><video>\nDescribe 2 videos.",
|
||||
"<video><video><video><video>\nDescribe 4 videos.",
|
||||
"<video>\nWhy is this video funny?",
|
||||
]
|
||||
formatted_prompts = [formatter(prompt) for prompt in video_prompts]
|
||||
aspect_ratio_videos = [
|
||||
[video, video],
|
||||
# Videos with different sizes and aspect-ratios
|
||||
[
|
||||
rescale_video_size(video, 0.1),
|
||||
video,
|
||||
],
|
||||
[
|
||||
video,
|
||||
rescale_video_size(video, 0.25),
|
||||
resize_video(video, (183, 488)),
|
||||
resize_video(video, (488, 183)),
|
||||
],
|
||||
video,
|
||||
]
|
||||
|
||||
return [
|
||||
PromptWithMultiModalInput(
|
||||
prompts=formatted_prompts,
|
||||
video_data=aspect_ratio_videos,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def different_patch_input_cases_internvl():
|
||||
images = [asset.pil_image.resize((896, 896)) for asset in IMAGE_ASSETS]
|
||||
formatter = (
|
||||
lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>Assistant\n" # noqa: E501
|
||||
)
|
||||
single_img_prompts = [
|
||||
"<image>\nWhat's the content in the center of the image?",
|
||||
"<image>\nWhat is the season?",
|
||||
]
|
||||
multi_img_prompts = [
|
||||
"Image-1: <image>\nImage-2: <image>\nDescribe the two images in detail.\n", # noqa: E501
|
||||
]
|
||||
formatted_sprompts = [formatter(prompt) for prompt in single_img_prompts]
|
||||
formatted_mprompts = [formatter(prompt) for prompt in multi_img_prompts]
|
||||
|
||||
wrapped_sf = ImageSizeWrapper(type=SizeType.SIZE_FACTOR, data=[0.5, 1.0])
|
||||
return [
|
||||
build_single_image_inputs(images, formatted_sprompts, wrapped_sf),
|
||||
build_multi_image_inputs([images], formatted_mprompts, wrapped_sf),
|
||||
]
|
||||
|
||||
|
||||
def windows_attention_image_qwen2_5_vl():
|
||||
# image from regression issue: https://github.com/vllm-project/vllm/issues/15122 # noqa: E501
|
||||
image = ImageAsset("hato").pil_image
|
||||
|
||||
question = "Describe the image."
|
||||
img_prompt = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
prompt = (
|
||||
f"<|im_start|>User\n{img_prompt}{question}<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
wrapped_sf = ImageSizeWrapper(type=SizeType.SIZE_FACTOR, data=[0.5])
|
||||
return build_single_image_inputs([image], [prompt], wrapped_sf)
|
||||
|
||||
|
||||
def video_with_metadata_glm4_1v():
|
||||
video_array = VIDEO_ASSETS[0].np_ndarrays
|
||||
metadata = VIDEO_ASSETS[0].metadata
|
||||
question = "Describe the video."
|
||||
video_prompt = "<|begin_of_video|><|video|><|end_of_video|>"
|
||||
formatted_prompt = f"[gMASK]<|user|>\n{video_prompt}{question}<|assistant|>\n"
|
||||
|
||||
scales = [0.1, 0.2, 0.25]
|
||||
video_input = [
|
||||
[(rescale_video_size(video_array, scale), metadata)] for scale in scales
|
||||
]
|
||||
prompts = [formatted_prompt] * len(video_input)
|
||||
|
||||
return [
|
||||
PromptWithMultiModalInput(
|
||||
prompts=prompts,
|
||||
video_data=video_input,
|
||||
)
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,190 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Entrypoints for wrapping the core run_test implementation for specific test
|
||||
types / modalities.
|
||||
"""
|
||||
|
||||
from pathlib import PosixPath
|
||||
|
||||
from .....conftest import (
|
||||
AudioTestAssets,
|
||||
HfRunner,
|
||||
ImageTestAssets,
|
||||
VideoTestAssets,
|
||||
VllmRunner,
|
||||
)
|
||||
from . import builders, core
|
||||
from .types import ExpandableVLMTestArgs, VLMTestInfo
|
||||
|
||||
|
||||
####### Entrypoints for running different test types
|
||||
def run_single_image_test(
|
||||
*,
|
||||
tmp_path: PosixPath,
|
||||
model_test_info: VLMTestInfo,
|
||||
test_case: ExpandableVLMTestArgs,
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
image_assets: ImageTestAssets,
|
||||
):
|
||||
assert test_case.size_wrapper is not None
|
||||
inputs = builders.build_single_image_inputs_from_test_info(
|
||||
model_test_info, image_assets, test_case.size_wrapper, tmp_path
|
||||
)
|
||||
|
||||
core.run_test(
|
||||
hf_runner=hf_runner,
|
||||
vllm_runner=vllm_runner,
|
||||
inputs=inputs,
|
||||
model=test_case.model,
|
||||
dtype=test_case.dtype,
|
||||
max_tokens=test_case.max_tokens,
|
||||
num_logprobs=test_case.num_logprobs,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
distributed_executor_backend=test_case.distributed_executor_backend,
|
||||
**model_test_info.get_non_parametrized_runner_kwargs(),
|
||||
)
|
||||
|
||||
|
||||
def run_multi_image_test(
|
||||
*,
|
||||
tmp_path: PosixPath,
|
||||
model_test_info: VLMTestInfo,
|
||||
test_case: ExpandableVLMTestArgs,
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
image_assets: ImageTestAssets,
|
||||
):
|
||||
assert test_case.size_wrapper is not None
|
||||
inputs = builders.build_multi_image_inputs_from_test_info(
|
||||
model_test_info, image_assets, test_case.size_wrapper, tmp_path
|
||||
)
|
||||
|
||||
core.run_test(
|
||||
hf_runner=hf_runner,
|
||||
vllm_runner=vllm_runner,
|
||||
inputs=inputs,
|
||||
model=test_case.model,
|
||||
dtype=test_case.dtype,
|
||||
max_tokens=test_case.max_tokens,
|
||||
num_logprobs=test_case.num_logprobs,
|
||||
limit_mm_per_prompt={"image": len(image_assets)},
|
||||
distributed_executor_backend=test_case.distributed_executor_backend,
|
||||
**model_test_info.get_non_parametrized_runner_kwargs(),
|
||||
)
|
||||
|
||||
|
||||
def run_embedding_test(
|
||||
*,
|
||||
model_test_info: VLMTestInfo,
|
||||
test_case: ExpandableVLMTestArgs,
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
image_assets: ImageTestAssets,
|
||||
):
|
||||
assert test_case.size_wrapper is not None
|
||||
inputs, vllm_embeddings = builders.build_embedding_inputs_from_test_info(
|
||||
model_test_info, image_assets, test_case.size_wrapper
|
||||
)
|
||||
|
||||
core.run_test(
|
||||
hf_runner=hf_runner,
|
||||
vllm_runner=vllm_runner,
|
||||
inputs=inputs,
|
||||
model=test_case.model,
|
||||
dtype=test_case.dtype,
|
||||
max_tokens=test_case.max_tokens,
|
||||
num_logprobs=test_case.num_logprobs,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
vllm_embeddings=vllm_embeddings,
|
||||
distributed_executor_backend=test_case.distributed_executor_backend,
|
||||
**model_test_info.get_non_parametrized_runner_kwargs(),
|
||||
)
|
||||
|
||||
|
||||
def run_video_test(
|
||||
*,
|
||||
model_test_info: VLMTestInfo,
|
||||
test_case: ExpandableVLMTestArgs,
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
video_assets: VideoTestAssets,
|
||||
):
|
||||
assert test_case.size_wrapper is not None
|
||||
assert test_case.num_video_frames is not None
|
||||
inputs = builders.build_video_inputs_from_test_info(
|
||||
model_test_info,
|
||||
video_assets,
|
||||
test_case.size_wrapper,
|
||||
test_case.num_video_frames,
|
||||
test_case.needs_video_metadata,
|
||||
)
|
||||
|
||||
core.run_test(
|
||||
hf_runner=hf_runner,
|
||||
vllm_runner=vllm_runner,
|
||||
inputs=inputs,
|
||||
model=test_case.model,
|
||||
dtype=test_case.dtype,
|
||||
max_tokens=test_case.max_tokens,
|
||||
num_logprobs=test_case.num_logprobs,
|
||||
limit_mm_per_prompt={"video": len(video_assets)},
|
||||
distributed_executor_backend=test_case.distributed_executor_backend,
|
||||
**model_test_info.get_non_parametrized_runner_kwargs(),
|
||||
)
|
||||
|
||||
|
||||
def run_audio_test(
|
||||
*,
|
||||
model_test_info: VLMTestInfo,
|
||||
test_case: ExpandableVLMTestArgs,
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
audio_assets: AudioTestAssets,
|
||||
):
|
||||
inputs = builders.build_audio_inputs_from_test_info(model_test_info, audio_assets)
|
||||
|
||||
core.run_test(
|
||||
hf_runner=hf_runner,
|
||||
vllm_runner=vllm_runner,
|
||||
inputs=inputs,
|
||||
model=test_case.model,
|
||||
dtype=test_case.dtype,
|
||||
max_tokens=test_case.max_tokens,
|
||||
num_logprobs=test_case.num_logprobs,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
distributed_executor_backend=test_case.distributed_executor_backend,
|
||||
**model_test_info.get_non_parametrized_runner_kwargs(),
|
||||
)
|
||||
|
||||
|
||||
def run_custom_inputs_test(
|
||||
*,
|
||||
model_test_info: VLMTestInfo,
|
||||
test_case: ExpandableVLMTestArgs,
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
):
|
||||
# Custom test cases can provide inputs directly, but they need to
|
||||
# explicitly provided a CustomTestConfig, which wraps the inputs and
|
||||
# the limit_mm_per_prompt
|
||||
assert test_case.custom_test_opts is not None
|
||||
|
||||
inputs = test_case.custom_test_opts.inputs
|
||||
limit_mm_per_prompt = test_case.custom_test_opts.limit_mm_per_prompt
|
||||
# Inputs and limit_mm_per_prompt should all be set
|
||||
assert inputs is not None
|
||||
assert limit_mm_per_prompt is not None
|
||||
|
||||
core.run_test(
|
||||
hf_runner=hf_runner,
|
||||
vllm_runner=vllm_runner,
|
||||
inputs=inputs,
|
||||
model=test_case.model,
|
||||
dtype=test_case.dtype,
|
||||
max_tokens=test_case.max_tokens,
|
||||
num_logprobs=test_case.num_logprobs,
|
||||
limit_mm_per_prompt=limit_mm_per_prompt,
|
||||
distributed_executor_backend=test_case.distributed_executor_backend,
|
||||
**model_test_info.get_non_parametrized_runner_kwargs(),
|
||||
)
|
||||
@@ -0,0 +1,220 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Types for writing multimodal model tests."""
|
||||
|
||||
from collections.abc import Callable, Iterable
|
||||
from enum import Enum
|
||||
from pathlib import PosixPath
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
import torch
|
||||
from pytest import MarkDecorator
|
||||
from transformers import AutoModelForCausalLM
|
||||
from transformers.models.auto.auto_factory import _BaseAutoModelClass
|
||||
|
||||
from vllm.config.model import RunnerOption
|
||||
from vllm.logprobs import SampleLogprobs
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
|
||||
from .....conftest import (
|
||||
AUDIO_ASSETS,
|
||||
IMAGE_ASSETS,
|
||||
HfRunner,
|
||||
ImageAsset,
|
||||
ImageTestAssets,
|
||||
PromptAudioInput,
|
||||
PromptImageInput,
|
||||
PromptVideoInput,
|
||||
)
|
||||
from ....utils import check_logprobs_close
|
||||
|
||||
# meta image tag; will be replaced by the appropriate tag for the model
|
||||
TEST_IMG_PLACEHOLDER = "<vlm_image>"
|
||||
TEST_VIDEO_PLACEHOLDER = "<vlm_video>"
|
||||
TEST_AUDIO_PLACEHOLDER = "<lmm_audio>"
|
||||
|
||||
SINGLE_IMAGE_BASE_PROMPTS = IMAGE_ASSETS.prompts(
|
||||
{
|
||||
"stop_sign": f"{TEST_IMG_PLACEHOLDER}What's the content of the image?",
|
||||
"cherry_blossom": f"{TEST_IMG_PLACEHOLDER}What is the season?",
|
||||
}
|
||||
)
|
||||
SINGLE_AUDIO_BASE_PROMPT = AUDIO_ASSETS.prompts(
|
||||
{
|
||||
"mary_had_lamb": f"{TEST_AUDIO_PLACEHOLDER}Transcribe this audio into English.", # noqa: E501
|
||||
"winning_call": f"{TEST_AUDIO_PLACEHOLDER}What is happening in this audio clip?", # noqa: E501
|
||||
}
|
||||
)
|
||||
|
||||
MULTI_IMAGE_BASE_PROMPT = f"Image-1: {TEST_IMG_PLACEHOLDER}Image-2: {TEST_IMG_PLACEHOLDER}Describe the two images in detail.\n" # noqa: E501
|
||||
VIDEO_BASE_PROMPT = f"{TEST_VIDEO_PLACEHOLDER}Why is this video funny?"
|
||||
|
||||
|
||||
IMAGE_SIZE_FACTORS = [(1.0,), (1.0, 1.0, 1.0), (0.25, 0.5, 1.0)]
|
||||
EMBEDDING_SIZE_FACTORS = [(1.0,), (1.0, 1.0, 1.0)]
|
||||
RunnerOutput = tuple[list[int], str, SampleLogprobs | None]
|
||||
|
||||
|
||||
class PromptWithMultiModalInput(NamedTuple):
|
||||
"""Holds the multimodal input for a single test case."""
|
||||
|
||||
prompts: list[str]
|
||||
image_data: PromptImageInput | None = None
|
||||
video_data: PromptVideoInput | None = None
|
||||
audio_data: PromptAudioInput | None = None
|
||||
|
||||
|
||||
class VLMTestType(Enum):
|
||||
IMAGE = 1
|
||||
MULTI_IMAGE = 2
|
||||
EMBEDDING = 3
|
||||
VIDEO = 4
|
||||
AUDIO = 5
|
||||
CUSTOM_INPUTS = 6
|
||||
|
||||
|
||||
class SizeType(Enum):
|
||||
SIZE_FACTOR = 1
|
||||
FIXED_SIZE = 2
|
||||
|
||||
|
||||
class CustomTestOptions(NamedTuple):
|
||||
inputs: list[PromptWithMultiModalInput]
|
||||
limit_mm_per_prompt: dict[str, int]
|
||||
|
||||
|
||||
class ImageSizeWrapper(NamedTuple):
|
||||
type: SizeType
|
||||
# A size factor is a wrapper of 0+ floats,
|
||||
# while a fixed size contains an iterable of integer pairs
|
||||
data: Iterable[float] | Iterable[tuple[int, int]]
|
||||
|
||||
|
||||
class VLMTestInfo(NamedTuple):
|
||||
"""Holds the configuration for 1+ tests for one model architecture."""
|
||||
|
||||
models: list[str]
|
||||
test_type: VLMTestType | Iterable[VLMTestType]
|
||||
|
||||
# Should be None only if this is a CUSTOM_INPUTS test
|
||||
prompt_formatter: Callable[[str], str] | None = None
|
||||
img_idx_to_prompt: Callable[[int], str] = lambda idx: "<image>\n"
|
||||
video_idx_to_prompt: Callable[[int], str] = lambda idx: "<video>\n"
|
||||
audio_idx_to_prompt: Callable[[int], str] = lambda idx: "<audio>\n"
|
||||
|
||||
# Most models work on the single / multi-image prompts above, but in some
|
||||
# cases the log prob check fails, e.g., for paligemma. We allow passing
|
||||
# an override for the single image prompts / multi-image prompt for this
|
||||
# reason.
|
||||
single_image_prompts: Iterable[str] = SINGLE_IMAGE_BASE_PROMPTS
|
||||
multi_image_prompt: str = MULTI_IMAGE_BASE_PROMPT
|
||||
|
||||
# Function for converting ImageAssets to image embeddings;
|
||||
# We need to define this explicitly for embedding tests
|
||||
convert_assets_to_embeddings: (
|
||||
Callable[[ImageTestAssets], list[torch.Tensor]] | None
|
||||
) = None
|
||||
|
||||
# Exposed options for vLLM runner; we change these in a several tests,
|
||||
# but the defaults are derived from VllmRunner & the engine defaults
|
||||
# These settings are chosen to avoid OOMs when running in the CI
|
||||
enforce_eager: bool = True
|
||||
max_model_len: int = 1024
|
||||
max_num_seqs: int = 256
|
||||
runner: RunnerOption = "auto"
|
||||
tensor_parallel_size: int = 1
|
||||
vllm_runner_kwargs: dict[str, Any] | None = None
|
||||
|
||||
# Optional callable which gets a list of token IDs from the model tokenizer
|
||||
get_stop_token_ids: Callable[[TokenizerLike], list[int]] | None = None
|
||||
# Optional list of strings to stop generation, useful when stop tokens are
|
||||
# not special tokens in the tokenizer
|
||||
stop_str: list[str] | None = None
|
||||
|
||||
# Exposed options for HF runner
|
||||
hf_model_kwargs: dict[str, Any] | None = None
|
||||
hf_processor: Callable[[str], Any] | None = None
|
||||
# Indicates we should explicitly pass the EOS from the tokenizer
|
||||
use_tokenizer_eos: bool = False
|
||||
auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM
|
||||
patch_hf_runner: Callable[[HfRunner], HfRunner] | None = None
|
||||
|
||||
# Post processors that if defined, will run oun the outputs of the
|
||||
# vLLM and HF runner, respectively (useful for sanitization, etc).
|
||||
vllm_output_post_proc: Callable[[RunnerOutput, str], Any] | None = None
|
||||
hf_output_post_proc: Callable[[RunnerOutput, str], Any] | None = None
|
||||
|
||||
# Consumes the output of the callables above and checks if they're equal
|
||||
comparator: Callable[..., None] = check_logprobs_close
|
||||
|
||||
# Default expandable params per test; these defaults can be overridden in
|
||||
# instances of this object; the complete set of test cases for the model
|
||||
# is all combinations of .models + all fields below
|
||||
max_tokens: int = 128
|
||||
num_logprobs: int = 5
|
||||
dtype: str = "auto"
|
||||
distributed_executor_backend: str | None = None
|
||||
# Only expanded in video tests
|
||||
num_video_frames: int | tuple[int] = 16
|
||||
needs_video_metadata: bool = False
|
||||
|
||||
# Fixed image sizes / image size factors; most tests use image_size_factors
|
||||
# The values provided for these two fields will be stacked and expanded
|
||||
# such that each model will consider each image size factor / image size
|
||||
# once per tests (much like concatenating and wrapping in one parametrize
|
||||
# call)
|
||||
image_size_factors: Iterable[Iterable[float]] = IMAGE_SIZE_FACTORS
|
||||
image_sizes: Iterable[Iterable[tuple[int, int]]] | None = None
|
||||
|
||||
# Hack for updating a prompt to take into a local path; currently only used
|
||||
# for Qwen-VL, which requires encoding the image path / url into the prompt
|
||||
# for HF runner
|
||||
prompt_path_encoder: (
|
||||
Callable[[PosixPath, str, list[ImageAsset] | ImageTestAssets], str] | None
|
||||
) = None # noqa: E501
|
||||
|
||||
# Allows configuring a test to run with custom inputs
|
||||
custom_test_opts: list[CustomTestOptions] | None = None
|
||||
|
||||
marks: list[MarkDecorator] | None = None
|
||||
|
||||
def get_non_parametrized_runner_kwargs(self):
|
||||
"""Returns a dictionary of expandable kwargs for items that are used
|
||||
in all test types, which are NOT used when creating the parametrized
|
||||
test cases.
|
||||
"""
|
||||
return {
|
||||
"enforce_eager": self.enforce_eager,
|
||||
"max_model_len": self.max_model_len,
|
||||
"max_num_seqs": self.max_num_seqs,
|
||||
"runner": self.runner,
|
||||
"tensor_parallel_size": self.tensor_parallel_size,
|
||||
"vllm_runner_kwargs": self.vllm_runner_kwargs,
|
||||
"hf_output_post_proc": self.hf_output_post_proc,
|
||||
"vllm_output_post_proc": self.vllm_output_post_proc,
|
||||
"auto_cls": self.auto_cls,
|
||||
"use_tokenizer_eos": self.use_tokenizer_eos,
|
||||
"comparator": self.comparator,
|
||||
"get_stop_token_ids": self.get_stop_token_ids,
|
||||
"hf_model_kwargs": self.hf_model_kwargs,
|
||||
"hf_processor": self.hf_processor,
|
||||
"stop_str": self.stop_str,
|
||||
"patch_hf_runner": self.patch_hf_runner,
|
||||
}
|
||||
|
||||
|
||||
class ExpandableVLMTestArgs(NamedTuple):
|
||||
"""The expanded kwargs which correspond to a single test case."""
|
||||
|
||||
model: str
|
||||
max_tokens: int
|
||||
num_logprobs: int
|
||||
dtype: str
|
||||
distributed_executor_backend: str | None
|
||||
# Sizes are used for everything except for custom input tests
|
||||
size_wrapper: ImageSizeWrapper | None = None
|
||||
# Video only
|
||||
num_video_frames: int | None = None
|
||||
needs_video_metadata: bool = False
|
||||
# Custom inputs only
|
||||
custom_test_opts: CustomTestOptions | None = None
|
||||
@@ -0,0 +1,18 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Pytest configuration for vLLM pooling tests."""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def siglip_attention_config():
|
||||
"""Return attention config for SigLIP tests on ROCm.
|
||||
|
||||
On ROCm, SigLIP tests require FLEX_ATTENTION backend.
|
||||
"""
|
||||
if current_platform.is_rocm():
|
||||
return {"backend": "FLEX_ATTENTION"}
|
||||
return None
|
||||
@@ -0,0 +1,143 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import CLIPModel
|
||||
|
||||
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
|
||||
from ...utils import check_embeddings_close
|
||||
|
||||
HF_TEXT_PROMPTS = [
|
||||
"a photo of a stop sign",
|
||||
"a photo of a cherry blossom",
|
||||
]
|
||||
|
||||
HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
|
||||
{
|
||||
"stop_sign": "",
|
||||
"cherry_blossom": "",
|
||||
}
|
||||
)
|
||||
|
||||
MODELS = ["openai/clip-vit-base-patch32"]
|
||||
|
||||
|
||||
def _run_test(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
input_texts: list[str],
|
||||
input_images: PromptImageInput,
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
# NOTE: take care of the order. run vLLM first, and then run HF.
|
||||
# vLLM needs a fresh new process without cuda initialization.
|
||||
# if we run HF first, the cuda initialization will be done and it
|
||||
# will hurt multiprocessing backend with fork method (the default method).
|
||||
with vllm_runner(
|
||||
model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.embed(input_texts, images=input_images)
|
||||
|
||||
with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model:
|
||||
all_inputs = hf_model.get_inputs(input_texts, images=input_images)
|
||||
|
||||
all_outputs = []
|
||||
for inputs in all_inputs:
|
||||
inputs = hf_model.wrap_device(inputs)
|
||||
|
||||
if "pixel_values" in inputs:
|
||||
pooled_output = hf_model.model.get_image_features(
|
||||
pixel_values=inputs.pixel_values,
|
||||
)
|
||||
else:
|
||||
pooled_output = hf_model.model.get_text_features(
|
||||
input_ids=inputs.input_ids,
|
||||
attention_mask=inputs.attention_mask,
|
||||
)
|
||||
|
||||
if not isinstance(pooled_output, torch.Tensor):
|
||||
pooled_output = pooled_output.pooler_output
|
||||
pooled_output = pooled_output.squeeze(0)
|
||||
all_outputs.append(pooled_output.tolist())
|
||||
|
||||
hf_outputs = all_outputs
|
||||
|
||||
check_embeddings_close(
|
||||
embeddings_0_lst=hf_outputs,
|
||||
embeddings_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["float"])
|
||||
def test_models_text(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
image_assets,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
|
||||
input_texts = [text for text, _ in input_texts_images]
|
||||
input_images = [image for _, image in input_texts_images]
|
||||
|
||||
_run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
input_texts,
|
||||
input_images, # type: ignore
|
||||
model,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["float"])
|
||||
def test_models_image(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
image_assets,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
input_texts_images = [
|
||||
(text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
|
||||
]
|
||||
input_texts = [text for text, _ in input_texts_images]
|
||||
input_images = [image for _, image in input_texts_images]
|
||||
|
||||
_run_test(
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
input_texts,
|
||||
input_images,
|
||||
model,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["float"])
|
||||
def test_models_text_image_no_crash(
|
||||
vllm_runner,
|
||||
image_assets,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
texts = [HF_TEXT_PROMPTS[0]]
|
||||
images = [image_assets[0].pil_image]
|
||||
|
||||
with vllm_runner(
|
||||
model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77
|
||||
) as vllm_model:
|
||||
with pytest.raises(ValueError, match="not both"):
|
||||
vllm_model.embed(texts, images=images)
|
||||
|
||||
# Should still be able to run subsequent requests
|
||||
vllm_model.embed(texts)
|
||||
vllm_model.embed([""], images=images)
|
||||
@@ -0,0 +1,115 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for ColModernVBERT multimodal late-interaction model.
|
||||
|
||||
ColModernVBERT combines SigLIP vision encoder + ModernBERT text encoder
|
||||
with a pixel shuffle connector and ColBERT-style 128-dim per-token
|
||||
embeddings for visual document retrieval.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
|
||||
|
||||
MODEL_NAME = "ModernVBERT/colmodernvbert-merged"
|
||||
COLBERT_DIM = 128
|
||||
DTYPE = "half"
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Text-only tests
|
||||
# -----------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_colmodernvbert_text_token_embed(vllm_runner):
|
||||
"""Text query produces per-token embeddings with shape (seq_len, 128)."""
|
||||
with vllm_runner(
|
||||
MODEL_NAME,
|
||||
runner="pooling",
|
||||
dtype=DTYPE,
|
||||
enforce_eager=True,
|
||||
) as vllm_model:
|
||||
outputs = vllm_model.token_embed(["What is machine learning?"])
|
||||
|
||||
assert len(outputs) == 1
|
||||
emb = torch.tensor(outputs[0])
|
||||
assert emb.dim() == 2
|
||||
assert emb.shape[1] == COLBERT_DIM
|
||||
assert emb.shape[0] > 1
|
||||
|
||||
|
||||
def test_colmodernvbert_text_relevance_ordering(vllm_runner):
|
||||
"""Relevant documents score higher than irrelevant ones."""
|
||||
query = "What is machine learning?"
|
||||
documents = [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"The weather in Paris is mild in spring.",
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
MODEL_NAME,
|
||||
runner="pooling",
|
||||
dtype=DTYPE,
|
||||
enforce_eager=True,
|
||||
) as vllm_model:
|
||||
scores = vllm_model.score(query, documents)
|
||||
|
||||
assert len(scores) == 2
|
||||
assert scores[0] > scores[1], "ML doc should score higher than weather doc"
|
||||
|
||||
|
||||
def test_colmodernvbert_text_late_interaction(vllm_runner):
|
||||
"""MaxSim scoring via vLLM matches manual computation."""
|
||||
query = "What is the capital of France?"
|
||||
doc = "The capital of France is Paris."
|
||||
|
||||
with vllm_runner(
|
||||
MODEL_NAME,
|
||||
runner="pooling",
|
||||
dtype=DTYPE,
|
||||
enforce_eager=True,
|
||||
) as vllm_model:
|
||||
q_out = vllm_model.token_embed([query])
|
||||
d_out = vllm_model.token_embed([doc])
|
||||
|
||||
q_emb = torch.tensor(q_out[0])
|
||||
d_emb = torch.tensor(d_out[0])
|
||||
manual_score = compute_maxsim_score(q_emb, d_emb).item()
|
||||
|
||||
vllm_scores = vllm_model.score(query, doc)
|
||||
|
||||
assert len(vllm_scores) == 1
|
||||
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Image tests
|
||||
# -----------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_colmodernvbert_image_token_embed(vllm_runner, image_assets):
|
||||
"""Image input produces per-token embeddings including vision tokens."""
|
||||
with vllm_runner(
|
||||
MODEL_NAME,
|
||||
runner="pooling",
|
||||
dtype=DTYPE,
|
||||
enforce_eager=True,
|
||||
) as vllm_model:
|
||||
image = image_assets[0].pil_image
|
||||
inputs = vllm_model.get_inputs(
|
||||
[""],
|
||||
images=[image],
|
||||
)
|
||||
req_outputs = vllm_model.llm.encode(
|
||||
inputs,
|
||||
pooling_task="token_embed",
|
||||
)
|
||||
outputs = [req_output.outputs.data for req_output in req_outputs]
|
||||
|
||||
assert len(outputs) == 1
|
||||
emb = torch.tensor(outputs[0])
|
||||
assert emb.dim() == 2
|
||||
assert emb.shape[1] == COLBERT_DIM
|
||||
# Should have at least the image tokens (64 after pixel shuffle)
|
||||
assert emb.shape[0] >= 64
|
||||
@@ -0,0 +1,323 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for ColPali late interaction model for multi-modal retrieval.
|
||||
|
||||
ColPali is a multi-vector retrieval model based on PaliGemma backbone
|
||||
(SigLIP + Gemma) with ColBERT-style late interaction scoring (MaxSim).
|
||||
It produces per-token embeddings for both text and image inputs.
|
||||
"""
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
import pybase64 as base64
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm.entrypoints.chat_utils import (
|
||||
ChatCompletionContentPartImageParam,
|
||||
ChatCompletionContentPartTextParam,
|
||||
)
|
||||
from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam
|
||||
|
||||
from ....conftest import VllmRunner
|
||||
|
||||
MODELS = [
|
||||
"vidore/colpali-v1.3-hf",
|
||||
]
|
||||
|
||||
EMBED_DIMS = {
|
||||
"vidore/colpali-v1.3-hf": 128,
|
||||
}
|
||||
|
||||
TEXT_QUERIES = [
|
||||
"What is the capital of France?",
|
||||
"Describe the contents of the document.",
|
||||
]
|
||||
|
||||
TEXT_DOCUMENTS = [
|
||||
"The capital of France is Paris.",
|
||||
"This document contains important financial data.",
|
||||
]
|
||||
|
||||
DTYPE = "half"
|
||||
GPU_MEMORY_UTILIZATION = 0.7
|
||||
|
||||
|
||||
def _make_base64_image(
|
||||
width: int = 64, height: int = 64, color: tuple[int, int, int] = (255, 0, 0)
|
||||
) -> str:
|
||||
"""Create a small solid-color PNG image and return its base64 data URI."""
|
||||
img = Image.new("RGB", (width, height), color)
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
b64 = base64.b64encode(buf.getvalue()).decode()
|
||||
return f"data:image/png;base64,{b64}"
|
||||
|
||||
|
||||
def _make_image_mm_param(
|
||||
image_uri: str,
|
||||
text: str | None = None,
|
||||
) -> ScoreMultiModalParam:
|
||||
"""Build a ScoreMultiModalParam containing an image (and optional text)."""
|
||||
content: list = [
|
||||
ChatCompletionContentPartImageParam(
|
||||
type="image_url",
|
||||
image_url={"url": image_uri},
|
||||
),
|
||||
]
|
||||
if text is not None:
|
||||
content.append(
|
||||
ChatCompletionContentPartTextParam(type="text", text=text),
|
||||
)
|
||||
return ScoreMultiModalParam(content=content)
|
||||
|
||||
|
||||
def _run_token_embed_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Verify per-token embedding shape and L2 normalization."""
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
|
||||
|
||||
assert len(outputs) == 1
|
||||
emb = torch.tensor(outputs[0])
|
||||
# Token embeddings should be 2D: [num_tokens, embed_dim]
|
||||
assert emb.dim() == 2
|
||||
assert emb.shape[1] == EMBED_DIMS[model]
|
||||
assert emb.shape[0] > 1
|
||||
|
||||
# Verify L2 normalization
|
||||
norms = torch.norm(emb, p=2, dim=-1)
|
||||
torch.testing.assert_close(
|
||||
norms,
|
||||
torch.ones_like(norms),
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
def _run_late_interaction_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Verify MaxSim scoring matches manual computation."""
|
||||
from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
|
||||
d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
|
||||
|
||||
q_emb = torch.tensor(q_outputs[0])
|
||||
d_emb = torch.tensor(d_outputs[0])
|
||||
|
||||
manual_score = compute_maxsim_score(q_emb, d_emb).item()
|
||||
|
||||
vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
|
||||
|
||||
assert len(vllm_scores) == 1
|
||||
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
|
||||
|
||||
|
||||
def _run_relevance_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Verify that relevant documents score higher than irrelevant ones."""
|
||||
query = "What is machine learning?"
|
||||
documents = [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"The weather forecast shows rain tomorrow.",
|
||||
"Deep learning uses neural networks for complex tasks.",
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
scores = vllm_model.score(query, documents)
|
||||
|
||||
assert len(scores) == 3
|
||||
assert scores[0] > scores[1], "ML doc should score higher than weather doc"
|
||||
assert scores[2] > scores[1], "DL doc should score higher than weather doc"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colpali_token_embed(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_token_embed_test(vllm_runner, model, dtype=dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colpali_late_interaction_scoring(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_late_interaction_test(vllm_runner, model, dtype=dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colpali_relevance_ordering(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_relevance_test(vllm_runner, model, dtype=dtype)
|
||||
|
||||
|
||||
# ── Multimodal scoring tests ────────────────────────────────
|
||||
|
||||
|
||||
def _run_multimodal_text_query_image_docs_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Score a text query against image documents via the multimodal path."""
|
||||
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
|
||||
blue_image = _make_base64_image(64, 64, color=(0, 0, 255))
|
||||
|
||||
query = "Describe the red object"
|
||||
image_docs = [
|
||||
_make_image_mm_param(red_image),
|
||||
_make_image_mm_param(blue_image),
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
scores = vllm_model.llm.score(query, image_docs)
|
||||
|
||||
assert len(scores) == 2
|
||||
for s in scores:
|
||||
assert isinstance(s.outputs.score, float)
|
||||
|
||||
|
||||
def _run_multimodal_mixed_docs_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Score a text query against a mix of text and image documents."""
|
||||
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
|
||||
|
||||
query = "What is the capital of France?"
|
||||
documents: list = [
|
||||
"The capital of France is Paris.",
|
||||
_make_image_mm_param(red_image),
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
scores = vllm_model.llm.score(query, documents)
|
||||
|
||||
assert len(scores) == 2
|
||||
for s in scores:
|
||||
assert isinstance(s.outputs.score, float)
|
||||
# Text document about France should score higher than a random image
|
||||
assert scores[0].outputs.score > scores[1].outputs.score
|
||||
|
||||
|
||||
def _run_multimodal_image_query_text_docs_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Score an image query against text documents."""
|
||||
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
|
||||
image_query = _make_image_mm_param(red_image, text="red color")
|
||||
|
||||
documents = [
|
||||
"A bright red sports car.",
|
||||
"The weather forecast shows rain tomorrow.",
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
scores = vllm_model.llm.score(image_query, documents)
|
||||
|
||||
assert len(scores) == 2
|
||||
for s in scores:
|
||||
assert isinstance(s.outputs.score, float)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colpali_multimodal_text_query_image_docs(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_multimodal_text_query_image_docs_test(vllm_runner, model, dtype=dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colpali_multimodal_mixed_docs(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_multimodal_mixed_docs_test(vllm_runner, model, dtype=dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colpali_multimodal_image_query_text_docs(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_multimodal_image_query_text_docs_test(vllm_runner, model, dtype=dtype)
|
||||
@@ -0,0 +1,352 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for ColQwen3 late interaction model for multi-modal retrieval.
|
||||
|
||||
ColQwen3 is a multi-vector retrieval model based on Qwen3-VL backbone with
|
||||
ColBERT-style late interaction scoring (MaxSim). It produces per-token
|
||||
embeddings for both text and image inputs.
|
||||
"""
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
import pybase64 as base64
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm.entrypoints.chat_utils import (
|
||||
ChatCompletionContentPartImageParam,
|
||||
ChatCompletionContentPartTextParam,
|
||||
)
|
||||
from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam
|
||||
|
||||
from ....conftest import VllmRunner
|
||||
|
||||
pytestmark = pytest.mark.skip(
|
||||
reason="ColQwen3 model's weight tying is incompatible with "
|
||||
"transformers v5 (missing all_tied_weights_keys)"
|
||||
)
|
||||
|
||||
MODELS = [
|
||||
"TomoroAI/tomoro-colqwen3-embed-4b",
|
||||
"OpenSearch-AI/Ops-Colqwen3-4B",
|
||||
"nvidia/nemotron-colembed-vl-4b-v2",
|
||||
]
|
||||
|
||||
EMBED_DIMS = {
|
||||
"TomoroAI/tomoro-colqwen3-embed-4b": 320,
|
||||
"OpenSearch-AI/Ops-Colqwen3-4B": 2560,
|
||||
"nvidia/nemotron-colembed-vl-4b-v2": 2560,
|
||||
}
|
||||
|
||||
TEXT_QUERIES = [
|
||||
"What is the capital of France?",
|
||||
"Describe the contents of the document.",
|
||||
]
|
||||
|
||||
TEXT_DOCUMENTS = [
|
||||
"The capital of France is Paris.",
|
||||
"This document contains important financial data.",
|
||||
]
|
||||
|
||||
DTYPE = "half"
|
||||
GPU_MEMORY_UTILIZATION = 0.7
|
||||
|
||||
|
||||
def _make_base64_image(
|
||||
width: int = 64, height: int = 64, color: tuple[int, int, int] = (255, 0, 0)
|
||||
) -> str:
|
||||
"""Create a small solid-color PNG image and return its base64 data URI."""
|
||||
img = Image.new("RGB", (width, height), color)
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
b64 = base64.b64encode(buf.getvalue()).decode()
|
||||
return f"data:image/png;base64,{b64}"
|
||||
|
||||
|
||||
def _make_image_mm_param(
|
||||
image_uri: str,
|
||||
text: str | None = None,
|
||||
) -> ScoreMultiModalParam:
|
||||
"""Build a ScoreMultiModalParam containing an image (and optional text)."""
|
||||
content: list = [
|
||||
ChatCompletionContentPartImageParam(
|
||||
type="image_url",
|
||||
image_url={"url": image_uri},
|
||||
),
|
||||
]
|
||||
if text is not None:
|
||||
content.append(
|
||||
ChatCompletionContentPartTextParam(type="text", text=text),
|
||||
)
|
||||
return ScoreMultiModalParam(content=content)
|
||||
|
||||
|
||||
def _make_text_mm_param(text: str) -> ScoreMultiModalParam:
|
||||
"""Build a ScoreMultiModalParam containing only text."""
|
||||
return ScoreMultiModalParam(
|
||||
content=[ChatCompletionContentPartTextParam(type="text", text=text)],
|
||||
)
|
||||
|
||||
|
||||
def _run_token_embed_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Verify per-token embedding shape and L2 normalization."""
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
|
||||
|
||||
assert len(outputs) == 1
|
||||
emb = torch.tensor(outputs[0])
|
||||
# Token embeddings should be 2D: [num_tokens, embed_dim]
|
||||
assert emb.dim() == 2
|
||||
assert emb.shape[1] == EMBED_DIMS[model]
|
||||
assert emb.shape[0] > 1
|
||||
|
||||
# Verify L2 normalization
|
||||
norms = torch.norm(emb, p=2, dim=-1)
|
||||
torch.testing.assert_close(
|
||||
norms,
|
||||
torch.ones_like(norms),
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
def _run_late_interaction_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Verify MaxSim scoring matches manual computation."""
|
||||
from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
|
||||
d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
|
||||
|
||||
q_emb = torch.tensor(q_outputs[0])
|
||||
d_emb = torch.tensor(d_outputs[0])
|
||||
|
||||
manual_score = compute_maxsim_score(q_emb, d_emb).item()
|
||||
|
||||
vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
|
||||
|
||||
assert len(vllm_scores) == 1
|
||||
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
|
||||
|
||||
|
||||
def _run_relevance_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Verify that relevant documents score higher than irrelevant ones."""
|
||||
query = "What is machine learning?"
|
||||
documents = [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"The weather forecast shows rain tomorrow.",
|
||||
"Deep learning uses neural networks for complex tasks.",
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
scores = vllm_model.score(query, documents)
|
||||
|
||||
assert len(scores) == 3
|
||||
assert scores[0] > scores[1], "ML doc should score higher than weather doc"
|
||||
assert scores[2] > scores[1], "DL doc should score higher than weather doc"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colqwen3_token_embed(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_token_embed_test(vllm_runner, model, dtype=dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colqwen3_late_interaction_scoring(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_late_interaction_test(vllm_runner, model, dtype=dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colqwen3_relevance_ordering(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_relevance_test(vllm_runner, model, dtype=dtype)
|
||||
|
||||
|
||||
# ── Multimodal scoring tests ────────────────────────────────
|
||||
|
||||
|
||||
def _run_multimodal_text_query_image_docs_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Score a text query against image documents via the multimodal path.
|
||||
|
||||
Verifies that score_data_to_prompts correctly handles image content
|
||||
and produces valid MaxSim scores.
|
||||
"""
|
||||
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
|
||||
blue_image = _make_base64_image(64, 64, color=(0, 0, 255))
|
||||
|
||||
query = "Describe the red object"
|
||||
image_docs = [
|
||||
_make_image_mm_param(red_image),
|
||||
_make_image_mm_param(blue_image),
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
scores = vllm_model.llm.score(query, image_docs)
|
||||
|
||||
assert len(scores) == 2
|
||||
for s in scores:
|
||||
assert isinstance(s.outputs.score, float)
|
||||
|
||||
|
||||
def _run_multimodal_mixed_docs_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Score a text query against a mix of text and image documents.
|
||||
|
||||
Ensures the late-interaction path handles heterogeneous document
|
||||
types (plain strings alongside ScoreMultiModalParam images) in
|
||||
a single call.
|
||||
"""
|
||||
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
|
||||
|
||||
query = "What is the capital of France?"
|
||||
documents: list = [
|
||||
"The capital of France is Paris.",
|
||||
_make_image_mm_param(red_image),
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
scores = vllm_model.llm.score(query, documents)
|
||||
|
||||
assert len(scores) == 2
|
||||
for s in scores:
|
||||
assert isinstance(s.outputs.score, float)
|
||||
# Text document about France should score higher than a random image
|
||||
assert scores[0].outputs.score > scores[1].outputs.score
|
||||
|
||||
|
||||
def _run_multimodal_image_query_text_docs_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
*,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
"""Score an image query against text documents.
|
||||
|
||||
Verifies the reverse direction: multimodal query with text-only
|
||||
documents through the late-interaction scoring path.
|
||||
"""
|
||||
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
|
||||
image_query = _make_image_mm_param(red_image, text="red color")
|
||||
|
||||
documents = [
|
||||
"A bright red sports car.",
|
||||
"The weather forecast shows rain tomorrow.",
|
||||
]
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
dtype=dtype,
|
||||
max_model_len=4096,
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
|
||||
) as vllm_model:
|
||||
scores = vllm_model.llm.score(image_query, documents)
|
||||
|
||||
assert len(scores) == 2
|
||||
for s in scores:
|
||||
assert isinstance(s.outputs.score, float)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colqwen3_multimodal_text_query_image_docs(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_multimodal_text_query_image_docs_test(vllm_runner, model, dtype=dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colqwen3_multimodal_mixed_docs(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_multimodal_mixed_docs_test(vllm_runner, model, dtype=dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", [DTYPE])
|
||||
def test_colqwen3_multimodal_image_query_text_docs(
|
||||
vllm_runner,
|
||||
model: str,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
_run_multimodal_image_query_text_docs_test(vllm_runner, model, dtype=dtype)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user