@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / afb4429b4f
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 12:15 UTC

[CI/Build] Reorganize models tests (#17459)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed May 1, 2025 at 14:03 UTC afb4429b4f13e744b1630b6c5a09156e5b1ececc
65 files changed +317 -324
.buildkite/test-pipeline.yaml
+38 -48
@@ -390,12 +390,15 @@ steps:
390 commands:
391 - pytest -v -s benchmarks/
392
393 -- label: Quantization Test # 33min
393 +- label: Quantization Test
394 source_file_dependencies:
395 - csrc/
396 - vllm/model_executor/layers/quantization
397 - tests/quantization
398 - command: VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization
398 + - tests/models/quantization
399 + commands:
400 + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization
401 + - pytest -v -s models/quantization
402
403 - label: LM Eval Small Models # 53min
404 working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
@@ -441,82 +444,70 @@ steps:
444 commands:
445 - pytest -v -s models/test_transformers.py
446 - pytest -v -s models/test_registry.py
447 + - pytest -v -s models/test_utils.py
448 + - pytest -v -s models/test_vision.py
449 # V1 Test: https://github.com/vllm-project/vllm/issues/14531
450 - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4 and not plamo2'
451 - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'llama4'
452 - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'plamo2'
453
449 -- label: Language Models Test (Standard) # 32min
454 +- label: Language Models Test (Standard)
455 #mirror_hardwares: [amd]
456 source_file_dependencies:
457 - vllm/
453 - - tests/models/decoder_only/language
454 - - tests/models/embedding/language
455 - - tests/models/encoder_decoder/language
458 + - tests/models/language
459 commands:
460 # Install causal-conv1d for plamo2 models here, as it is not compatible with pip-compile.
461 - pip install 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.0.post8'
459 - - pytest -v -s models/decoder_only/language -m 'core_model or quant_model'
460 - - pytest -v -s models/embedding/language -m core_model
462 + - pytest -v -s models/language -m core_model
463
462 -- label: Language Models Test (Extended) # 1h10min
464 +- label: Language Models Test (Extended)
465 optional: true
466 source_file_dependencies:
467 - vllm/
466 - - tests/models/decoder_only/language
467 - - tests/models/embedding/language
468 - - tests/models/encoder_decoder/language
468 + - tests/models/language
469 commands:
470 # Install causal-conv1d for plamo2 models here, as it is not compatible with pip-compile.
471 - - pip install causal-conv1d
472 - - pytest -v -s models/decoder_only/language -m 'not core_model and not quant_model'
473 - - pytest -v -s models/embedding/language -m 'not core_model'
471 + - pip install 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.0.post8'
472 + - pytest -v -s models/language -m 'not core_model'
473
475 -- label: Multi-Modal Models Test (Standard) # 40min
474 +- label: Multi-Modal Models Test (Standard)
475 #mirror_hardwares: [amd]
476 source_file_dependencies:
477 - vllm/
479 - - tests/models/decoder_only/audio_language
480 - - tests/models/decoder_only/vision_language
481 - - tests/models/embedding/vision_language
482 - - tests/models/encoder_decoder/audio_language
483 - - tests/models/encoder_decoder/vision_language
478 + - tests/models/multimodal
479 + commands:
480 + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
481 + - pytest -v -s models/multimodal/processing
482 + - pytest -v -s --ignore models/multimodal/generation/test_whisper.py models/multimodal -m core_model
483 + - cd .. && pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
484 +
485 +- label: Multi-Modal Models Test (Extended) 1
486 + optional: true
487 + source_file_dependencies:
488 + - vllm/
489 + - tests/models/multimodal
490 commands:
491 - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
486 - - pytest -v -s models/multimodal
487 - - pytest -v -s models/decoder_only/audio_language -m 'core_model or quant_model'
488 - - pytest -v -s models/decoder_only/vision_language -m 'core_model or quant_model'
489 - - pytest -v -s models/embedding/vision_language -m core_model
490 - - pytest -v -s models/encoder_decoder/audio_language -m core_model
491 - - pytest -v -s models/encoder_decoder/language -m core_model
492 - - pytest -v -s models/encoder_decoder/vision_language -m core_model
493 - - pytest -v -s models/decoder_only/vision_language/test_interleaved.py
494 -
495 -- label: Multi-Modal Models Test (Extended) 1 # 48m
492 + - pytest -v -s --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/processing models/multimodal -m 'not core_model'
493 +
494 +- label: Multi-Modal Models Test (Extended) 2
495 optional: true
496 source_file_dependencies:
497 - vllm/
499 - - tests/models/decoder_only/audio_language
500 - - tests/models/decoder_only/vision_language
501 - - tests/models/embedding/vision_language
502 - - tests/models/encoder_decoder/vision_language
498 + - tests/models/multimodal
499 commands:
500 - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
505 - - pytest -v -s models/decoder_only/audio_language -m 'not core_model and not quant_model'
506 - - pytest -v -s models/decoder_only/vision_language/test_models.py -m 'split(group=0) and not core_model and not quant_model'
507 - - pytest -v -s --ignore models/decoder_only/vision_language/test_models.py models/decoder_only/vision_language -m 'not core_model and not quant_model'
508 - - pytest -v -s models/embedding/vision_language -m 'not core_model'
509 - - pytest -v -s models/encoder_decoder/language -m 'not core_model'
510 - - pytest -v -s models/encoder_decoder/vision_language -m 'not core_model'
511 -
512 -- label: Multi-Modal Models Test (Extended) 2 # 38m
501 + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
502 +
503 +- label: Multi-Modal Models Test (Extended) 3
504 optional: true
505 source_file_dependencies:
506 - vllm/
516 - - tests/models/decoder_only/vision_language
507 + - tests/models/multimodal
508 commands:
509 - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
519 - - pytest -v -s models/decoder_only/vision_language/test_models.py -m 'split(group=1) and not core_model and not quant_model'
510 + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
511
512 # This test is used only in PR development phase to test individual models and should never run on main
513 - label: Custom Models Test
@@ -586,9 +577,8 @@ steps:
577 - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
578 # Avoid importing model tests that cause CUDA reinitialization error
579 - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
589 - - pytest models/encoder_decoder/language/test_bart.py -v -s -m 'distributed(num_gpus=2)'
590 - - pytest models/encoder_decoder/vision_language/test_broadcast.py -v -s -m 'distributed(num_gpus=2)'
591 - - pytest models/decoder_only/vision_language/test_models.py -v -s -m 'distributed(num_gpus=2)'
580 + - pytest models/language -v -s -m 'distributed(num_gpus=2)'
581 + - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)'
582 # test sequence parallel
583 - pytest -v -s distributed/test_sequence_parallel.py
584 # this test fails consistently.
pyproject.toml
-1
@@ -158,7 +158,6 @@ markers = [
158 "skip_global_cleanup",
159 "core_model: enable this model test in each PR instead of only nightly",
160 "cpu_model: enable this model test in CPU tests",
161 - "quant_model: run this model test under Quantized category",
161 "split: run this test as part of a split",
162 "distributed: run this test only in distributed GPU tests",
163 "skip_v1: do not run this test with v1",
tests/entrypoints/openai/test_embedding.py
+6 -6
@@ -11,7 +11,7 @@ import requests
11 from vllm.entrypoints.openai.protocol import EmbeddingResponse
12 from vllm.transformers_utils.tokenizer import get_tokenizer
13
14 -from ...models.embedding.utils import correctness_test
14 +from ...models.utils import run_embedding_correctness_test
15 from ...utils import RemoteOpenAIServer
16
17 MODEL_NAME = "intfloat/multilingual-e5-small"
@@ -76,7 +76,7 @@ async def test_single_embedding(hf_model, client: openai.AsyncOpenAI,
76 assert embeddings.usage.total_tokens == 11
77
78 vllm_outputs = [d.embedding for d in embeddings.data]
79 - correctness_test(hf_model, input_texts, vllm_outputs)
79 + run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
80
81 # test using token IDs
82 input_tokens = [1, 1, 1, 1, 1]
@@ -121,7 +121,7 @@ async def test_batch_embedding(hf_model, client: openai.AsyncOpenAI,
121 assert embeddings.usage.total_tokens == 33
122
123 vllm_outputs = [d.embedding for d in embeddings.data]
124 - correctness_test(hf_model, input_texts, vllm_outputs)
124 + run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
125
126 # test list[list[int]]
127 input_tokens = [[4, 5, 7, 9, 20], [15, 29, 499], [24, 24, 24, 24, 24],
@@ -208,7 +208,7 @@ async def test_batch_base64_embedding(hf_model, client: openai.AsyncOpenAI,
208 model=model_name,
209 encoding_format="float")
210 float_data = [d.embedding for d in responses_float.data]
211 - correctness_test(hf_model, input_texts, float_data)
211 + run_embedding_correctness_test(hf_model, input_texts, float_data)
212
213 responses_base64 = await client.embeddings.create(input=input_texts,
214 model=model_name,
@@ -219,13 +219,13 @@ async def test_batch_base64_embedding(hf_model, client: openai.AsyncOpenAI,
219 np.frombuffer(base64.b64decode(data.embedding),
220 dtype="float32").tolist())
221
222 - correctness_test(hf_model, input_texts, base64_data)
222 + run_embedding_correctness_test(hf_model, input_texts, base64_data)
223
224 # Default response is float32 decoded from base64 by OpenAI Client
225 responses_default = await client.embeddings.create(input=input_texts,
226 model=model_name)
227 default_data = [d.embedding for d in responses_default.data]
228 - correctness_test(hf_model, input_texts, default_data)
228 + run_embedding_correctness_test(hf_model, input_texts, default_data)
229
230
231 @pytest.mark.asyncio
tests/entrypoints/openai/test_embedding_dimensions.py
+3 -2
@@ -11,7 +11,7 @@ import pytest
11 from vllm.entrypoints.openai.protocol import EmbeddingResponse
12
13 from ...conftest import HfRunner
14 -from ...models.embedding.utils import EmbedModelInfo, correctness_test
14 +from ...models.utils import EmbedModelInfo, run_embedding_correctness_test
15 from ...utils import RemoteOpenAIServer
16
17 MODELS = [
@@ -95,7 +95,8 @@ async def test_matryoshka(model_info: EmbedModelInfo,
95 assert len(embeddings.data[0].embedding) == dimensions
96
97 vllm_outputs = [d.embedding for d in embeddings.data]
98 - correctness_test(hf_model, prompts, vllm_outputs, dimensions)
98 + run_embedding_correctness_test(hf_model, prompts, vllm_outputs,
99 + dimensions)
100
101 if model_info.is_matryoshka:
102 valid_dimensions: list[Optional[int]] = [None]
tests/models/embedding/utils.py deleted
-66
@@ -1,66 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -
3 -from collections.abc import Sequence
4 -from typing import NamedTuple, Optional
5 -
6 -import torch
7 -import torch.nn.functional as F
8 -
9 -
10 -def check_embeddings_close(
11 - *,
12 - embeddings_0_lst: Sequence[list[float]],
13 - embeddings_1_lst: Sequence[list[float]],
14 - name_0: str,
15 - name_1: str,
16 - tol: float = 1e-3,
17 -) -> None:
18 - assert len(embeddings_0_lst) == len(embeddings_1_lst)
19 -
20 - for prompt_idx, (embeddings_0, embeddings_1) in enumerate(
21 - zip(embeddings_0_lst, embeddings_1_lst)):
22 - assert len(embeddings_0) == len(embeddings_1), (
23 - f"Length mismatch: {len(embeddings_0)} vs. {len(embeddings_1)}")
24 -
25 - sim = F.cosine_similarity(torch.tensor(embeddings_0),
26 - torch.tensor(embeddings_1),
27 - dim=0)
28 -
29 - fail_msg = (f"Test{prompt_idx}:"
30 - f"\n{name_0}:\t{embeddings_0[:16]!r}"
31 - f"\n{name_1}:\t{embeddings_1[:16]!r}")
32 -
33 - assert sim >= 1 - tol, fail_msg
34 -
35 -
36 -def matryoshka_fy(tensor, dimensions):
37 - tensor = torch.tensor(tensor)
38 - tensor = tensor[..., :dimensions]
39 - tensor = F.normalize(tensor, p=2, dim=1)
40 - return tensor
41 -
42 -
43 -class EmbedModelInfo(NamedTuple):
44 - name: str
45 - is_matryoshka: bool
46 - matryoshka_dimensions: Optional[list[int]] = None
47 - architecture: str = ""
48 - enable_test: bool = True
49 -
50 -
51 -def correctness_test(hf_model,
52 - inputs,
53 - vllm_outputs: Sequence[list[float]],
54 - dimensions: Optional[int] = None):
55 -
56 - hf_outputs = hf_model.encode(inputs)
57 - if dimensions:
58 - hf_outputs = matryoshka_fy(hf_outputs, dimensions)
59 -
60 - check_embeddings_close(
61 - embeddings_0_lst=hf_outputs,
62 - embeddings_1_lst=vllm_outputs,
63 - name_0="hf",
64 - name_1="vllm",
65 - tol=1e-2,
66 - )
tests/models/embedding/vision_language/__init__.py
tests/models/encoder_decoder/__init__.py
tests/models/encoder_decoder/audio_language/__init__.py
tests/models/encoder_decoder/language/__init__.py
tests/models/encoder_decoder/vision_language/__init__.py
tests/models/encoder_decoder/vision_language/test_broadcast.py deleted
-37
@@ -1,37 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -
3 -import pytest
4 -
5 -from ....utils import multi_gpu_test
6 -
7 -
8 -@multi_gpu_test(num_gpus=2)
9 -@pytest.mark.parametrize("distributed_executor_backend", ["ray", "mp"])
10 -@pytest.mark.parametrize("model", [
11 - "meta-llama/Llama-3.2-11B-Vision-Instruct",
12 -])
13 -def test_models(hf_runner, vllm_runner, image_assets,
14 - distributed_executor_backend, model) -> None:
15 -
16 - dtype = "half"
17 - max_tokens = 5
18 - num_logprobs = 5
19 - tensor_parallel_size = 2
20 -
21 - if model.startswith("meta-llama/Llama-3.2-11B-Vision-Instruct"):
22 - from .test_mllama import models, run_test
23 - else:
24 - raise NotImplementedError(f"Unsupported model: {model}")
25 -
26 - run_test(
27 - hf_runner,
28 - vllm_runner,
29 - image_assets,
30 - model=models[0],
31 - size_factors=[0.25, 0.5, 1.0],
32 - dtype=dtype,
33 - max_tokens=max_tokens,
34 - num_logprobs=num_logprobs,
35 - tensor_parallel_size=tensor_parallel_size,
36 - distributed_executor_backend=distributed_executor_backend,
37 - )
tests/models/language/__init__.py renamed
tests/models/language/generation/__init__.py renamed
tests/models/language/generation/test_bart.py renamed
-4
@@ -1,8 +1,4 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -"""Compare the outputs of HF and vLLM for BART models using greedy sampling.
3 -
4 -Run `pytest tests/models/encoder_decoder/language/test_bart.py`.
5 -"""
2 from typing import Optional
3
4 import pytest
tests/models/language/generation/test_granite.py renamed
tests/models/language/generation/test_hybrid.py renamed
+10 -8
@@ -289,23 +289,25 @@ def test_multistep_correctness(
289 @multi_gpu_test(num_gpus=2)
290 @pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
291 @pytest.mark.parametrize("max_tokens", [64])
292 -def test_hybrid_distributed_produces_identical_generation(
292 +@pytest.mark.parametrize("num_logprobs", [5])
293 +def test_distributed_correctness(
294 vllm_runner,
295 example_prompts,
296 model: str,
297 max_tokens: int,
298 + num_logprobs: int,
299 ) -> None:
298 - with vllm_runner(model, tensor_parallel_size=2,
300 + with vllm_runner(model, tensor_parallel_size=1,
301 max_num_seqs=2) as vllm_model:
300 - vllm_outputs_tp_2 = vllm_model.generate_greedy(example_prompts,
301 - max_tokens)
302 + vllm_outputs_tp_1 = vllm_model.generate_greedy_logprobs(
303 + example_prompts, max_tokens, num_logprobs)
304
303 - with vllm_runner(model, tensor_parallel_size=1,
305 + with vllm_runner(model, tensor_parallel_size=2,
306 max_num_seqs=2) as vllm_model:
305 - vllm_outputs_tp_1 = vllm_model.generate_greedy(example_prompts,
306 - max_tokens)
307 + vllm_outputs_tp_2 = vllm_model.generate_greedy_logprobs(
308 + example_prompts, max_tokens, num_logprobs)
309
308 - check_outputs_equal(
310 + check_logprobs_close(
311 outputs_0_lst=vllm_outputs_tp_1,
312 outputs_1_lst=vllm_outputs_tp_2,
313 name_0="vllm_tp_1",
tests/models/language/generation/test_mistral.py renamed
tests/models/language/generation/test_models.py renamed
tests/models/language/generation/test_phimoe.py renamed
tests/models/language/pooling/__init__.py renamed
tests/models/language/pooling/test_cls_models.py renamed
tests/models/language/pooling/test_embedding.py renamed
+1 -1
@@ -8,7 +8,7 @@ import pytest
8 from vllm.config import PoolerConfig
9 from vllm.platforms import current_platform
10
11 -from ..utils import check_embeddings_close
11 +from ...utils import check_embeddings_close
12
13
14 @pytest.mark.parametrize(
tests/models/language/pooling/test_gritlm.py renamed
+95 -94
@@ -7,11 +7,10 @@ from array import array
7
8 import openai
9 import pytest
10 -import pytest_asyncio
10 from scipy.spatial.distance import cosine
11
13 -import vllm
14 -import vllm.config
12 +from vllm import LLM, SamplingParams
13 +from vllm.config import ModelConfig
14 from vllm.utils import STR_BACKEND_ENV_VAR
15
16 from ....utils import RemoteOpenAIServer
@@ -31,73 +30,45 @@ def _arr(arr):
30 return array("i", arr)
31
32
34 -def test_find_array(monkeypatch: pytest.MonkeyPatch):
35 - # GritLM embedding implementation is only supported by XFormers backend.
36 - with monkeypatch.context() as m:
37 - m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
38 -
39 - from vllm.model_executor.models.gritlm import GritLMPooler
40 -
41 - # Create an LLM object to get the model config.
42 - llm = vllm.LLM(MODEL_NAME, task="embed", max_model_len=MAX_MODEL_LEN)
43 - pooler = GritLMPooler(model_config=llm.llm_engine.model_config)
44 -
45 - arr = _arr([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
46 -
47 - assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=0) == 3
48 - assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=1) == 3
49 - assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=5) == -1
50 - assert pooler._find_array(arr, _arr([3, 5]), start_idx=0) == -1
51 -
52 - with pytest.raises(ValueError):
53 - pooler._find_array(arr, _arr([3, 4, 5]), start_idx=-1)
54 -
55 -
56 -@pytest.fixture(scope="module")
57 -def server_embedding():
58 - # GritLM embedding implementation is only supported by XFormers backend.
59 - args = ["--task", "embed", "--max_model_len", str(MAX_MODEL_LEN)]
60 - with pytest.MonkeyPatch.context() as m:
61 - m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
62 - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
63 - yield remote_server
64 -
65 -
66 -@pytest.fixture(scope="module")
67 -def server_generate():
68 - args = ["--task", "generate", "--max_model_len", str(MAX_MODEL_LEN)]
69 - with pytest.MonkeyPatch.context() as m:
70 - m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
71 - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
72 - yield remote_server
33 +def test_find_array():
34 + from vllm.model_executor.models.gritlm import GritLMPooler
35
36 + model_config = ModelConfig(
37 + MODEL_NAME,
38 + task="embed",
39 + tokenizer=MODEL_NAME,
40 + tokenizer_mode="auto",
41 + trust_remote_code=False,
42 + dtype="bfloat16",
43 + seed=0,
44 + )
45 + pooler = GritLMPooler(model_config=model_config)
46
75 -@pytest_asyncio.fixture
76 -async def client_embedding(server_embedding: RemoteOpenAIServer):
77 - async with server_embedding.get_async_client() as async_client:
78 - yield async_client
47 + arr = _arr([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
48
49 + assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=0) == 3
50 + assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=1) == 3
51 + assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=5) == -1
52 + assert pooler._find_array(arr, _arr([3, 5]), start_idx=0) == -1
53
81 -@pytest_asyncio.fixture
82 -async def client_generate(server_generate: RemoteOpenAIServer):
83 - async with server_generate.get_async_client() as async_client:
84 - yield async_client
54 + with pytest.raises(ValueError):
55 + pooler._find_array(arr, _arr([3, 4, 5]), start_idx=-1)
56
57
58 def run_llm_encode(
88 - llm: vllm.LLM,
59 + llm: LLM,
60 queries: list[str],
61 instruction: str,
91 -) -> list[float]:
92 - outputs = llm.encode([instruction + q for q in queries], )
62 +) -> list[list[float]]:
63 + outputs = llm.embed([instruction + q for q in queries])
64 return [output.outputs.embedding for output in outputs]
65
66
67 async def run_client_embeddings(
97 - client: vllm.LLM,
68 + client: openai.AsyncOpenAI,
69 queries: list[str],
70 instruction: str,
100 -) -> list[float]:
71 +) -> list[list[float]]:
72 outputs = await client.embeddings.create(
73 model=MODEL_NAME,
74 input=[instruction + q for q in queries],
@@ -132,7 +103,7 @@ def get_test_data():
103 return queries, q_instruction, documents, d_instruction
104
105
135 -def validate_embed_output(q_rep: list[float], d_rep: list[float]):
106 +def validate_embed_output(q_rep: list[list[float]], d_rep: list[list[float]]):
107 cosine_sim_q0_d0 = 1 - cosine(q_rep[0], d_rep[0])
108 assert math.isclose(cosine_sim_q0_d0, 0.609, abs_tol=0.001)
109
@@ -143,70 +114,100 @@ def validate_embed_output(q_rep: list[float], d_rep: list[float]):
114 assert math.isclose(cosine_sim_q1_d0, 0.120, abs_tol=0.001)
115
116 cosine_sim_q1_d1 = 1 - cosine(q_rep[1], d_rep[1])
146 - assert math.isclose(cosine_sim_q1_d1, 0.532, abs_tol=0.001)
117 + assert math.isclose(cosine_sim_q1_d1, 0.534, abs_tol=0.001)
118
119
149 -def test_gritlm_offline_embedding(monkeypatch: pytest.MonkeyPatch):
120 +def test_gritlm_offline_embedding(monkeypatch: pytest.MonkeyPatch,
121 + vllm_runner):
122 # GritLM embedding implementation is only supported by XFormers backend.
123 with monkeypatch.context() as m:
124 m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
125
126 queries, q_instruction, documents, d_instruction = get_test_data()
127
156 - llm = vllm.LLM(MODEL_NAME, task="embed", max_model_len=MAX_MODEL_LEN)
128 + with vllm_runner(
129 + MODEL_NAME,
130 + task="embed",
131 + max_model_len=MAX_MODEL_LEN,
132 + ) as vllm_model:
133 + llm = vllm_model.model
134 +
135 + d_rep = run_llm_encode(
136 + llm,
137 + documents,
138 + d_instruction,
139 + )
140 + q_rep = run_llm_encode(
141 + llm,
142 + queries,
143 + q_instruction,
144 + )
145 +
146 + validate_embed_output(q_rep, d_rep)
147 +
148 +
149 +@pytest.mark.asyncio
150 +async def test_gritlm_api_server_embedding():
151 + queries, q_instruction, documents, d_instruction = get_test_data()
152 +
153 + # GritLM embedding implementation is only supported by XFormers backend.
154 + args = ["--task", "embed", "--max_model_len", str(MAX_MODEL_LEN)]
155 + env_dict = {STR_BACKEND_ENV_VAR: "XFORMERS"}
156 +
157 + with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as server:
158 + client_embedding = server.get_async_client()
159
158 - d_rep = run_llm_encode(
159 - llm,
160 + d_rep = await run_client_embeddings(
161 + client_embedding,
162 documents,
163 d_instruction,
164 )
163 - q_rep = run_llm_encode(
164 - llm,
165 + q_rep = await run_client_embeddings(
166 + client_embedding,
167 queries,
168 q_instruction,
169 )
170
169 - validate_embed_output(q_rep, d_rep)
170 -
171 -
172 -@pytest.mark.asyncio
173 -async def test_gritlm_api_server_embedding(
174 - client_embedding: openai.AsyncOpenAI, ):
175 - queries, q_instruction, documents, d_instruction = get_test_data()
171 + validate_embed_output(q_rep, d_rep)
172
177 - d_rep = await run_client_embeddings(
178 - client_embedding,
179 - documents,
180 - d_instruction,
181 - )
182 - q_rep = await run_client_embeddings(
183 - client_embedding,
184 - queries,
185 - q_instruction,
186 - )
173
188 - validate_embed_output(q_rep, d_rep)
174 +def test_gritlm_offline_generate(monkeypatch: pytest.MonkeyPatch, vllm_runner):
175 + # GritLM embedding implementation is only supported by XFormers backend.
176 + with monkeypatch.context() as m:
177 + m.setenv("VLLM_USE_V1", "0")
178 + m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
179
180 + input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
181
191 -def test_gritlm_offline_gen():
192 - input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
182 + with vllm_runner(
183 + MODEL_NAME,
184 + task="generate",
185 + max_model_len=MAX_MODEL_LEN,
186 + ) as vllm_model:
187 + llm = vllm_model.model
188
194 - llm = vllm.LLM(MODEL_NAME, max_model_len=MAX_MODEL_LEN)
195 - sampling_params = vllm.SamplingParams(temperature=0.0, max_tokens=256)
196 - outputs = llm.generate(input, sampling_params=sampling_params)
189 + sampling_params = SamplingParams(temperature=0.0, max_tokens=256)
190 + outputs = llm.generate(input, sampling_params=sampling_params)
191
198 - assert outputs[0].outputs[0].text == "The capital of France is Paris."
192 + assert outputs[0].outputs[0].text == "The capital of France is Paris."
193
194
195 @pytest.mark.asyncio
202 -async def test_gritlm_api_server_gen(client_generate: openai.AsyncOpenAI):
196 +async def test_gritlm_api_server_generate():
197 input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
198
205 - outputs = await client_generate.completions.create(
206 - model=MODEL_NAME,
207 - prompt=input,
208 - max_tokens=256,
209 - temperature=0.0,
210 - )
199 + # GritLM embedding implementation is only supported by XFormers backend.
200 + args = ["--task", "generate", "--max_model_len", str(MAX_MODEL_LEN)]
201 + env_dict = {"VLLM_USE_V1": "0", STR_BACKEND_ENV_VAR: "XFORMERS"}
202 +
203 + with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as server:
204 + client_generate = server.get_async_client()
205 +
206 + outputs = await client_generate.completions.create(
207 + model=MODEL_NAME,
208 + prompt=input,
209 + max_tokens=256,
210 + temperature=0.0,
211 + )
212
213 assert outputs.choices[0].text == "The capital of France is Paris."
tests/models/language/pooling/test_jina.py renamed
+2 -1
@@ -8,9 +8,10 @@ import math
8
9 import pytest
10
11 -from tests.models.embedding.utils import check_embeddings_close, matryoshka_fy
11 from vllm import PoolingParams
12
13 +from ...utils import check_embeddings_close, matryoshka_fy
14 +
15 SCORING_MODELS = [
16 "jinaai/jina-reranker-v2-base-multilingual", # Roberta
17 ]
tests/models/language/pooling/test_scoring.py renamed
tests/models/language/pooling/test_snowflake_arctic_embed.py renamed
+1 -3
@@ -5,9 +5,7 @@ Run `pytest tests/models/embedding/language/test_snowflake_arctic_embed.py`.
5 """
6 import pytest
7
8 -from tests.models.embedding.utils import EmbedModelInfo
9 -
10 -from ..utils import check_embeddings_close
8 +from ...utils import EmbedModelInfo, check_embeddings_close
9
10 EMBEDDING_PROMPTS = [
11 'what is snowflake?', 'Where can I get the best tacos?', 'The Data Cloud!',
tests/models/language/pooling/test_truncation_control.py renamed
tests/models/multimodal/generation/__init__.py renamed
tests/models/multimodal/generation/test_common.py renamed
+7
@@ -267,6 +267,7 @@ VLM_TEST_SETTINGS = {
267 multi_image_prompt="<start_of_image><start_of_image>Describe the two images in detail.", # noqa: E501
268 max_model_len=4096,
269 max_num_seqs=2,
270 + dtype="bfloat16",
271 auto_cls=AutoModelForImageTextToText,
272 vllm_runner_kwargs={"mm_processor_kwargs": {"do_pan_and_scan": True}},
273 patch_hf_runner=model_utils.gemma3_patch_hf_runner,
@@ -423,6 +424,8 @@ VLM_TEST_SETTINGS = {
424 get_stop_token_ids=lambda tok: [tok.eos_id, tok.eot_id],
425 hf_output_post_proc=model_utils.minicpmv_trunc_hf_output,
426 patch_hf_runner=model_utils.minicpmv_25_patch_hf_runner,
427 + # FIXME: https://huggingface.co/openbmb/MiniCPM-V-2_6/discussions/55
428 + marks=[pytest.mark.skip("HF import fails")],
429 ),
430 "minicpmo_26": VLMTestInfo(
431 models=["openbmb/MiniCPM-o-2_6"],
@@ -434,6 +437,8 @@ VLM_TEST_SETTINGS = {
437 get_stop_token_ids=lambda tok: tok.convert_tokens_to_ids(['<|im_end|>', '<|endoftext|>']), # noqa: E501
438 hf_output_post_proc=model_utils.minicpmv_trunc_hf_output,
439 patch_hf_runner=model_utils.minicpmo_26_patch_hf_runner,
440 + # FIXME: https://huggingface.co/openbmb/MiniCPM-V-2_6/discussions/55
441 + marks=[pytest.mark.skip("HF import fails")],
442 ),
443 "minicpmv_26": VLMTestInfo(
444 models=["openbmb/MiniCPM-V-2_6"],
@@ -445,6 +450,8 @@ VLM_TEST_SETTINGS = {
450 get_stop_token_ids=lambda tok: tok.convert_tokens_to_ids(['<|im_end|>', '<|endoftext|>']), # noqa: E501
451 hf_output_post_proc=model_utils.minicpmv_trunc_hf_output,
452 patch_hf_runner=model_utils.minicpmv_26_patch_hf_runner,
453 + # FIXME: https://huggingface.co/openbmb/MiniCPM-V-2_6/discussions/55
454 + marks=[pytest.mark.skip("HF import fails")],
455 ),
456 "minimax_vl_01": VLMTestInfo(
457 models=["MiniMaxAI/MiniMax-VL-01"],
tests/models/multimodal/generation/test_florence2.py renamed
tests/models/multimodal/generation/test_granite_speech.py renamed
tests/models/multimodal/generation/test_interleaved.py renamed
+1
@@ -16,6 +16,7 @@ INTERLEAVED_PROMPT = base_prompt("<image><video><image>\n")
16 NONINTERLEAVED_PROMPT = base_prompt("<image><image><video>\n")
17
18
19 +@pytest.mark.core_model
20 @pytest.mark.parametrize("model", models)
21 @pytest.mark.parametrize("dtype", ["float16"])
22 @pytest.mark.parametrize("max_tokens", [128])
tests/models/multimodal/generation/test_intern_vit.py renamed
tests/models/multimodal/generation/test_mllama.py renamed
+33 -1
@@ -17,7 +17,8 @@ from vllm.sequence import SampleLogprobs
17 from ....conftest import (IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner,
18 _ImageAssets)
19 from ....quantization.utils import is_quant_method_supported
20 -from ....utils import large_gpu_test
20 +from ....utils import (create_new_process_for_each_test, large_gpu_test,
21 + multi_gpu_test)
22 from ...utils import check_logprobs_close
23
24 _LIMIT_IMAGE_PER_PROMPT = 3
@@ -393,6 +394,37 @@ def test_models_interleaved_images(hf_runner, vllm_runner, image_assets, model,
394 )
395
396
397 +@create_new_process_for_each_test()
398 +@multi_gpu_test(num_gpus=2)
399 +@pytest.mark.parametrize("distributed_executor_backend", ["ray", "mp"])
400 +@pytest.mark.parametrize("model", models)
401 +@pytest.mark.parametrize("dtype", ["bfloat16"])
402 +@pytest.mark.parametrize("max_tokens", [64])
403 +@pytest.mark.parametrize("num_logprobs", [5])
404 +def test_models_distributed(
405 + hf_runner,
406 + vllm_runner,
407 + image_assets,
408 + distributed_executor_backend,
409 + model,
410 + dtype,
411 + max_tokens,
412 + num_logprobs,
413 +) -> None:
414 + run_test(
415 + hf_runner,
416 + vllm_runner,
417 + image_assets,
418 + model=model,
419 + size_factors=[0.25, 0.5, 1.0],
420 + dtype=dtype,
421 + max_tokens=max_tokens,
422 + num_logprobs=num_logprobs,
423 + tensor_parallel_size=2,
424 + distributed_executor_backend=distributed_executor_backend,
425 + )
426 +
427 +
428 @large_gpu_test(min_gb=48)
429 @pytest.mark.core_model
430 @pytest.mark.parametrize("model", models)
tests/models/multimodal/generation/test_phi4mm.py renamed
tests/models/multimodal/generation/test_pixtral.py renamed
tests/models/multimodal/generation/test_qwen2_vl.py renamed
tests/models/multimodal/generation/test_ultravox.py renamed
tests/models/multimodal/generation/test_whisper.py renamed
+35 -23
@@ -1,15 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -"""Compare the outputs of HF and vLLM for Whisper models using greedy sampling.
3 -
4 -Run `pytest tests/models/encoder_decoder/audio/test_whisper.py`.
5 -"""
2 from typing import Optional
3
4 import pytest
5
10 -from vllm import LLM, SamplingParams
6 +from vllm import SamplingParams
7 from vllm.assets.audio import AudioAsset
8
9 +from ....conftest import VllmRunner
10 from ....utils import create_new_process_for_each_test, multi_gpu_test
11
12 PROMPTS = [
@@ -92,6 +89,7 @@ EXPECTED = {
89
90
91 def run_test(
92 + vllm_runner: type[VllmRunner],
93 model: str,
94 *,
95 tensor_parallel_size: int,
@@ -100,38 +98,52 @@ def run_test(
98 prompt_list = PROMPTS * 10
99 expected_list = EXPECTED[model] * 10
100
103 - llm = LLM(
104 - model=model,
105 - tensor_parallel_size=tensor_parallel_size,
106 - distributed_executor_backend=distributed_executor_backend,
107 - )
101 + with vllm_runner(
102 + model,
103 + max_model_len=448,
104 + tensor_parallel_size=tensor_parallel_size,
105 + distributed_executor_backend=distributed_executor_backend,
106 + ) as vllm_model:
107 + llm = vllm_model.model
108
109 - sampling_params = SamplingParams(
110 - temperature=0,
111 - top_p=1.0,
112 - max_tokens=200,
113 - )
109 + sampling_params = SamplingParams(
110 + temperature=0,
111 + top_p=1.0,
112 + max_tokens=200,
113 + )
114
115 - outputs = llm.generate(prompt_list, sampling_params)
115 + outputs = llm.generate(prompt_list, sampling_params)
116
117 for output, expected in zip(outputs, expected_list):
118 print(output.outputs[0].text)
119 assert output.outputs[0].text == expected
120
121
122 -@create_new_process_for_each_test()
122 +@create_new_process_for_each_test("spawn")
123 @pytest.mark.core_model
124 @pytest.mark.parametrize(
125 "model", ["openai/whisper-small", "openai/whisper-large-v3-turbo"])
126 -def test_models(model) -> None:
127 - run_test(model, tensor_parallel_size=1)
126 +def test_models(vllm_runner, model) -> None:
127 + run_test(
128 + vllm_runner,
129 + model,
130 + tensor_parallel_size=1,
131 + )
132
133
134 +@create_new_process_for_each_test("spawn")
135 @multi_gpu_test(num_gpus=2)
136 @pytest.mark.core_model
137 @pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
138 @pytest.mark.parametrize("distributed_executor_backend", ["ray", "mp"])
134 -def test_models_distributed(model, distributed_executor_backend) -> None:
135 - run_test(model,
136 - tensor_parallel_size=2,
137 - distributed_executor_backend=distributed_executor_backend)
139 +def test_models_distributed(
140 + vllm_runner,
141 + model,
142 + distributed_executor_backend,
143 +) -> None:
144 + run_test(
145 + vllm_runner,
146 + model,
147 + tensor_parallel_size=2,
148 + distributed_executor_backend=distributed_executor_backend,
149 + )
tests/models/multimodal/generation/vlm_utils/__init__.py renamed
tests/models/multimodal/generation/vlm_utils/builders.py renamed
tests/models/multimodal/generation/vlm_utils/case_filtering.py renamed
tests/models/multimodal/generation/vlm_utils/core.py renamed
tests/models/multimodal/generation/vlm_utils/custom_inputs.py renamed
tests/models/multimodal/generation/vlm_utils/model_utils.py renamed
tests/models/multimodal/generation/vlm_utils/runners.py renamed
tests/models/multimodal/generation/vlm_utils/types.py renamed
tests/models/multimodal/pooling/__init__.py renamed
tests/models/multimodal/pooling/test_dse_qwen2_vl.py renamed
+1 -1
@@ -10,7 +10,7 @@ from transformers import Qwen2VLForConditionalGeneration
10
11 from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
12 from ....utils import large_gpu_test
13 -from ..utils import check_embeddings_close
13 +from ...utils import check_embeddings_close
14
15 HF_TEXT_PROMPTS = [
16 # T -> X
tests/models/multimodal/pooling/test_llava_next.py renamed
+1 -1
@@ -8,7 +8,7 @@ from vllm.platforms import current_platform
8
9 from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
10 from ....utils import large_gpu_test
11 -from ..utils import check_embeddings_close
11 +from ...utils import check_embeddings_close
12
13 # Llava Next embedding implementation is only supported by CUDA.
14 # If run on ROCm, hf_model.model.resize_token_embeddings will
tests/models/multimodal/pooling/test_phi3v.py renamed
+1 -1
@@ -9,7 +9,7 @@ from vllm.assets.image import VLM_IMAGES_DIR
9
10 from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
11 from ....utils import large_gpu_test
12 -from ..utils import check_embeddings_close
12 +from ...utils import check_embeddings_close
13
14 HF_TEXT_PROMPTS = [
15 # T -> X
tests/models/quantization/__init__.py renamed
tests/models/quantization/test_aqlm.py renamed
-1
@@ -39,7 +39,6 @@ ground_truth_generations = [
39 ]
40
41
42 -@pytest.mark.quant_model
42 @pytest.mark.skipif(not is_quant_method_supported("aqlm"),
43 reason="AQLM is not supported on this GPU type.")
44 @pytest.mark.parametrize("model", ["ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf"])
tests/models/quantization/test_awq.py renamed
+2 -3
@@ -7,8 +7,8 @@ import torch
7
8 from vllm.multimodal.image import rescale_image_size
9
10 -from ....conftest import IMAGE_ASSETS, VllmRunner, _ImageAssets
11 -from ...utils import check_logprobs_close
10 +from ...conftest import IMAGE_ASSETS, VllmRunner, _ImageAssets
11 +from ..utils import check_logprobs_close
12
13 HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts({
14 "stop_sign":
@@ -85,7 +85,6 @@ def run_awq_test(
85 )
86
87
88 -@pytest.mark.quant_model
88 @pytest.mark.parametrize(
89 ("source_model", "quant_model"),
90 [("OpenGVLab/InternVL2-2B", "OpenGVLab/InternVL2-2B-AWQ")],
tests/models/quantization/test_bitblas.py renamed
+1 -1
@@ -15,7 +15,7 @@ from dataclasses import dataclass
15
16 import pytest
17
18 -from .utils import check_logprobs_close
18 +from ..utils import check_logprobs_close
19
20
21 @dataclass
tests/models/quantization/test_fp8.py renamed
+1 -6
@@ -4,20 +4,15 @@
4 """Tests fp8 models against ground truth generation
5 Note: these tests will only pass on L4 GPU.
6 """
7 -import os
8 -from typing import Optional
9 -
7 import pytest
8
12 -from tests.kernels.utils import override_backend_env_variable
9 from tests.quantization.utils import is_quant_method_supported
10 from vllm.platforms import current_platform
11 from vllm.utils import STR_BACKEND_ENV_VAR
12
17 -from ...utils import check_logprobs_close
13 +from ..utils import check_logprobs_close
14
15
20 -@pytest.mark.quant_model
16 @pytest.mark.skipif(not is_quant_method_supported("fp8"),
17 reason="fp8 is not supported on this GPU type.")
18 @pytest.mark.parametrize(
tests/models/quantization/test_gguf.py renamed
+3 -4
@@ -14,9 +14,9 @@ from transformers import AutoTokenizer
14
15 from tests.quantization.utils import is_quant_method_supported
16
17 -from ....conftest import VllmRunner
18 -from ....utils import multi_gpu_test
19 -from ...utils import check_logprobs_close
17 +from ...conftest import VllmRunner
18 +from ...utils import multi_gpu_test
19 +from ..utils import check_logprobs_close
20
21 os.environ["TOKENIZERS_PARALLELISM"] = "true"
22
@@ -38,7 +38,6 @@ LLAMA_CONFIG = GGUFTestConfig(
38 original_model="meta-llama/Llama-3.2-1B-Instruct",
39 gguf_repo="bartowski/Llama-3.2-1B-Instruct-GGUF",
40 gguf_filename="Llama-3.2-1B-Instruct-IQ4_XS.gguf",
41 - marks=[pytest.mark.quant_model],
41 )
42
43 QWEN2_CONFIG = GGUFTestConfig(
tests/models/quantization/test_gptq_bitblas.py renamed
+1 -1
@@ -15,7 +15,7 @@ from dataclasses import dataclass
15
16 import pytest
17
18 -from .utils import check_logprobs_close
18 +from ..utils import check_logprobs_close
19
20
21 @dataclass
tests/models/quantization/test_gptq_marlin.py renamed
+1 -2
@@ -16,7 +16,7 @@ import pytest
16 from tests.quantization.utils import is_quant_method_supported
17 from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT
18
19 -from ...utils import check_logprobs_close
19 +from ..utils import check_logprobs_close
20
21 os.environ["TOKENIZERS_PARALLELISM"] = "true"
22
@@ -34,7 +34,6 @@ MODELS = [
34 ]
35
36
37 -@pytest.mark.quant_model
37 @pytest.mark.flaky(reruns=3)
38 @pytest.mark.skipif(not is_quant_method_supported("gptq_marlin"),
39 reason="gptq_marlin is not supported on this GPU type.")
tests/models/quantization/test_gptq_marlin_24.py renamed
+1 -2
@@ -13,7 +13,7 @@ import pytest
13
14 from tests.quantization.utils import is_quant_method_supported
15
16 -from ...utils import check_logprobs_close
16 +from ..utils import check_logprobs_close
17
18
19 @dataclass
@@ -39,7 +39,6 @@ model_pairs = [
39 ]
40
41
42 -@pytest.mark.quant_model
42 @pytest.mark.flaky(reruns=2)
43 @pytest.mark.skipif(not is_quant_method_supported("gptq_marlin_24"),
44 reason="Marlin24 is not supported on this GPU type.")
tests/models/quantization/test_modelopt.py renamed
-1
@@ -40,7 +40,6 @@ EXPECTED_STRS_MAP = {
40 @pytest.mark.skip(
41 reason=
42 "Prevent unstable test based on golden strings from breaking the build.")
43 -@pytest.mark.quant_model
43 @pytest.mark.skipif(not is_quant_method_supported("fp8"),
44 reason="fp8 is not supported on this GPU type.")
45 @pytest.mark.parametrize("model_name", MODELS)
tests/models/quantization/test_nvfp4.py renamed
-1
@@ -41,7 +41,6 @@ EXPECTED_STRS_MAP = {
41 reason=
42 "Prevent unstable test based on golden strings from breaking the build "
43 " and test input model being too large and hanging the system.")
44 -@pytest.mark.quant_model
44 @pytest.mark.skipif(not is_quant_method_supported("nvfp4"),
45 reason="nvfp4 is not supported on this GPU type.")
46 @pytest.mark.parametrize("model_name", MODELS)
tests/models/utils.py
+65 -1
@@ -2,9 +2,10 @@
2
3 import warnings
4 from collections.abc import Sequence
5 -from typing import Any, Optional, Union
5 +from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union
6
7 import torch
8 +import torch.nn.functional as F
9
10 from vllm.config import ModelConfig, TaskOption
11 from vllm.inputs import InputContext
@@ -12,6 +13,9 @@ from vllm.sequence import Logprob, PromptLogprobs, SampleLogprobs
13
14 from .registry import HF_EXAMPLE_MODELS
15
16 +if TYPE_CHECKING:
17 + from ..conftest import HfRunner
18 +
19 TokensText = tuple[list[int], str]
20
21
@@ -291,3 +295,63 @@ def build_model_context(
295 **model_config_kwargs,
296 )
297 return InputContext(model_config)
298 +
299 +
300 +def check_embeddings_close(
301 + *,
302 + embeddings_0_lst: Sequence[list[float]],
303 + embeddings_1_lst: Sequence[list[float]],
304 + name_0: str,
305 + name_1: str,
306 + tol: float = 1e-3,
307 +) -> None:
308 + assert len(embeddings_0_lst) == len(embeddings_1_lst)
309 +
310 + for prompt_idx, (embeddings_0, embeddings_1) in enumerate(
311 + zip(embeddings_0_lst, embeddings_1_lst)):
312 + assert len(embeddings_0) == len(embeddings_1), (
313 + f"Length mismatch: {len(embeddings_0)} vs. {len(embeddings_1)}")
314 +
315 + sim = F.cosine_similarity(torch.tensor(embeddings_0),
316 + torch.tensor(embeddings_1),
317 + dim=0)
318 +
319 + fail_msg = (f"Test{prompt_idx}:"
320 + f"\n{name_0}:\t{embeddings_0[:16]!r}"
321 + f"\n{name_1}:\t{embeddings_1[:16]!r}")
322 +
323 + assert sim >= 1 - tol, fail_msg
324 +
325 +
326 +def matryoshka_fy(tensor: torch.Tensor, dimensions: int):
327 + tensor = torch.tensor(tensor)
328 + tensor = tensor[..., :dimensions]
329 + tensor = F.normalize(tensor, p=2, dim=1)
330 + return tensor
331 +
332 +
333 +class EmbedModelInfo(NamedTuple):
334 + name: str
335 + is_matryoshka: bool
336 + matryoshka_dimensions: Optional[list[int]] = None
337 + architecture: str = ""
338 + enable_test: bool = True
339 +
340 +
341 +def run_embedding_correctness_test(
342 + hf_model: "HfRunner",
343 + inputs: list[str],
344 + vllm_outputs: Sequence[list[float]],
345 + dimensions: Optional[int] = None,
346 +):
347 + hf_outputs = hf_model.encode(inputs)
348 + if dimensions:
349 + hf_outputs = matryoshka_fy(hf_outputs, dimensions)
350 +
351 + check_embeddings_close(
352 + embeddings_0_lst=hf_outputs,
353 + embeddings_1_lst=vllm_outputs,
354 + name_0="hf",
355 + name_1="vllm",
356 + tol=1e-2,
357 + )
vllm/config.py
+3 -1
@@ -1043,8 +1043,10 @@ class ModelConfig:
1043 if self.is_attention_free:
1044 return 0
1045
1046 - if hasattr(self.hf_text_config, "head_dim"):
1046 + # NOTE: Some configs may set head_dim=None in the config
1047 + if getattr(self.hf_text_config, "head_dim", None) is not None:
1048 return self.hf_text_config.head_dim
1049 +
1050 # FIXME(woosuk): This may not be true for all models.
1051 return (self.hf_text_config.hidden_size //
1052 self.hf_text_config.num_attention_heads)
vllm/model_executor/models/llama.py
+4 -2
@@ -127,8 +127,10 @@ class LlamaAttention(nn.Module):
127 assert tp_size % self.total_num_kv_heads == 0
128 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
129 # MistralConfig has an optional head_dim introduced by Mistral-Nemo
130 - self.head_dim = getattr(config, "head_dim",
131 - self.hidden_size // self.total_num_heads)
130 + head_dim = getattr(config, "head_dim", None)
131 + if head_dim is None:
132 + head_dim = self.hidden_size // self.total_num_heads
133 + self.head_dim = head_dim
134 # Phi models introduced a partial_rotary_factor parameter in the config
135 self.partial_rotary_factor = getattr(config, "partial_rotary_factor",
136 1)