| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import subprocess |
| 4 | import sys |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | import vllm |
| 9 | import vllm.config |
| 10 | from vllm import LLM |
| 11 | from vllm.lora.request import LoRARequest |
| 12 | from vllm.model_executor.model_loader.tensorizer import TensorizerConfig |
| 13 | from vllm.platforms import current_platform |
| 14 | |
| 15 | from ..utils import VLLM_PATH, create_new_process_for_each_test, multi_gpu_test |
| 16 | |
| 17 | PROMPT_TEMPLATE = """<|eot_id|><|start_header_id|>user<|end_header_id|> |
| 18 | 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. |
| 19 | " |
| 20 | ##Instruction: |
| 21 | candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key. |
| 22 | Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key. |
| 23 | The People_ID of candidate is the foreign key of People_ID of people. |
| 24 | ###Input: |
| 25 | {context} |
| 26 | ###Response:<|eot_id|><|start_header_id|>assistant<|end_header_id|> |
| 27 | """ # noqa: E501 |
| 28 | |
| 29 | EXPECTED_LORA_OUTPUT = [ |
| 30 | "SELECT count(*) FROM candidate", |
| 31 | "SELECT count(*) FROM candidate", |
| 32 | "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501 |
| 33 | "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501 |
| 34 | ] |
| 35 | |
| 36 | MODEL_PATH = "meta-llama/Llama-3.2-3B-Instruct" |
| 37 | |
| 38 | |
| 39 | def do_sample( |
| 40 | llm: vllm.LLM, |
| 41 | lora_path: str, |
| 42 | lora_id: int, |
| 43 | tensorizer_config_dict: dict | None = None, |
| 44 | ) -> list[str]: |
| 45 | prompts = [ |
| 46 | PROMPT_TEMPLATE.format(context="How many candidates are there?"), |
| 47 | PROMPT_TEMPLATE.format(context="Count the number of candidates."), |
| 48 | PROMPT_TEMPLATE.format( |
| 49 | context="Which poll resource provided the most number of candidate information?" # noqa: E501 |
| 50 | ), |
| 51 | PROMPT_TEMPLATE.format( |
| 52 | context="Return the poll resource associated with the most candidates." |
| 53 | ), |
| 54 | ] |
| 55 | |
| 56 | sampling_params = vllm.SamplingParams( |
| 57 | temperature=0, max_tokens=64, stop=["<|im_end|>"] |
| 58 | ) |
| 59 | if tensorizer_config_dict is not None: |
| 60 | outputs = llm.generate( |
| 61 | prompts, |
| 62 | sampling_params, |
| 63 | lora_request=LoRARequest( |
| 64 | str(lora_id), |
| 65 | lora_id, |
| 66 | lora_path, |
| 67 | tensorizer_config_dict=tensorizer_config_dict, |
| 68 | ) |
| 69 | if lora_id |
| 70 | else None, |
| 71 | ) |
| 72 | else: |
| 73 | outputs = llm.generate( |
| 74 | prompts, |
| 75 | sampling_params, |
| 76 | lora_request=LoRARequest(str(lora_id), lora_id, lora_path) |
| 77 | if lora_id |
| 78 | else None, |
| 79 | ) |
| 80 | lora_request = LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None |
| 81 | generated_texts: list[str] = [] |
| 82 | for output in outputs: |
| 83 | prompt = output.prompt |
| 84 | generated_text = output.outputs[0].text |
| 85 | # The output should include correct lora_request info |
| 86 | if lora_request is not None: |
| 87 | assert output.lora_request.lora_name == lora_request.lora_name |
| 88 | assert output.lora_request.lora_int_id == lora_request.lora_int_id |
| 89 | assert output.lora_request.lora_path == lora_request.lora_path |
| 90 | else: |
| 91 | assert output.lora_request is None |
| 92 | generated_texts.append(generated_text) |
| 93 | print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") |
| 94 | return generated_texts |
| 95 | |
| 96 | |
| 97 | def generate_and_test( |
| 98 | llm, llama32_lora_files, tensorizer_config_dict: dict | None = None |
| 99 | ): |
| 100 | print("lora adapter created") |
| 101 | print("lora 1") |
| 102 | assert ( |
| 103 | do_sample( |
| 104 | llm, |
| 105 | llama32_lora_files, |
| 106 | tensorizer_config_dict=tensorizer_config_dict, |
| 107 | lora_id=1, |
| 108 | ) |
| 109 | == EXPECTED_LORA_OUTPUT |
| 110 | ) |
| 111 | |
| 112 | print("lora 2") |
| 113 | assert ( |
| 114 | do_sample( |
| 115 | llm, |
| 116 | llama32_lora_files, |
| 117 | tensorizer_config_dict=tensorizer_config_dict, |
| 118 | lora_id=2, |
| 119 | ) |
| 120 | == EXPECTED_LORA_OUTPUT |
| 121 | ) |
| 122 | |
| 123 | print("removing lora") |
| 124 | |
| 125 | |
| 126 | @create_new_process_for_each_test() |
| 127 | @pytest.mark.parametrize("cudagraph_specialize_lora", [True, False]) |
| 128 | def test_llama_lora(llama32_lora_files, cudagraph_specialize_lora: bool): |
| 129 | llm = vllm.LLM( |
| 130 | MODEL_PATH, |
| 131 | enable_lora=True, |
| 132 | # also test odd max_num_seqs |
| 133 | max_num_seqs=7, |
| 134 | max_model_len=1024, |
| 135 | max_loras=4, |
| 136 | compilation_config=vllm.config.CompilationConfig( |
| 137 | cudagraph_specialize_lora=cudagraph_specialize_lora, |
| 138 | ), |
| 139 | ) |
| 140 | generate_and_test(llm, llama32_lora_files) |
| 141 | |
| 142 | |
| 143 | @pytest.mark.skipif( |
| 144 | current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests" |
| 145 | ) |
| 146 | @multi_gpu_test(num_gpus=4) |
| 147 | def test_llama_lora_tp4(llama32_lora_files): |
| 148 | llm = vllm.LLM( |
| 149 | MODEL_PATH, |
| 150 | enable_lora=True, |
| 151 | max_num_seqs=7, |
| 152 | max_model_len=1024, |
| 153 | max_loras=4, |
| 154 | tensor_parallel_size=4, |
| 155 | ) |
| 156 | generate_and_test(llm, llama32_lora_files) |
| 157 | |
| 158 | |
| 159 | @multi_gpu_test(num_gpus=4) |
| 160 | def test_llama_lora_tp4_fully_sharded_loras(llama32_lora_files): |
| 161 | llm = vllm.LLM( |
| 162 | MODEL_PATH, |
| 163 | enable_lora=True, |
| 164 | max_num_seqs=8, |
| 165 | max_loras=4, |
| 166 | max_model_len=1024, |
| 167 | tensor_parallel_size=4, |
| 168 | fully_sharded_loras=True, |
| 169 | ) |
| 170 | generate_and_test(llm, llama32_lora_files) |
| 171 | |
| 172 | |
| 173 | @multi_gpu_test(num_gpus=2) |
| 174 | def test_tp2_serialize_and_deserialize_lora( |
| 175 | tmp_path, |
| 176 | llama32_lora_files, |
| 177 | ): |
| 178 | # Run the tensorizing of the LoRA adapter and the model in a subprocess |
| 179 | # to guarantee cleanup |
| 180 | |
| 181 | tp_size = 2 |
| 182 | model_name = "model-rank-%03d.tensors" |
| 183 | |
| 184 | model_ref = MODEL_PATH |
| 185 | lora_path = llama32_lora_files |
| 186 | suffix = "test" |
| 187 | try: |
| 188 | result = subprocess.run( |
| 189 | [ |
| 190 | sys.executable, |
| 191 | f"{VLLM_PATH}/examples/features/tensorize_vllm_model.py", |
| 192 | "--model", |
| 193 | MODEL_PATH, |
| 194 | "--lora-path", |
| 195 | lora_path, |
| 196 | "--tensor-parallel-size", |
| 197 | str(tp_size), |
| 198 | "serialize", |
| 199 | "--serialized-directory", |
| 200 | str(tmp_path), |
| 201 | "--suffix", |
| 202 | suffix, |
| 203 | "--serialization-kwargs", |
| 204 | '{"limit_cpu_concurrency": 4}', |
| 205 | ], |
| 206 | check=True, |
| 207 | capture_output=True, |
| 208 | text=True, |
| 209 | ) |
| 210 | except subprocess.CalledProcessError as e: |
| 211 | print("Tensorizing failed.") |
| 212 | print("STDOUT:\n", e.stdout) |
| 213 | print("STDERR:\n", e.stderr) |
| 214 | raise |
| 215 | |
| 216 | print("STDOUT:\n", result.stdout) |
| 217 | |
| 218 | model_uri = tmp_path / "vllm" / model_ref / suffix / model_name |
| 219 | tensorizer_config = TensorizerConfig(tensorizer_uri=str(model_uri)) |
| 220 | |
| 221 | loaded_llm = LLM( |
| 222 | model=model_ref, |
| 223 | load_format="tensorizer", |
| 224 | enable_lora=True, |
| 225 | enforce_eager=True, |
| 226 | model_loader_extra_config=tensorizer_config, |
| 227 | max_num_seqs=7, |
| 228 | max_model_len=1024, |
| 229 | tensor_parallel_size=2, |
| 230 | max_loras=2, |
| 231 | # Leave headroom for LoRA adapter loading and Triton JIT |
| 232 | # compilation, which can allocate GPU memory concurrently |
| 233 | # during the first inference step. |
| 234 | gpu_memory_utilization=0.85, |
| 235 | ) |
| 236 | |
| 237 | tc_as_dict = tensorizer_config.to_serializable() |
| 238 | |
| 239 | print("lora adapter created") |
| 240 | print("lora 1") |
| 241 | assert ( |
| 242 | do_sample( |
| 243 | loaded_llm, llama32_lora_files, tensorizer_config_dict=tc_as_dict, lora_id=1 |
| 244 | ) |
| 245 | == EXPECTED_LORA_OUTPUT |
| 246 | ) |