@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 213 lines 7.36 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass
4 from typing import Any
5
6 import pytest
7 from pydantic import TypeAdapter, ValidationError
8
9 from tests.models.utils import EmbedModelInfo
10 from vllm import PoolingParams
11 from vllm.config import ModelConfig, PoolerConfig
12 from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest
13 from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest
14 from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest
15 from vllm.exceptions import VLLMValidationError
16
17 EMBEDDING_MODELS = [
18 EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False),
19 EmbedModelInfo(
20 "Snowflake/snowflake-arctic-embed-m-v1.5",
21 is_matryoshka=True,
22 matryoshka_dimensions=[256],
23 ),
24 ]
25
26 classify_parameters = ["use_activation"]
27 embed_parameters = ["dimensions", "use_activation"]
28 step_pooling_parameters = ["step_tag_id", "returned_token_ids"]
29
30
31 @dataclass()
32 class MockModelConfig:
33 pooler_config: PoolerConfig
34
35
36 @pytest.mark.parametrize(
37 ("parameter", "value", "message"),
38 [
39 (
40 "normalize",
41 False,
42 "Parameter `normalize` was removed; use `use_activation` instead.",
43 ),
44 ("task", "score", "`score` task was removed; use `classify` instead."),
45 (
46 "task",
47 "encode",
48 "`encode` task was removed; use `token_embed` or `token_classify` instead.",
49 ),
50 ],
51 )
52 def test_removed_pooling_parameters(parameter: str, value: Any, message: str):
53 data = {"input": "hello", parameter: value}
54 for request_type in (EmbeddingRequest, ClassificationRequest, PoolingRequest):
55 with pytest.raises(VLLMValidationError, match=message):
56 TypeAdapter(request_type).validate_python(data)
57
58 # PoolerConfig still raises bare ValueError for `normalize`
59 # (wrapped to ValidationError by Pydantic), but `check_removed_pooling_task`
60 # raises VLLMValidationError for removed tasks.
61 if parameter == "normalize":
62 with pytest.raises(ValidationError, match=message) as exc_info:
63 TypeAdapter(PoolerConfig).validate_python({parameter: value})
64 assert len(exc_info.value.errors()) == 1
65 else:
66 with pytest.raises(VLLMValidationError, match=message):
67 TypeAdapter(PoolerConfig).validate_python({parameter: value})
68
69 if parameter == "task":
70 with pytest.raises(VLLMValidationError, match=message):
71 PoolingParams(task=value)
72
73
74 def test_embed():
75 task = "embed"
76 model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
77
78 pooling_params = PoolingParams(task=task, use_activation=None)
79 pooling_params.verify(model_config)
80
81 pooling_params = PoolingParams(task=task, use_activation=True)
82 pooling_params.verify(model_config)
83
84 pooling_params = PoolingParams(task=task, use_activation=False)
85 pooling_params.verify(model_config)
86
87 invalid_parameters = classify_parameters + step_pooling_parameters
88 for p in set(invalid_parameters) - set(embed_parameters):
89 with pytest.raises(VLLMValidationError):
90 pooling_params = PoolingParams(task=task, **{p: True})
91 pooling_params.verify(model_config)
92
93
94 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
95 def test_embed_dimensions(model_info: EmbedModelInfo):
96 task = "embed"
97 model_config = ModelConfig(
98 model_info.name,
99 tokenizer=model_info.name,
100 tokenizer_mode="auto",
101 trust_remote_code=False,
102 seed=0,
103 dtype="float16",
104 )
105
106 pooling_params = PoolingParams(task=task, dimensions=None)
107 pooling_params.verify(model_config)
108
109 with pytest.raises(VLLMValidationError):
110 pooling_params = PoolingParams(task=task, dimensions=1)
111 pooling_params.verify(model_config)
112
113 if model_info.is_matryoshka:
114 assert model_info.matryoshka_dimensions is not None
115 pooling_params = PoolingParams(
116 task=task, dimensions=model_info.matryoshka_dimensions[0]
117 )
118 pooling_params.verify(model_config)
119
120
121 @dataclass()
122 class MockMatryoshkaModelConfig:
123 pooler_config: PoolerConfig
124 is_matryoshka: bool = True
125 matryoshka_dimensions: list[int] | None = None
126 served_model_name: str = "mock-matryoshka-model"
127 embedding_size: int = 32
128
129
130 def test_embed_dimensions_matryoshka_without_list_upper_bound():
131 task = "embed"
132 model_config = MockMatryoshkaModelConfig(
133 pooler_config=PoolerConfig(seq_pooling_type="CLS"),
134 matryoshka_dimensions=None,
135 embedding_size=32,
136 )
137
138 PoolingParams(task=task, dimensions=16).verify(model_config)
139
140 with pytest.raises(VLLMValidationError):
141 PoolingParams(task=task, dimensions=64).verify(model_config)
142
143
144 @pytest.mark.parametrize("task", ["classify"])
145 def test_classify(task):
146 model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
147
148 pooling_params = PoolingParams(task=task, use_activation=None)
149 pooling_params.verify(model_config)
150
151 pooling_params = PoolingParams(task=task, use_activation=True)
152 pooling_params.verify(model_config)
153
154 pooling_params = PoolingParams(task=task, use_activation=False)
155 pooling_params.verify(model_config)
156
157 invalid_parameters = embed_parameters + step_pooling_parameters
158 for p in set(invalid_parameters) - set(classify_parameters):
159 with pytest.raises(VLLMValidationError):
160 pooling_params = PoolingParams(task=task, **{p: True})
161 pooling_params.verify(model_config)
162
163
164 @pytest.mark.parametrize("pooling_type", ["ALL", "STEP"])
165 def test_token_embed(pooling_type: str):
166 task = "token_embed"
167 model_config = MockModelConfig(
168 pooler_config=PoolerConfig(tok_pooling_type=pooling_type)
169 )
170
171 pooling_params = PoolingParams(task=task, use_activation=None)
172 pooling_params.verify(model_config)
173
174 pooling_params = PoolingParams(task=task, use_activation=True)
175 pooling_params.verify(model_config)
176
177 pooling_params = PoolingParams(task=task, use_activation=False)
178 pooling_params.verify(model_config)
179
180 invalid_parameters = classify_parameters
181 if pooling_type != "STEP":
182 invalid_parameters = classify_parameters + step_pooling_parameters
183
184 for p in set(invalid_parameters) - set(embed_parameters):
185 with pytest.raises(VLLMValidationError):
186 pooling_params = PoolingParams(task=task, **{p: True})
187 pooling_params.verify(model_config)
188
189
190 @pytest.mark.parametrize("pooling_type", ["ALL", "STEP"])
191 def test_token_classify(pooling_type: str):
192 task = "token_classify"
193 model_config = MockModelConfig(
194 pooler_config=PoolerConfig(tok_pooling_type=pooling_type)
195 )
196
197 pooling_params = PoolingParams(task=task, use_activation=None)
198 pooling_params.verify(model_config)
199
200 pooling_params = PoolingParams(task=task, use_activation=True)
201 pooling_params.verify(model_config)
202
203 pooling_params = PoolingParams(task=task, use_activation=False)
204 pooling_params.verify(model_config)
205
206 invalid_parameters = embed_parameters
207 if pooling_type != "STEP":
208 invalid_parameters = embed_parameters + step_pooling_parameters
209
210 for p in set(invalid_parameters) - set(classify_parameters):
211 with pytest.raises(VLLMValidationError):
212 pooling_params = PoolingParams(task=task, **{p: True})
213 pooling_params.verify(model_config)