@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 132 lines 5.32 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 import vllm.config
8 from vllm.lora.request import LoRARequest
9 from vllm.platforms import current_platform
10
11 from ..utils import create_new_process_for_each_test, multi_gpu_test
12
13 MODEL_PATH = "zai-org/chatglm3-6b"
14
15 PROMPT_TEMPLATE = """I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.\n"\n##Instruction:\nconcert_singer contains tables such as stadium, singer, concert, singer_in_concert. Table stadium has columns such as Stadium_ID, Location, Name, Capacity, Highest, Lowest, Average. Stadium_ID is the primary key.\nTable singer has columns such as Singer_ID, Name, Country, Song_Name, Song_release_year, Age, Is_male. Singer_ID is the primary key.\nTable concert has columns such as concert_ID, concert_Name, Theme, Stadium_ID, Year. concert_ID is the primary key.\nTable singer_in_concert has columns such as concert_ID, Singer_ID. concert_ID is the primary key.\nThe Stadium_ID of concert is the foreign key of Stadium_ID of stadium.\nThe Singer_ID of singer_in_concert is the foreign key of Singer_ID of singer.\nThe concert_ID of singer_in_concert is the foreign key of concert_ID of concert.\n\n###Input:\n{query}\n\n###Response:""" # noqa: E501
16
17 EXPECTED_LORA_OUTPUT = [
18 "SELECT count(*) FROM singer",
19 "SELECT avg(age) , min(age) , max(age) FROM singer WHERE country = 'France'",
20 "SELECT name , country , age FROM singer ORDER BY age",
21 ]
22
23
24 def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
25 prompts = [
26 PROMPT_TEMPLATE.format(query="How many singers do we have?"),
27 PROMPT_TEMPLATE.format(
28 query=(
29 "What is the average, minimum, and maximum "
30 "age of all singers from France?"
31 )
32 ),
33 PROMPT_TEMPLATE.format(
34 query=(
35 "Show name, country, age for all singers ordered "
36 "by age from the oldest to the youngest."
37 )
38 ),
39 ]
40 sampling_params = vllm.SamplingParams(temperature=0, max_tokens=32)
41 outputs = llm.generate(
42 prompts,
43 sampling_params,
44 lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
45 )
46 # Print the outputs.
47 generated_texts: list[str] = []
48 for output in outputs:
49 prompt = output.prompt
50 generated_text = output.outputs[0].text.strip()
51 generated_texts.append(generated_text)
52 print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
53 return generated_texts
54
55
56 @pytest.mark.skipif(
57 current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
58 )
59 @create_new_process_for_each_test()
60 def test_chatglm3_lora(chatglm3_lora_files):
61 llm = vllm.LLM(
62 MODEL_PATH,
63 max_model_len=512,
64 enable_lora=True,
65 max_loras=2,
66 max_num_seqs=16,
67 max_lora_rank=64,
68 trust_remote_code=True,
69 )
70
71 output1 = do_sample(llm, chatglm3_lora_files, lora_id=1)
72 for i in range(len(EXPECTED_LORA_OUTPUT)):
73 assert output1[i] == EXPECTED_LORA_OUTPUT[i]
74 output2 = do_sample(llm, chatglm3_lora_files, lora_id=2)
75 for i in range(len(EXPECTED_LORA_OUTPUT)):
76 assert output2[i] == EXPECTED_LORA_OUTPUT[i]
77
78
79 @pytest.mark.skipif(
80 current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
81 )
82 @multi_gpu_test(num_gpus=4)
83 def test_chatglm3_lora_tp4(chatglm3_lora_files):
84 llm = vllm.LLM(
85 MODEL_PATH,
86 max_model_len=512,
87 enable_lora=True,
88 max_loras=2,
89 max_lora_rank=64,
90 max_num_seqs=16,
91 tensor_parallel_size=4,
92 trust_remote_code=True,
93 fully_sharded_loras=False,
94 compilation_config=vllm.config.CompilationConfig( # Avoid OOM
95 cudagraph_specialize_lora=False,
96 ),
97 )
98
99 output1 = do_sample(llm, chatglm3_lora_files, lora_id=1)
100 for i in range(len(EXPECTED_LORA_OUTPUT)):
101 assert output1[i] == EXPECTED_LORA_OUTPUT[i]
102 output2 = do_sample(llm, chatglm3_lora_files, lora_id=2)
103 for i in range(len(EXPECTED_LORA_OUTPUT)):
104 assert output2[i] == EXPECTED_LORA_OUTPUT[i]
105
106
107 @multi_gpu_test(num_gpus=4)
108 def test_chatglm3_lora_tp4_fully_sharded_loras(chatglm3_lora_files):
109 # https://github.com/NVIDIA/nccl/issues/1790, set a lower value for
110 # gpu_memory_utilization here because NCCL >= 2.26.3 seems to use
111 # more GPU memory causing vLLM to OOM
112 llm = vllm.LLM(
113 MODEL_PATH,
114 max_model_len=512,
115 enable_lora=True,
116 max_loras=2,
117 max_lora_rank=64,
118 max_num_seqs=16,
119 tensor_parallel_size=4,
120 trust_remote_code=True,
121 fully_sharded_loras=True,
122 gpu_memory_utilization=0.8,
123 compilation_config=vllm.config.CompilationConfig( # Avoid OOM
124 cudagraph_specialize_lora=False,
125 ),
126 )
127 output1 = do_sample(llm, chatglm3_lora_files, lora_id=1)
128 for i in range(len(EXPECTED_LORA_OUTPUT)):
129 assert output1[i] == EXPECTED_LORA_OUTPUT[i]
130 output2 = do_sample(llm, chatglm3_lora_files, lora_id=2)
131 for i in range(len(EXPECTED_LORA_OUTPUT)):
132 assert output2[i] == EXPECTED_LORA_OUTPUT[i]