@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 215 lines 6.41 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable
5
6 import pytest
7 import torch
8 import torch.nn.functional as F
9 from PIL import Image
10 from transformers import Qwen2VLForConditionalGeneration
11
12 from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
13 from ....utils import large_gpu_test
14 from ...utils import check_embeddings_close
15
16 HF_TEXT_PROMPTS = [
17 # T -> X
18 (
19 "Query: Find me an everyday image that matches the given caption: The label of the object is stop sign", # noqa: E501,
20 Image.new("RGB", (56, 56)),
21 ),
22 # T -> X
23 (
24 "Query: Retrieve an image of this caption: cherry blossom",
25 Image.new("RGB", (56, 56)),
26 ),
27 ]
28
29 HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
30 {
31 "stop_sign": "What is shown in this image?",
32 "cherry_blossom": "What is shown in this image?",
33 }
34 )
35
36 MODELS = ["MrLight/dse-qwen2-2b-mrl-v1"]
37
38
39 def get_messages(image: Image.Image, text: str, embed_text: bool):
40 # assert False, 'remember to use outer [] as required'
41 if embed_text:
42 messages = [
43 {
44 "role": "user",
45 "content": [
46 {
47 "type": "image",
48 "image": Image.new("RGB", (56, 56)),
49 "resized_height": 1,
50 "resized_width": 1,
51 }, # need a dummy image here for an easier process.
52 {"type": "text", "text": text},
53 ],
54 }
55 ]
56 else:
57 messages = [
58 {
59 "role": "user",
60 "content": [
61 {"type": "image", "image": image},
62 {"type": "text", "text": text},
63 ],
64 }
65 ]
66 return messages
67
68
69 def apply_chat_template_and_add_eos(
70 messages: list[dict],
71 apply_chat_template_fn: Callable,
72 ):
73 prompt = (
74 apply_chat_template_fn(messages, tokenize=False, add_generation_prompt=True)
75 + "<|endoftext|>"
76 )
77 return prompt
78
79
80 def _run_test(
81 hf_runner: type[HfRunner],
82 vllm_runner: type[VllmRunner],
83 input_texts: list[str],
84 input_images: PromptImageInput,
85 embed_texts: list[bool],
86 model: str,
87 *,
88 dtype: str,
89 ) -> None:
90 """SET PYTHONPATH."""
91 # NOTE: take care of the order. run vLLM first, and then run HF.
92 # vLLM needs a fresh new process without cuda initialization.
93 # if we run HF first, the cuda initialization will be done and it
94 # will hurt multiprocessing backend with fork method (the default method).
95 with vllm_runner(
96 model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=8192
97 ) as vllm_model:
98 tokenizer = vllm_model.llm.get_tokenizer()
99 texts = [
100 # this is necessary because vllm_model.embed will not apply any
101 # templating to the prompt, and therefore lacks an image_pad
102 # token unless one is inserted beforehand (the (28,28) image
103 # above is converted to an image pad token by the chat template).
104 apply_chat_template_and_add_eos(
105 get_messages(image, text, False),
106 apply_chat_template_fn=tokenizer.apply_chat_template,
107 )
108 for text, image in zip(input_texts, input_images)
109 # vllm will replace the pad token with the actual image,
110 # which may be a placeholder image, later.
111 ]
112 vllm_outputs = vllm_model.embed(texts, images=input_images)
113
114 hf_outputs = []
115 with hf_runner(
116 model, dtype=dtype, auto_cls=Qwen2VLForConditionalGeneration
117 ) as hf_model:
118 prompts = []
119 for text, image, embed_text in zip(input_texts, input_images, embed_texts):
120 # dse requires non-standard input processing
121 # because it needs an image_pad token
122 messages = get_messages(image, text, embed_text)
123 prompt = apply_chat_template_and_add_eos(
124 messages, hf_model.processor.apply_chat_template
125 )
126
127 prompts.append(prompt)
128
129 all_inputs = hf_model.get_inputs(
130 prompts=prompts,
131 images=input_images,
132 )
133
134 with torch.no_grad():
135 all_outputs = []
136 for inputs in all_inputs:
137 inputs = hf_model.model.prepare_inputs_for_generation(
138 **inputs,
139 cache_position=torch.arange(1), # 1 for batch size
140 use_cache=False,
141 )
142 outputs = hf_model.model(
143 **hf_model.wrap_device(inputs),
144 return_dict=True,
145 output_hidden_states=True,
146 )
147 pooled_output = F.normalize(
148 outputs.hidden_states[-1][0, -1], p=2, dim=-1
149 )
150
151 all_outputs.append(pooled_output.tolist())
152
153 hf_outputs = all_outputs
154
155 check_embeddings_close(
156 embeddings_0_lst=hf_outputs,
157 embeddings_1_lst=vllm_outputs,
158 name_0="hf",
159 name_1="vllm",
160 )
161
162
163 @pytest.mark.parametrize("model", MODELS)
164 @pytest.mark.parametrize("dtype", ["bfloat16"])
165 def test_models_text(
166 hf_runner,
167 vllm_runner,
168 image_assets,
169 model: str,
170 dtype: str,
171 ) -> None:
172 input_texts_images = [
173 (text, image_placeholder) for text, image_placeholder in HF_TEXT_PROMPTS
174 ]
175 input_texts = [text for text, _ in input_texts_images]
176 input_images = [image for _, image in input_texts_images]
177 embed_texts = [True] * len(input_texts)
178
179 _run_test(
180 hf_runner,
181 vllm_runner,
182 input_texts,
183 input_images, # type: ignore
184 embed_texts,
185 model,
186 dtype=dtype,
187 )
188
189
190 @large_gpu_test(min_gb=48)
191 @pytest.mark.parametrize("model", MODELS)
192 @pytest.mark.parametrize("dtype", ["bfloat16"])
193 def test_models_image(
194 hf_runner,
195 vllm_runner,
196 image_assets,
197 model: str,
198 dtype: str,
199 ) -> None:
200 input_texts_images = [
201 (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
202 ]
203 input_texts = [text for text, _ in input_texts_images]
204 input_images = [image for _, image in input_texts_images]
205 embed_texts = [False] * len(input_texts)
206
207 _run_test(
208 hf_runner,
209 vllm_runner,
210 input_texts,
211 input_images,
212 embed_texts,
213 model,
214 dtype=dtype,
215 )