@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 124 lines 3.79 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 import vllm
7 from vllm.assets.image import ImageAsset
8 from vllm.lora.request import LoRARequest
9 from vllm.platforms import current_platform
10
11 from ..utils import multi_gpu_test
12
13 MODEL_PATH = "openbmb/MiniCPM-Llama3-V-2_5"
14
15 PROMPT_TEMPLATE = (
16 "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n"
17 "(<image>./</image>)\nWhat is in the image?<|eot_id|>"
18 "<|start_header_id|>assistant<|end_header_id|>\n\n"
19 )
20
21 IMAGE_ASSETS = [
22 ImageAsset("stop_sign"),
23 ]
24
25 # After fine-tuning with LoRA, all generated content should start begin `A`.
26 EXPECTED_OUTPUT = [
27 "A red and white stop sign with a Chinese archway in the background featuring red lanterns and gold accents.", # noqa: E501
28 ]
29
30
31 def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
32 sampling_params = vllm.SamplingParams(
33 temperature=0,
34 max_tokens=5,
35 stop_token_ids=[128001, 128009], # eos_id, eot_id
36 )
37
38 inputs = [
39 {
40 "prompt": PROMPT_TEMPLATE,
41 "multi_modal_data": {"image": asset.pil_image},
42 }
43 for asset in IMAGE_ASSETS
44 ]
45
46 outputs = llm.generate(
47 inputs,
48 sampling_params,
49 lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
50 )
51 # Print the outputs.
52 generated_texts: list[str] = []
53 for output in outputs:
54 generated_text = output.outputs[0].text.strip()
55 generated_texts.append(generated_text)
56 print(f"Generated text: {generated_text!r}")
57 return generated_texts
58
59
60 @pytest.mark.skipif(
61 current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
62 )
63 def test_minicpmv_lora(minicpmv_lora_files):
64 llm = vllm.LLM(
65 MODEL_PATH,
66 max_num_seqs=2,
67 enable_lora=True,
68 max_loras=2,
69 max_lora_rank=8,
70 enforce_eager=True,
71 max_model_len=2048,
72 limit_mm_per_prompt={"image": 2, "video": 0},
73 trust_remote_code=True,
74 )
75 output1 = do_sample(llm, minicpmv_lora_files, lora_id=1)
76 for i in range(len(EXPECTED_OUTPUT)):
77 assert EXPECTED_OUTPUT[i].startswith(output1[i])
78 output2 = do_sample(llm, minicpmv_lora_files, lora_id=2)
79 for i in range(len(EXPECTED_OUTPUT)):
80 assert EXPECTED_OUTPUT[i].startswith(output2[i])
81
82
83 @pytest.mark.skipif(
84 current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
85 )
86 @multi_gpu_test(num_gpus=4)
87 def test_minicpmv_tp4_wo_fully_sharded_loras(minicpmv_lora_files):
88 llm = vllm.LLM(
89 MODEL_PATH,
90 enable_lora=True,
91 max_num_seqs=2,
92 max_loras=4,
93 max_lora_rank=64,
94 tensor_parallel_size=4,
95 limit_mm_per_prompt={"image": 2, "video": 0},
96 trust_remote_code=True,
97 )
98 output_tp = do_sample(llm, minicpmv_lora_files, lora_id=1)
99 for i in range(len(EXPECTED_OUTPUT)):
100 assert EXPECTED_OUTPUT[i].startswith(output_tp[i])
101
102
103 @pytest.mark.skipif(
104 current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
105 )
106 @multi_gpu_test(num_gpus=4)
107 def test_minicpmv_tp4_fully_sharded_loras(minicpmv_lora_files):
108 llm = vllm.LLM(
109 MODEL_PATH,
110 enable_lora=True,
111 max_num_seqs=2,
112 max_loras=2,
113 max_lora_rank=8,
114 tensor_parallel_size=4,
115 trust_remote_code=True,
116 limit_mm_per_prompt={"image": 1, "video": 0},
117 fully_sharded_loras=True,
118 )
119 output_tp = do_sample(llm, minicpmv_lora_files, lora_id=1)
120 for i in range(len(EXPECTED_OUTPUT)):
121 assert EXPECTED_OUTPUT[i].startswith(output_tp[i])
122 output_tp = do_sample(llm, minicpmv_lora_files, lora_id=2)
123 for i in range(len(EXPECTED_OUTPUT)):
124 assert EXPECTED_OUTPUT[i].startswith(output_tp[i])