@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 167 lines 5.15 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import os
5
6 os.environ["TOKENIZERS_PARALLELISM"] = "true"
7
8 from typing import Any, NamedTuple
9
10 import pytest
11 from pytest import MarkDecorator
12 from transformers import AutoModelForImageTextToText
13
14 from vllm.assets.image import ImageAsset
15 from vllm.multimodal.image import rescale_image_size
16 from vllm.transformers_utils.repo_utils import hf_api
17 from vllm.utils.torch_utils import set_default_torch_num_threads
18
19 from ...conftest import IMAGE_ASSETS, HfRunner, VllmRunner
20 from ...models.utils import check_logprobs_close
21
22
23 class GGUFMMTestConfig(NamedTuple):
24 original_model: str
25 gguf_repo: str
26 gguf_backbone: str
27 gguf_mmproj: str
28 prompt: list[str]
29 image_names: list[str] # Store names, load PIL images at runtime
30 max_model_len: int = 4096
31 marks: list[MarkDecorator] = []
32 mm_processor_kwargs: dict[str, Any] = {}
33
34 @property
35 def gguf_model(self):
36 hf_api().hf_hub_download(self.gguf_repo, filename=self.gguf_mmproj)
37 return hf_api().hf_hub_download(self.gguf_repo, filename=self.gguf_backbone)
38
39
40 # Common prompts aligned with test_common.py "gemma3" entry format
41 _GEMMA3_PROMPTS = IMAGE_ASSETS.prompts(
42 {
43 "stop_sign": (
44 "<bos><start_of_turn>user\n"
45 "<start_of_image>What's the content in the center of the image?"
46 "<end_of_turn>\n<start_of_turn>model\n"
47 ),
48 "cherry_blossom": (
49 "<bos><start_of_turn>user\n"
50 "<start_of_image>What is the season?"
51 "<end_of_turn>\n<start_of_turn>model\n"
52 ),
53 }
54 )
55
56 # Image asset names - load at runtime to avoid pickle issues with subprocess
57 _GEMMA3_IMAGE_NAMES = ["stop_sign", "cherry_blossom"]
58
59 # Regular multimodal (no pan-and-scan) - uses QAT Q4_0 GGUF
60 GEMMA3_CONFIG = GGUFMMTestConfig(
61 original_model="google/gemma-3-4b-it",
62 gguf_repo="google/gemma-3-4b-it-qat-q4_0-gguf",
63 gguf_backbone="gemma-3-4b-it-q4_0.gguf",
64 gguf_mmproj="mmproj-model-f16-4B.gguf",
65 prompt=_GEMMA3_PROMPTS,
66 image_names=_GEMMA3_IMAGE_NAMES,
67 max_model_len=4096,
68 mm_processor_kwargs={},
69 )
70
71 # Pan-and-scan multimodal - uses unquantized BF16 GGUF
72 GEMMA3_CONFIG_PAN_AND_SCAN = GGUFMMTestConfig(
73 original_model="google/gemma-3-4b-it",
74 gguf_repo="google/gemma-3-4b-it-qat-q4_0-gguf",
75 gguf_backbone="gemma-3-4b-it-q4_0.gguf",
76 gguf_mmproj="mmproj-model-f16-4B.gguf",
77 prompt=_GEMMA3_PROMPTS,
78 image_names=_GEMMA3_IMAGE_NAMES,
79 max_model_len=4096,
80 mm_processor_kwargs={"do_pan_and_scan": True},
81 )
82
83 MODELS_TO_TEST = [GEMMA3_CONFIG, GEMMA3_CONFIG_PAN_AND_SCAN]
84
85
86 def run_multimodal_gguf_test(
87 hf_runner: type[HfRunner],
88 vllm_runner: type[VllmRunner],
89 model: GGUFMMTestConfig,
90 dtype: str,
91 max_tokens: int,
92 num_logprobs: int,
93 ):
94 # Load images at runtime (inside subprocess) to avoid pickle issues
95 images = [ImageAsset(name).pil_image for name in model.image_names]
96 size_factors = [0.25, 0.5, 1.0]
97 inputs_per_image = [
98 (
99 [prompt for _ in size_factors],
100 [rescale_image_size(image, factor) for factor in size_factors],
101 )
102 for image, prompt in zip(images, model.prompt)
103 ]
104
105 # NOTE: Run vLLM first to avoid CUDA init issues with multiprocessing fork.
106 # Run GGUF model via vLLM.
107 with (
108 set_default_torch_num_threads(1),
109 vllm_runner(
110 model_name=model.gguf_model,
111 enforce_eager=True,
112 tokenizer_name=model.original_model,
113 dtype=dtype,
114 max_model_len=model.max_model_len,
115 mm_processor_kwargs=model.mm_processor_kwargs,
116 ) as gguf_model,
117 ):
118 gguf_outputs_per_case = [
119 gguf_model.generate_greedy_logprobs(
120 prompts,
121 max_tokens,
122 num_logprobs=num_logprobs,
123 images=images,
124 )
125 for prompts, images in inputs_per_image
126 ]
127
128 # Then run HfRunner for HuggingFace baseline comparison.
129 with hf_runner(
130 model.original_model,
131 dtype=dtype,
132 auto_cls=AutoModelForImageTextToText,
133 ) as hf_model:
134 hf_outputs_per_case = [
135 hf_model.generate_greedy_logprobs_limit(
136 prompts,
137 max_tokens,
138 num_logprobs=num_logprobs,
139 images=images,
140 )
141 for prompts, images in inputs_per_image
142 ]
143
144 for hf_outputs, gguf_outputs in zip(hf_outputs_per_case, gguf_outputs_per_case):
145 check_logprobs_close(
146 outputs_0_lst=hf_outputs,
147 outputs_1_lst=gguf_outputs,
148 name_0="hf",
149 name_1="gguf",
150 )
151
152
153 @pytest.mark.parametrize("model", MODELS_TO_TEST)
154 @pytest.mark.parametrize("dtype", ["bfloat16"])
155 @pytest.mark.parametrize("max_tokens", [32])
156 @pytest.mark.parametrize("num_logprobs", [10])
157 def test_gemma3_mm_gguf(
158 hf_runner: type[HfRunner],
159 vllm_runner: type[VllmRunner],
160 model: GGUFMMTestConfig,
161 dtype: str,
162 max_tokens: int,
163 num_logprobs: int,
164 ) -> None:
165 run_multimodal_gguf_test(
166 hf_runner, vllm_runner, model, dtype, max_tokens, num_logprobs
167 )