@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 125 lines 4.02 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """This example shows how to use LoRA with different quantization techniques
4 for offline inference.
5
6 Requires HuggingFace credentials for access.
7 The bitsandbytes example also requires installing vllm-bnb-plugin.
8 """
9
10 import gc
11
12 import torch
13 from huggingface_hub import snapshot_download
14
15 from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams
16 from vllm.lora.request import LoRARequest
17
18
19 def create_test_prompts(
20 lora_path: str,
21 ) -> list[tuple[str, SamplingParams, LoRARequest | None]]:
22 return [
23 # this is an example of using quantization without LoRA
24 (
25 "My name is",
26 SamplingParams(temperature=0.0, logprobs=1, max_tokens=128),
27 None,
28 ),
29 # the next three examples use quantization with LoRA
30 (
31 "my name is",
32 SamplingParams(temperature=0.0, logprobs=1, max_tokens=128),
33 LoRARequest("lora-test-1", 1, lora_path),
34 ),
35 (
36 "The capital of USA is",
37 SamplingParams(temperature=0.0, logprobs=1, max_tokens=128),
38 LoRARequest("lora-test-2", 1, lora_path),
39 ),
40 (
41 "The capital of France is",
42 SamplingParams(temperature=0.0, logprobs=1, max_tokens=128),
43 LoRARequest("lora-test-3", 1, lora_path),
44 ),
45 ]
46
47
48 def process_requests(
49 engine: LLMEngine,
50 test_prompts: list[tuple[str, SamplingParams, LoRARequest | None]],
51 ):
52 """Continuously process a list of prompts and handle the outputs."""
53 request_id = 0
54
55 while test_prompts or engine.has_unfinished_requests():
56 if test_prompts:
57 prompt, sampling_params, lora_request = test_prompts.pop(0)
58 engine.add_request(
59 str(request_id), prompt, sampling_params, lora_request=lora_request
60 )
61 request_id += 1
62
63 request_outputs: list[RequestOutput] = engine.step()
64 for request_output in request_outputs:
65 if request_output.finished:
66 print("----------------------------------------------------")
67 print(f"Prompt: {request_output.prompt}")
68 print(f"Output: {request_output.outputs[0].text}")
69
70
71 def initialize_engine(
72 model: str, quantization: str, lora_repo: str | None
73 ) -> LLMEngine:
74 """Initialize the LLMEngine."""
75 engine_args = EngineArgs(
76 model=model,
77 quantization=quantization,
78 enable_lora=True,
79 max_lora_rank=64,
80 max_loras=4,
81 )
82 return LLMEngine.from_engine_args(engine_args)
83
84
85 def main():
86 """Main function that sets up and runs the prompt processing."""
87 test_configs = [
88 # QLoRA (https://arxiv.org/abs/2305.14314)
89 {
90 "name": "qlora_inference_example",
91 "model": "huggyllama/llama-7b",
92 "quantization": "bitsandbytes",
93 "lora_repo": "timdettmers/qlora-flan-7b",
94 },
95 {
96 "name": "AWQ_inference_with_lora_example",
97 "model": "TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ",
98 "quantization": "awq",
99 "lora_repo": "jashing/tinyllama-colorist-lora",
100 },
101 {
102 "name": "GPTQ_inference_with_lora_example",
103 "model": "TheBloke/TinyLlama-1.1B-Chat-v0.3-GPTQ",
104 "quantization": "gptq",
105 "lora_repo": "jashing/tinyllama-colorist-lora",
106 },
107 ]
108
109 for test_config in test_configs:
110 print(f"~~~~~~~~~~~~~~~~ Running: {test_config['name']} ~~~~~~~~~~~~~~~~")
111 engine = initialize_engine(
112 test_config["model"], test_config["quantization"], test_config["lora_repo"]
113 )
114 lora_path = snapshot_download(repo_id=test_config["lora_repo"])
115 test_prompts = create_test_prompts(lora_path)
116 process_requests(engine, test_prompts)
117
118 # Clean up the GPU memory for the next test
119 del engine
120 gc.collect()
121 torch.accelerator.empty_cache()
122
123
124 if __name__ == "__main__":
125 main()