| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This example shows how to use the multi-LoRA functionality |
| 4 | for offline inference. |
| 5 | |
| 6 | Requires HuggingFace credentials for access to Llama2. |
| 7 | """ |
| 8 | |
| 9 | from huggingface_hub import snapshot_download |
| 10 | |
| 11 | from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams |
| 12 | from vllm.lora.request import LoRARequest |
| 13 | |
| 14 | |
| 15 | def create_test_prompts( |
| 16 | lora_path: str, |
| 17 | ) -> list[tuple[str, SamplingParams, LoRARequest | None]]: |
| 18 | """Create a list of test prompts with their sampling parameters. |
| 19 | |
| 20 | 2 requests for base model, 4 requests for the LoRA. We define 2 |
| 21 | different LoRA adapters (using the same model for demo purposes). |
| 22 | Since we also set `max_loras=1`, the expectation is that the requests |
| 23 | with the second LoRA adapter will be run after all requests with the |
| 24 | first adapter have finished. |
| 25 | """ |
| 26 | return [ |
| 27 | ( |
| 28 | "A robot may not injure a human being", |
| 29 | SamplingParams(temperature=0.0, logprobs=1, max_tokens=128), |
| 30 | None, |
| 31 | ), |
| 32 | ( |
| 33 | "To be or not to be,", |
| 34 | SamplingParams( |
| 35 | temperature=0.8, top_k=5, presence_penalty=0.2, max_tokens=128 |
| 36 | ), |
| 37 | None, |
| 38 | ), |
| 39 | ( |
| 40 | "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]", # noqa: E501 |
| 41 | SamplingParams(temperature=0.0, logprobs=1, max_tokens=128), |
| 42 | LoRARequest("sql-lora", 1, lora_path), |
| 43 | ), |
| 44 | ( |
| 45 | "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]", # noqa: E501 |
| 46 | SamplingParams(temperature=0.0, logprobs=1, max_tokens=128), |
| 47 | LoRARequest("sql-lora2", 2, lora_path), |
| 48 | ), |
| 49 | ] |
| 50 | |
| 51 | |
| 52 | def process_requests( |
| 53 | engine: LLMEngine, |
| 54 | test_prompts: list[tuple[str, SamplingParams, LoRARequest | None]], |
| 55 | ): |
| 56 | """Continuously process a list of prompts and handle the outputs.""" |
| 57 | request_id = 0 |
| 58 | |
| 59 | print("-" * 50) |
| 60 | while test_prompts or engine.has_unfinished_requests(): |
| 61 | if test_prompts: |
| 62 | prompt, sampling_params, lora_request = test_prompts.pop(0) |
| 63 | engine.add_request( |
| 64 | str(request_id), prompt, sampling_params, lora_request=lora_request |
| 65 | ) |
| 66 | request_id += 1 |
| 67 | |
| 68 | request_outputs: list[RequestOutput] = engine.step() |
| 69 | |
| 70 | for request_output in request_outputs: |
| 71 | if request_output.finished: |
| 72 | print(request_output) |
| 73 | print("-" * 50) |
| 74 | |
| 75 | |
| 76 | def initialize_engine() -> LLMEngine: |
| 77 | """Initialize the LLMEngine.""" |
| 78 | # max_loras: controls the number of LoRAs that can be used in the same |
| 79 | # batch. Larger numbers will cause higher memory usage, as each LoRA |
| 80 | # slot requires its own preallocated tensor. |
| 81 | # max_lora_rank: controls the maximum supported rank of all LoRAs. Larger |
| 82 | # numbers will cause higher memory usage. If you know that all LoRAs will |
| 83 | # use the same rank, it is recommended to set this as low as possible. |
| 84 | # max_cpu_loras: controls the size of the CPU LoRA cache. |
| 85 | engine_args = EngineArgs( |
| 86 | model="meta-llama/Llama-3.2-3B-Instruct", |
| 87 | enable_lora=True, |
| 88 | max_loras=1, |
| 89 | max_lora_rank=8, |
| 90 | max_cpu_loras=2, |
| 91 | max_num_seqs=256, |
| 92 | ) |
| 93 | return LLMEngine.from_engine_args(engine_args) |
| 94 | |
| 95 | |
| 96 | def main(): |
| 97 | """Main function that sets up and runs the prompt processing.""" |
| 98 | engine = initialize_engine() |
| 99 | lora_path = snapshot_download(repo_id="jeeejeee/llama32-3b-text2sql-spider") |
| 100 | test_prompts = create_test_prompts(lora_path) |
| 101 | process_requests(engine, test_prompts) |
| 102 | |
| 103 | |
| 104 | if __name__ == "__main__": |
| 105 | main() |