@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 03f48b3db6
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 22:15 UTC

[Core] LoRA V1 - Add add/pin/list/remove_lora functions (#13705)

Varun Sundar Rabindranath committed Feb 25, 2025 at 13:48 UTC 03f48b3db6cff93ffff3ca440383aea1fe4b1a4c
8 files changed +270 -22
tests/lora/test_add_lora.py
+9 -4
@@ -7,6 +7,7 @@ from typing import List
7 import pytest
8 from huggingface_hub import snapshot_download
9
10 +import vllm.envs as env
11 from vllm.engine.arg_utils import AsyncEngineArgs
12 from vllm.inputs import TextPrompt
13 from vllm.lora.request import LoRARequest
@@ -144,10 +145,14 @@ async def test_add_lora():
145 await requests_processing_time(llm, dummy_run_requests)
146
147 # Run with warmup
147 - for lr in warmup_run_requests:
148 - await llm.add_lora(lr)
149 - # Wait for the add_lora function to complete on the server side.
150 - await asyncio.sleep(30)
148 + add_lora_tasks = [llm.add_lora(lr) for lr in warmup_run_requests]
149 + add_lora_results = await asyncio.gather(*add_lora_tasks)
150 + if env.VLLM_USE_V1:
151 + # Test that all all_lora calls are successful.
152 + assert all(add_lora_results)
153 + else:
154 + # No way to check V0 engine results as the calls just return None.
155 + pass
156 time_with_add_lora = await requests_processing_time(
157 llm, warmup_run_requests)
158
tests/lora/test_lora_functions.py new
+137
@@ -0,0 +1,137 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +Script to test add_lora, remove_lora, pin_lora, list_loras functions.
4 +"""
5 +
6 +import os
7 +from typing import List
8 +
9 +import pytest
10 +
11 +from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
12 +from vllm.entrypoints.llm import LLM
13 +from vllm.lora.request import LoRARequest
14 +
15 +MODEL_PATH = "meta-llama/Llama-2-7b-hf"
16 +LORA_MODULE_PATH = "yard1/llama-2-7b-sql-lora-test"
17 +LORA_RANK = 8
18 +
19 +
20 +@pytest.fixture(autouse=True)
21 +def v1(run_with_both_engines_lora):
22 + # Simple autouse wrapper to run both engines for each test
23 + # This can be promoted up to conftest.py to run for every
24 + # test in a package
25 + pass
26 +
27 +
28 +def make_lora_request(lora_id: int):
29 + return LoRARequest(lora_name=f"{lora_id}",
30 + lora_int_id=lora_id,
31 + lora_path=LORA_MODULE_PATH)
32 +
33 +
34 +def test_lora_functions_sync():
35 +
36 + max_loras = 4
37 + # Create engine in eager-mode. Due to high max_loras, the CI can
38 + # OOM during cuda-graph capture.
39 + engine_args = EngineArgs(model=MODEL_PATH,
40 + enable_lora=True,
41 + max_loras=max_loras,
42 + max_lora_rank=LORA_RANK,
43 + max_model_len=128,
44 + gpu_memory_utilization=0.8,
45 + enforce_eager=True)
46 +
47 + llm = LLM.get_engine_class().from_engine_args(engine_args)
48 +
49 + def run_check(fn, args, expected: List):
50 + fn(args)
51 + assert set(llm.list_loras()) == set(expected)
52 +
53 + run_check(llm.add_lora, make_lora_request(1), [1])
54 + run_check(llm.add_lora, make_lora_request(2), [1, 2])
55 +
56 + # Pin LoRA 1 and test that it is never removed on subsequent adds.
57 + run_check(llm.pin_lora, 1, [1, 2])
58 + run_check(llm.add_lora, make_lora_request(3), [1, 2, 3])
59 + run_check(llm.add_lora, make_lora_request(4), [1, 2, 3, 4])
60 + run_check(llm.add_lora, make_lora_request(5), [1, 5, 3, 4])
61 + run_check(llm.add_lora, make_lora_request(6), [1, 5, 6, 4])
62 + run_check(llm.add_lora, make_lora_request(7), [1, 5, 6, 7])
63 + run_check(llm.add_lora, make_lora_request(8), [1, 8, 6, 7])
64 + run_check(llm.add_lora, make_lora_request(9), [1, 8, 9, 7])
65 + run_check(llm.add_lora, make_lora_request(10), [1, 8, 9, 10])
66 +
67 + # Remove LoRA 1 and continue adding.
68 + run_check(llm.remove_lora, 1, [8, 9, 10])
69 + run_check(llm.add_lora, make_lora_request(11), [8, 9, 10, 11])
70 + run_check(llm.add_lora, make_lora_request(12), [12, 9, 10, 11])
71 + run_check(llm.add_lora, make_lora_request(13), [12, 13, 10, 11])
72 +
73 + # Remove all LoRAs
74 + run_check(llm.remove_lora, 13, [12, 10, 11])
75 + run_check(llm.remove_lora, 12, [10, 11])
76 + run_check(llm.remove_lora, 11, [10])
77 + run_check(llm.remove_lora, 10, [])
78 +
79 +
80 +@pytest.mark.asyncio
81 +async def test_lora_functions_async():
82 +
83 + if os.getenv("VLLM_USE_V1") == "0":
84 + pytest.skip(
85 + reason=
86 + "V0 AsyncLLMEngine does not expose remove/list/pin LoRA functions")
87 +
88 + # The run_with_both_engines_lora fixture sets up the `VLLM_USE_V1`
89 + # environment variable. reload vllm.enging.async_llm_engine as
90 + # vllm.engine.async_llm_engine.AsyncLLMEgnine changes depending on the
91 + # env var.
92 + import importlib
93 +
94 + import vllm.engine.async_llm_engine
95 + importlib.reload(vllm.engine.async_llm_engine)
96 + from vllm.entrypoints.openai.api_server import (
97 + build_async_engine_client_from_engine_args)
98 +
99 + max_loras = 4
100 + engine_args = AsyncEngineArgs(model=MODEL_PATH,
101 + enable_lora=True,
102 + max_loras=max_loras,
103 + max_lora_rank=LORA_RANK,
104 + max_model_len=128,
105 + gpu_memory_utilization=0.8,
106 + enforce_eager=True)
107 +
108 + async def run_check(fn, args, expected: List):
109 + await fn(args)
110 + assert set(await llm.list_loras()) == set(expected)
111 +
112 + async with build_async_engine_client_from_engine_args(engine_args) as llm:
113 + await run_check(llm.add_lora, make_lora_request(1), [1])
114 + await run_check(llm.add_lora, make_lora_request(2), [1, 2])
115 +
116 + # Pin LoRA 1 and test that it is never removed on subsequent adds.
117 + await run_check(llm.pin_lora, 1, [1, 2])
118 + await run_check(llm.add_lora, make_lora_request(3), [1, 2, 3])
119 + await run_check(llm.add_lora, make_lora_request(4), [1, 2, 3, 4])
120 + await run_check(llm.add_lora, make_lora_request(5), [1, 5, 3, 4])
121 + await run_check(llm.add_lora, make_lora_request(6), [1, 5, 6, 4])
122 + await run_check(llm.add_lora, make_lora_request(7), [1, 5, 6, 7])
123 + await run_check(llm.add_lora, make_lora_request(8), [1, 8, 6, 7])
124 + await run_check(llm.add_lora, make_lora_request(9), [1, 8, 9, 7])
125 + await run_check(llm.add_lora, make_lora_request(10), [1, 8, 9, 10])
126 +
127 + # Remove LoRA 1 and continue adding.
128 + await run_check(llm.remove_lora, 1, [8, 9, 10])
129 + await run_check(llm.add_lora, make_lora_request(11), [8, 9, 10, 11])
130 + await run_check(llm.add_lora, make_lora_request(12), [12, 9, 10, 11])
131 + await run_check(llm.add_lora, make_lora_request(13), [12, 13, 10, 11])
132 +
133 + # Remove all LoRAs
134 + await run_check(llm.remove_lora, 13, [12, 10, 11])
135 + await run_check(llm.remove_lora, 12, [10, 11])
136 + await run_check(llm.remove_lora, 11, [10])
137 + await run_check(llm.remove_lora, 10, [])
vllm/v1/engine/async_llm.py
+15 -3
@@ -2,7 +2,7 @@
2
3 import asyncio
4 import os
5 -from typing import AsyncGenerator, List, Mapping, Optional, Type, Union
5 +from typing import AsyncGenerator, List, Mapping, Optional, Set, Type, Union
6
7 import numpy as np
8
@@ -392,9 +392,21 @@ class AsyncLLM(EngineClient):
392 async def wake_up(self) -> None:
393 await self.engine_core.wake_up_async()
394
395 - async def add_lora(self, lora_request: LoRARequest) -> None:
395 + async def add_lora(self, lora_request: LoRARequest) -> bool:
396 """Load a new LoRA adapter into the engine for future requests."""
397 - await self.engine_core.add_lora_async(lora_request)
397 + return await self.engine_core.add_lora_async(lora_request)
398 +
399 + async def remove_lora(self, lora_id: int) -> bool:
400 + """Remove an already loaded LoRA adapter."""
401 + return await self.engine_core.remove_lora_async(lora_id)
402 +
403 + async def list_loras(self) -> Set[int]:
404 + """List all registered adapters."""
405 + return await self.engine_core.list_loras_async()
406 +
407 + async def pin_lora(self, lora_id: int) -> bool:
408 + """Prevent an adapter from being evicted."""
409 + return await self.engine_core.pin_lora_async(lora_id)
410
411 @property
412 def is_running(self) -> bool:
vllm/v1/engine/core.py
+12 -3
@@ -7,7 +7,7 @@ import time
7 from concurrent.futures import Future
8 from inspect import isclass, signature
9 from multiprocessing.connection import Connection
10 -from typing import Any, List, Optional, Tuple, Type
10 +from typing import Any, List, Optional, Set, Tuple, Type
11
12 import msgspec
13 import psutil
@@ -222,8 +222,17 @@ class EngineCore:
222 def execute_dummy_batch(self):
223 self.model_executor.collective_rpc("execute_dummy_batch")
224
225 - def add_lora(self, lora_request: LoRARequest) -> None:
226 - self.model_executor.add_lora(lora_request)
225 + def add_lora(self, lora_request: LoRARequest) -> bool:
226 + return self.model_executor.add_lora(lora_request)
227 +
228 + def remove_lora(self, lora_id: int) -> bool:
229 + return self.model_executor.remove_lora(lora_id)
230 +
231 + def list_loras(self) -> Set[int]:
232 + return self.model_executor.list_loras()
233 +
234 + def pin_lora(self, lora_id: int) -> bool:
235 + return self.model_executor.pin_lora(lora_id)
236
237
238 class EngineCoreProc(EngineCore):
vllm/v1/engine/core_client.py
+54 -9
@@ -10,7 +10,7 @@ from abc import ABC, abstractmethod
10 from concurrent.futures import Future
11 from dataclasses import dataclass
12 from threading import Thread
13 -from typing import Any, Dict, List, Optional, Type, Union
13 +from typing import Any, Dict, List, Optional, Set, Type, Union
14
15 import zmq
16 import zmq.asyncio
@@ -97,7 +97,16 @@ class EngineCoreClient(ABC):
97 def abort_requests(self, request_ids: List[str]) -> None:
98 raise NotImplementedError
99
100 - def add_lora(self, lora_request: LoRARequest) -> None:
100 + def add_lora(self, lora_request: LoRARequest) -> bool:
101 + raise NotImplementedError
102 +
103 + def remove_lora(self, lora_id: int) -> bool:
104 + raise NotImplementedError
105 +
106 + def list_loras(self) -> Set[int]:
107 + raise NotImplementedError
108 +
109 + def pin_lora(self, lora_id: int) -> bool:
110 raise NotImplementedError
111
112 async def get_output_async(self) -> EngineCoreOutputs:
@@ -121,7 +130,16 @@ class EngineCoreClient(ABC):
130 async def abort_requests_async(self, request_ids: List[str]) -> None:
131 raise NotImplementedError
132
124 - async def add_lora_async(self, lora_request: LoRARequest) -> None:
133 + async def add_lora_async(self, lora_request: LoRARequest) -> bool:
134 + raise NotImplementedError
135 +
136 + async def remove_lora_async(self, lora_id: int) -> bool:
137 + raise NotImplementedError
138 +
139 + async def list_loras_async(self) -> Set[int]:
140 + raise NotImplementedError
141 +
142 + async def pin_lora_async(self, lora_id: int) -> bool:
143 raise NotImplementedError
144
145
@@ -166,8 +184,17 @@ class InprocClient(EngineCoreClient):
184 def execute_dummy_batch(self) -> None:
185 self.engine_core.execute_dummy_batch()
186
169 - def add_lora(self, lora_request: LoRARequest) -> None:
170 - self.engine_core.add_lora(lora_request)
187 + def add_lora(self, lora_request: LoRARequest) -> bool:
188 + return self.engine_core.add_lora(lora_request)
189 +
190 + def remove_lora(self, lora_id: int) -> bool:
191 + return self.engine_core.remove_lora(lora_id)
192 +
193 + def list_loras(self) -> Set[int]:
194 + return self.engine_core.list_loras()
195 +
196 + def pin_lora(self, lora_id: int) -> bool:
197 + return self.engine_core.pin_lora(lora_id)
198
199
200 @dataclass
@@ -356,8 +383,17 @@ class SyncMPClient(MPClient):
383 def reset_prefix_cache(self) -> None:
384 self._call_utility("reset_prefix_cache")
385
359 - def add_lora(self, lora_request: LoRARequest) -> None:
360 - self._call_utility("add_lora", lora_request)
386 + def add_lora(self, lora_request: LoRARequest) -> bool:
387 + return self._call_utility("add_lora", lora_request)
388 +
389 + def remove_lora(self, lora_id: int) -> bool:
390 + return self._call_utility("remove_lora", lora_id)
391 +
392 + def list_loras(self) -> Set[int]:
393 + return self._call_utility("list_loras")
394 +
395 + def pin_lora(self, lora_id: int) -> bool:
396 + return self._call_utility("pin_lora", lora_id)
397
398 def sleep(self, level: int = 1) -> None:
399 self._call_utility("sleep", level)
@@ -454,5 +490,14 @@ class AsyncMPClient(MPClient):
490 async def execute_dummy_batch_async(self) -> None:
491 await self._call_utility_async("execute_dummy_batch")
492
457 - async def add_lora_async(self, lora_request: LoRARequest) -> None:
458 - await self._call_utility_async("add_lora", lora_request)
493 + async def add_lora_async(self, lora_request: LoRARequest) -> bool:
494 + return await self._call_utility_async("add_lora", lora_request)
495 +
496 + async def remove_lora_async(self, lora_id: int) -> bool:
497 + return await self._call_utility_async("remove_lora", lora_id)
498 +
499 + async def list_loras_async(self) -> Set[int]:
500 + return await self._call_utility_async("list_loras")
501 +
502 + async def pin_lora_async(self, lora_id: int) -> bool:
503 + return await self._call_utility_async("pin_lora", lora_id)
vllm/v1/engine/llm_engine.py
+17 -1
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -from typing import Dict, List, Mapping, Optional, Type, Union
3 +from typing import Dict, List, Mapping, Optional, Set, Type, Union
4
5 from typing_extensions import TypeVar
6
@@ -254,3 +254,19 @@ class LLMEngine:
254 f"found type: {type(tokenizer_group)}")
255
256 return tokenizer_group
257 +
258 + def add_lora(self, lora_request: LoRARequest) -> bool:
259 + """Load a new LoRA adapter into the engine for future requests."""
260 + return self.engine_core.add_lora(lora_request)
261 +
262 + def remove_lora(self, lora_id: int) -> bool:
263 + """Remove an already loaded LoRA adapter."""
264 + return self.engine_core.remove_lora(lora_id)
265 +
266 + def list_loras(self) -> Set[int]:
267 + """List all registered adapters."""
268 + return self.engine_core.list_loras()
269 +
270 + def pin_lora(self, lora_id: int) -> bool:
271 + """Prevent an adapter from being evicted."""
272 + return self.engine_core.pin_lora(lora_id)
vllm/v1/worker/gpu_worker.py
+10 -1
@@ -2,7 +2,7 @@
2 """A GPU worker class."""
3 import gc
4 import os
5 -from typing import TYPE_CHECKING, Optional
5 +from typing import TYPE_CHECKING, Optional, Set
6
7 import torch
8 import torch.distributed
@@ -240,6 +240,15 @@ class Worker(WorkerBase):
240 def add_lora(self, lora_request: LoRARequest) -> bool:
241 return self.model_runner.add_lora(lora_request)
242
243 + def remove_lora(self, lora_id: int) -> bool:
244 + return self.model_runner.remove_lora(lora_id)
245 +
246 + def list_loras(self) -> Set[int]:
247 + return self.model_runner.list_loras()
248 +
249 + def pin_lora(self, lora_id: int) -> bool:
250 + return self.model_runner.pin_lora(lora_id)
251 +
252 def check_health(self) -> None:
253 # worker will always be healthy as long as it's running.
254 return
vllm/v1/worker/lora_model_runner_mixin.py
+16 -1
@@ -131,4 +131,19 @@ class LoRAModelRunnerMixin:
131 def add_lora(self, lora_request: LoRARequest) -> bool:
132 if not self.lora_manager:
133 raise RuntimeError("LoRA is not enabled.")
134 - return self.lora_manager.add_adapter(lora_request)
\ No newline at end of file
134 + return self.lora_manager.add_adapter(lora_request)
135 +
136 + def remove_lora(self, lora_id: int) -> bool:
137 + if not self.lora_manager:
138 + raise RuntimeError("LoRA is not enabled.")
139 + return self.lora_manager.remove_adapter(lora_id)
140 +
141 + def pin_lora(self, lora_id: int) -> bool:
142 + if not self.lora_manager:
143 + raise RuntimeError("LoRA is not enabled.")
144 + return self.lora_manager.pin_adapter(lora_id)
145 +
146 + def list_loras(self) -> Set[int]:
147 + if not self.lora_manager:
148 + raise RuntimeError("LoRA is not enabled.")
149 + return self.lora_manager.list_adapters()
\ No newline at end of file