@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 176 lines 5.17 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Generation tests for Moondream3 query and caption support."""
4
5 import pytest
6
7 from tests.models.registry import HF_EXAMPLE_MODELS
8 from vllm.platforms import current_platform
9
10 from ....conftest import IMAGE_ASSETS, ImageTestAssets
11 from ....utils import large_gpu_mark, multi_gpu_test
12
13 MOONDREAM3_MODEL_ID = "moondream/moondream3-preview"
14 MOONDREAM3_TOKENIZER = "moondream/starmie-v1"
15
16 HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
17 {
18 "stop_sign": "<|endoftext|><image><|md_reserved_0|>query<|md_reserved_1|>What color is the stop sign?<|md_reserved_2|>", # noqa: E501
19 "cherry_blossom": "<|endoftext|><image><|md_reserved_0|>query<|md_reserved_1|>What color are the flowers?<|md_reserved_2|>", # noqa: E501
20 }
21 )
22
23
24 def make_query_prompt(question: str) -> str:
25 """Create a direct-answer query prompt for Moondream3."""
26 return (
27 "<|endoftext|><image><|md_reserved_0|>query<|md_reserved_1|>"
28 f"{question}<|md_reserved_2|>"
29 )
30
31
32 def make_caption_prompt(length: str = "normal") -> str:
33 """Create a caption prompt for Moondream3."""
34 return (
35 "<|endoftext|><image><|md_reserved_0|>"
36 f"describe<|md_reserved_1|>{length}<|md_reserved_2|>"
37 )
38
39
40 @multi_gpu_test(num_gpus=2)
41 @large_gpu_mark(min_gb=80)
42 def test_tensor_parallel(image_assets: ImageTestAssets):
43 import gc
44
45 from vllm import LLM, SamplingParams
46 from vllm.distributed.parallel_state import destroy_model_parallel
47
48 destroy_model_parallel()
49 gc.collect()
50 current_platform.empty_cache()
51
52 llm = LLM(
53 model=MOONDREAM3_MODEL_ID,
54 tokenizer=MOONDREAM3_TOKENIZER,
55 trust_remote_code=True,
56 dtype="bfloat16",
57 tensor_parallel_size=2,
58 max_model_len=1024,
59 enforce_eager=True,
60 limit_mm_per_prompt={"image": 1},
61 gpu_memory_utilization=0.45,
62 )
63
64 image = image_assets[0].pil_image
65 prompt = make_query_prompt("What color is the stop sign?")
66
67 try:
68 outputs = llm.generate(
69 {"prompt": prompt, "multi_modal_data": {"image": image}},
70 SamplingParams(max_tokens=20, temperature=0),
71 )
72
73 assert len(outputs) > 0
74 assert outputs[0].outputs[0].text is not None
75 finally:
76 del llm
77 gc.collect()
78 current_platform.empty_cache()
79
80
81 @pytest.fixture(scope="module")
82 def llm():
83 model_info = HF_EXAMPLE_MODELS.get_hf_info("Moondream3ForCausalLM")
84 model_info.check_transformers_version(on_fail="skip")
85
86 from vllm import LLM
87
88 try:
89 return LLM(
90 model=MOONDREAM3_MODEL_ID,
91 tokenizer=MOONDREAM3_TOKENIZER,
92 trust_remote_code=True,
93 dtype="bfloat16",
94 max_model_len=2048,
95 enforce_eager=True,
96 limit_mm_per_prompt={"image": 1},
97 gpu_memory_utilization=0.45,
98 )
99 except Exception as exc:
100 pytest.skip(f"Failed to load {MOONDREAM3_MODEL_ID}: {exc}")
101
102
103 @large_gpu_mark(min_gb=48)
104 def test_model_loading(llm):
105 assert llm is not None
106
107
108 @large_gpu_mark(min_gb=48)
109 def test_query_skill(llm, image_assets: ImageTestAssets):
110 from vllm import SamplingParams
111
112 image = image_assets[0].pil_image
113 prompt = make_query_prompt("What color is the stop sign?")
114
115 outputs = llm.generate(
116 {"prompt": prompt, "multi_modal_data": {"image": image}},
117 SamplingParams(max_tokens=50, temperature=0),
118 )
119
120 output_text = outputs[0].outputs[0].text
121 assert output_text is not None
122 assert len(output_text) > 0
123
124
125 @large_gpu_mark(min_gb=48)
126 def test_caption_skill(llm, image_assets: ImageTestAssets):
127 from vllm import SamplingParams
128
129 image = image_assets[1].pil_image
130 prompt = make_caption_prompt()
131
132 outputs = llm.generate(
133 {"prompt": prompt, "multi_modal_data": {"image": image}},
134 SamplingParams(max_tokens=100, temperature=0),
135 )
136
137 output_text = outputs[0].outputs[0].text
138 assert output_text is not None
139 assert len(output_text) > 0
140
141
142 @large_gpu_mark(min_gb=48)
143 def test_batched_inference(llm, image_assets: ImageTestAssets):
144 from vllm import SamplingParams
145
146 images = [asset.pil_image for asset in image_assets]
147 prompts = [
148 {"prompt": prompt, "multi_modal_data": {"image": img}}
149 for img, prompt in zip(images, HF_IMAGE_PROMPTS)
150 ]
151
152 outputs = llm.generate(prompts, SamplingParams(max_tokens=50, temperature=0))
153
154 assert len(outputs) == len(images)
155 for output in outputs:
156 assert output.outputs[0].text is not None
157 assert len(output.outputs[0].text) > 0
158
159
160 @pytest.mark.parametrize("asset_name", ["stop_sign", "cherry_blossom"])
161 @large_gpu_mark(min_gb=48)
162 def test_image_assets(llm, image_assets: ImageTestAssets, asset_name: str):
163 from vllm import SamplingParams
164
165 asset_idx = 0 if asset_name == "stop_sign" else 1
166 image = image_assets[asset_idx].pil_image
167 prompt = HF_IMAGE_PROMPTS[asset_idx]
168
169 outputs = llm.generate(
170 {"prompt": prompt, "multi_modal_data": {"image": image}},
171 SamplingParams(max_tokens=50, temperature=0),
172 )
173
174 output_text = outputs[0].outputs[0].text
175 assert output_text is not None
176 assert len(output_text) > 0