@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,105 lines 38.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import os
5
6 import pytest
7 import torch
8 from safetensors.torch import load_file
9 from torch import nn
10
11 from vllm.config import ModelConfig, VllmConfig
12 from vllm.config.lora import LoRAConfig
13 from vllm.lora.layers import (
14 ColumnParallelLinearWithLoRA,
15 MergedColumnParallelLinearWithLoRA,
16 ReplicatedLinearWithLoRA,
17 RowParallelLinearWithLoRA,
18 )
19 from vllm.lora.lora_model import LoRAModel
20 from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
21 from vllm.lora.model_manager import (
22 DEFAULT_LANGUAGE_WRAPPER_KEY,
23 LoRAMapping,
24 LoRAModelManager,
25 LRUCacheLoRAModelManager,
26 )
27 from vllm.lora.peft_helper import PEFTHelper
28 from vllm.lora.request import LoRARequest
29 from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager, WorkerLoRAManager
30 from vllm.model_executor.layers.fused_moe import GateLinear
31 from vllm.platforms import current_platform
32
33 from .utils import create_peft_lora
34
35 EMBEDDING_MODULES = {
36 "embed_tokens": "input_embeddings",
37 "lm_head": "output_embeddings",
38 }
39
40 DEVICE_TYPE = current_platform.device_type
41 DEVICES = (
42 [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
43 if (current_platform.is_cuda_alike() or current_platform.is_xpu())
44 else ["cpu"]
45 )
46
47 DEFAULT_DTYPE = torch.get_default_dtype()
48
49
50 @pytest.mark.parametrize("device", DEVICES)
51 def test_from_lora_tensors(qwen3_lora_files, device):
52 tensors = load_file(os.path.join(qwen3_lora_files, "adapter_model.safetensors"))
53
54 peft_helper = PEFTHelper.from_local_dir(
55 qwen3_lora_files, max_position_embeddings=4096
56 )
57 lora_model = LoRAModel.from_lora_tensors(
58 1,
59 tensors,
60 peft_helper=peft_helper,
61 device=device,
62 )
63 for module_name, lora in lora_model.loras.items():
64 assert lora.module_name == module_name
65 assert lora.rank == 8
66 assert lora.lora_alpha == 32
67 assert lora.lora_a is not None
68 assert lora.lora_b is not None
69 assert lora.lora_a.device == torch.device(device)
70 assert lora.lora_b.device == torch.device(device)
71 assert lora.lora_a.shape[0] == lora.lora_b.shape[1], (
72 f"{lora.lora_a.shape=}, {lora.lora_b.shape=}"
73 )
74 assert lora.lora_a.shape[0] == 8
75
76
77 def create_lora(
78 lora_id: int, model: nn.Module, sub_modules: list[str], device: torch.device
79 ) -> LoRAModel:
80 loras: dict[str, LoRALayerWeights] = {}
81 for name in sub_modules:
82 w = model.get_submodule(name).weight
83 loras[name] = LoRALayerWeights(
84 name,
85 8,
86 16,
87 torch.rand([8, w.shape[1]], device=device),
88 torch.rand([w.shape[0], 8], device=device),
89 )
90 return LoRAModel(lora_id, 8, loras)
91
92
93 def create_packed_lora(
94 lora_id: int,
95 model: nn.Module,
96 module_name,
97 replaced_module_names,
98 device: torch.device,
99 empty_replaced_module_name=None,
100 ) -> LoRAModel:
101 w = model.get_submodule(module_name).weight
102 loras: dict[str, LoRALayerWeights] = {}
103 for replaced_module_name in replaced_module_names:
104 if replaced_module_name == empty_replaced_module_name:
105 continue
106 loras[replaced_module_name] = LoRALayerWeights(
107 replaced_module_name,
108 8,
109 16,
110 torch.rand([8, w.shape[1]], device=device),
111 torch.rand([w.shape[0] // len(replaced_module_names), 8], device=device),
112 )
113 return LoRAModel(lora_id, 8, loras)
114
115
116 def test_replace_submodules(default_vllm_config, dist_init, dummy_model):
117 model = dummy_model
118 manager = LoRAModelManager(
119 model,
120 1,
121 1,
122 1,
123 LoRAConfig(
124 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
125 ),
126 torch.device(DEVICES[0]),
127 default_vllm_config,
128 )
129 model = manager.model
130 assert isinstance(model.get_submodule("dense1"), ColumnParallelLinearWithLoRA)
131 assert isinstance(
132 model.get_submodule("layer1.dense1"), ColumnParallelLinearWithLoRA
133 )
134 assert isinstance(model.get_submodule("dense2"), RowParallelLinearWithLoRA)
135 assert isinstance(model.get_submodule("layer1.dense2"), RowParallelLinearWithLoRA)
136
137
138 def test_wrap_replicated_linear_subclasses(default_vllm_config, dist_init, dummy_model):
139 from vllm.model_executor.layers.linear import ReplicatedLinear
140
141 class CustomReplicatedLinear(ReplicatedLinear):
142 pass
143
144 model = dummy_model
145 model.add_module("custom_gate", CustomReplicatedLinear(10, 10, bias=False))
146
147 manager = LoRAModelManager(
148 model,
149 1,
150 1,
151 1,
152 LoRAConfig(
153 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
154 ),
155 torch.device(DEVICES[0]),
156 default_vllm_config,
157 )
158
159 assert isinstance(
160 manager.model.get_submodule("custom_gate"), ReplicatedLinearWithLoRA
161 )
162
163
164 def test_wrap_gate_linear(default_vllm_config, dist_init, dummy_model):
165 model = dummy_model
166 model.add_module("router_gate", GateLinear(10, 4, bias=False))
167
168 manager = LoRAModelManager(
169 model,
170 1,
171 1,
172 1,
173 LoRAConfig(
174 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
175 ),
176 torch.device(DEVICES[0]),
177 default_vllm_config,
178 )
179
180 assert isinstance(
181 manager.model.get_submodule("router_gate"), ReplicatedLinearWithLoRA
182 )
183
184
185 def test_dedup_shared_module_across_paths(default_vllm_config, dist_init, dummy_model):
186 """A module reachable from two attribute paths (e.g. a MoE gate held
187 both directly on the block and inside its inner runner) must produce a
188 single LoRA wrapper. Both paths must end up pointing to that same
189 wrapper instance, and only the canonical path should live in
190 `manager.modules` — otherwise activate_adapter would call `reset_lora`
191 on the alias and clobber weights set under the canonical name.
192 """
193 from vllm.model_executor.layers.linear import ReplicatedLinear
194
195 class AliasContainer(nn.Module):
196 def __init__(self, gate: nn.Module):
197 super().__init__()
198 self.gate = gate # canonical path: "moe.gate"
199
200 # Inner submodule holding the SAME gate instance under another
201 # path. This mirrors how MoERunner.runner.gate references the
202 # block's gate in qwen3_moe.
203 class _Runner(nn.Module):
204 def __init__(self, g):
205 super().__init__()
206 self.gate = g # alias path: "moe.runner.gate"
207
208 self.runner = _Runner(gate)
209
210 gate = ReplicatedLinear(10, 4, bias=False)
211 model = dummy_model
212 model.add_module("moe", AliasContainer(gate))
213
214 assert model.moe.gate is model.moe.runner.gate
215
216 manager = LoRAModelManager(
217 model,
218 1,
219 1,
220 1,
221 LoRAConfig(
222 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
223 ),
224 torch.device(DEVICES[0]),
225 default_vllm_config,
226 )
227
228 canonical = manager.model.get_submodule("moe.gate")
229 alias = manager.model.get_submodule("moe.runner.gate")
230
231 # Same wrapper instance on both paths so forward through either side
232 # sees the LoRA-augmented module.
233 assert isinstance(canonical, ReplicatedLinearWithLoRA)
234 assert alias is canonical
235
236 # Only the canonical path is tracked as a LoRA target. Tracking the
237 # alias would cause activate_adapter to reset_lora on it after the
238 # canonical entry already populated the weights.
239 assert "moe.gate" in manager.modules
240 assert "moe.runner.gate" not in manager.modules
241
242
243 def test_lm_head_exempt_from_dedup(default_vllm_config, dist_init, dummy_model):
244 """The dedup logic must NOT collapse `lm_head` even when it is reachable
245 from another attribute path (tied-embedding models do
246 `self.lm_head = self.model.embed_tokens`, sharing the same nn.Module
247 instance). The lm_head branch additionally rewires `logits_processor`
248 into a `LogitsProcessorWithLoRA`, so skipping it would silently break
249 LoRA on lm_head.
250 """
251 from vllm.lora.layers import LogitsProcessorWithLoRA
252
253 # Add a non-lm_head alias to the same module instance as lm_head. The
254 # dedup keys on id(module); without the lm_head exemption the alias
255 # would consume the wrapped_by_id slot first and lm_head would be
256 # silently skipped, so logits_processor would never be wrapped.
257 model = dummy_model
258 model.add_module("embed_tokens", model.lm_head)
259 assert model.embed_tokens is model.lm_head
260
261 manager = LoRAModelManager(
262 model,
263 1,
264 1,
265 1,
266 LoRAConfig(
267 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
268 ),
269 torch.device(DEVICES[0]),
270 default_vllm_config,
271 )
272
273 # lm_head's special handling still ran: logits_processor got wrapped
274 # and the lm_head entry is tracked under self.modules.
275 assert isinstance(
276 manager.model.get_submodule("logits_processor"), LogitsProcessorWithLoRA
277 )
278 assert "lm_head" in manager.modules
279
280
281 def test_skip_unsupported_matched_modules(default_vllm_config, dist_init, dummy_model):
282 class UnsupportedContainer(nn.Module):
283 def __init__(self):
284 super().__init__()
285 # This name matches a supported target suffix ("dense1"),
286 # but nn.Linear is not currently a LoRA-wrappable layer type.
287 self.dense1 = nn.Linear(10, 10, bias=False)
288
289 model = dummy_model
290 model.add_module("unsupported", UnsupportedContainer())
291
292 manager = LoRAModelManager(
293 model,
294 1,
295 1,
296 1,
297 LoRAConfig(
298 max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
299 ),
300 torch.device(DEVICES[0]),
301 default_vllm_config,
302 )
303
304 # Should not crash and should keep unsupported matched modules unchanged.
305 assert isinstance(manager.model.get_submodule("unsupported.dense1"), nn.Linear)
306 assert "unsupported.dense1" not in manager.modules
307
308
309 def test_target_modules_fail_closed_on_unsupported_matched_modules(
310 default_vllm_config, dist_init, dummy_model
311 ):
312 class UnsupportedContainer(nn.Module):
313 def __init__(self):
314 super().__init__()
315 self.dense1 = nn.Linear(10, 10, bias=False)
316
317 model = dummy_model
318 model.add_module("unsupported", UnsupportedContainer())
319
320 with pytest.raises(ValueError, match="unsupported.dense1"):
321 LoRAModelManager(
322 model,
323 1,
324 1,
325 1,
326 LoRAConfig(
327 max_lora_rank=8,
328 max_cpu_loras=8,
329 max_loras=8,
330 lora_dtype=DEFAULT_DTYPE,
331 target_modules=["dense1"],
332 ),
333 torch.device(DEVICES[0]),
334 default_vllm_config,
335 )
336
337
338 def test_get_dummy_lora_warmup_rank_for_fully_sharded_moe():
339 manager = LoRAModelManager.__new__(LoRAModelManager)
340 manager.lora_config = LoRAConfig(
341 max_lora_rank=64,
342 max_cpu_loras=1,
343 max_loras=1,
344 lora_dtype=DEFAULT_DTYPE,
345 fully_sharded_loras=True,
346 )
347
348 class DummyModule:
349 def __init__(self, tp_size: int, fully_sharded: bool):
350 self.tp_size = tp_size
351 self.fully_sharded = fully_sharded
352
353 manager.modules = {
354 "model.layers.0.self_attn.q_proj": DummyModule(
355 tp_size=32,
356 fully_sharded=True,
357 ),
358 "model.layers.0.mlp.experts": DummyModule(
359 tp_size=32,
360 fully_sharded=True,
361 ),
362 }
363
364 assert manager.get_dummy_lora_warmup_rank(8) == 32
365
366
367 @pytest.mark.parametrize("device", DEVICES)
368 def test_lora_model_manager(default_vllm_config, dist_init, dummy_model, device):
369 model = dummy_model
370 model_lora1 = create_lora(
371 1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
372 )
373 model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
374 model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
375 manager = LoRAModelManager(
376 model,
377 2,
378 2,
379 2,
380 LoRAConfig(
381 max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
382 ),
383 device=device,
384 vllm_config=default_vllm_config,
385 )
386 assert all(x is None for x in manager.lora_index_to_id)
387 assert manager.add_adapter(model_lora1)
388 assert manager.activate_adapter(1)
389 assert manager.lora_index_to_id[0] == 1
390 assert not manager.add_adapter(model_lora1)
391 assert not manager.activate_adapter(1)
392 assert manager.add_adapter(model_lora2)
393 assert manager.activate_adapter(2)
394 assert manager.lora_index_to_id[0] == 1
395 assert manager.lora_index_to_id[1] == 2
396 assert not manager.add_adapter(model_lora2)
397 assert not manager.activate_adapter(2)
398 assert manager.add_adapter(model_lora3)
399 assert manager.lora_index_to_id[0] == 1
400 assert manager.lora_index_to_id[1] == 2
401 with pytest.raises(ValueError):
402 assert manager.activate_adapter(3)
403 assert manager.lora_index_to_id[0] == 1
404 assert manager.lora_index_to_id[1] == 2
405 assert manager.remove_adapter(model_lora2.id)
406 assert manager.lora_index_to_id[1] is None
407 assert not manager.remove_adapter(model_lora2.id)
408 assert manager.remove_adapter(model_lora1.id)
409 assert not manager.remove_adapter(model_lora1.id)
410 assert manager.add_adapter(model_lora1)
411 assert manager.lora_index_to_id[0] is None
412 assert manager.lora_index_to_id[1] is None
413 assert manager.add_adapter(model_lora2)
414 assert manager.activate_adapter(3)
415 assert manager.lora_index_to_id[0] == 3
416 assert manager.lora_index_to_id[1] is None
417 assert manager.activate_adapter(2)
418 assert manager.lora_index_to_id[0] == 3
419 assert manager.lora_index_to_id[1] == 2
420 assert manager.device == device
421 assert (
422 manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
423 == device
424 )
425 assert hasattr(manager, "supported_lora_modules")
426 assert sorted(manager.supported_lora_modules) == [
427 "dense1",
428 "dense2",
429 "lm_head",
430 "output",
431 ]
432
433
434 @pytest.mark.parametrize("device", DEVICES)
435 def test_lora_lru_cache_model_manager(
436 default_vllm_config, dist_init, dummy_model, device
437 ):
438 model = dummy_model
439 model_lora1 = create_lora(
440 1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
441 )
442 model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
443 model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
444 manager = LRUCacheLoRAModelManager(
445 model,
446 2,
447 2,
448 2,
449 LoRAConfig(
450 max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
451 ),
452 device=device,
453 vllm_config=default_vllm_config,
454 )
455 assert all(x is None for x in manager.lora_index_to_id)
456 assert manager.add_adapter(model_lora1)
457 assert manager.activate_adapter(1)
458 assert manager.lora_index_to_id[0] == 1
459 assert not manager.add_adapter(model_lora1)
460 assert not manager.activate_adapter(1)
461 assert manager.add_adapter(model_lora2)
462 assert manager.activate_adapter(2)
463 assert manager.lora_index_to_id[0] == 1
464 assert manager.lora_index_to_id[1] == 2
465 assert not manager.add_adapter(model_lora2)
466 assert not manager.activate_adapter(2)
467 assert manager.add_adapter(model_lora3)
468 assert manager.lora_index_to_id[0] == 1
469 assert manager.lora_index_to_id[1] == 2
470 assert manager.activate_adapter(3)
471 assert manager.lora_index_to_id[0] == 3
472 assert manager.lora_index_to_id[1] == 2
473 assert manager.remove_adapter(model_lora2.id)
474 assert manager.lora_index_to_id[1] is None
475 assert not manager.remove_adapter(model_lora2.id)
476 assert manager.remove_adapter(model_lora1.id)
477 assert not manager.remove_adapter(model_lora1.id)
478 assert manager.add_adapter(model_lora1)
479 assert manager.activate_adapter(1)
480 assert manager.lora_index_to_id[0] == 3
481 assert manager.lora_index_to_id[1] == 1
482 assert manager.add_adapter(model_lora2)
483 assert manager.deactivate_adapter(3)
484 assert manager.lora_index_to_id[0] is None
485 assert manager.lora_index_to_id[1] == 1
486 assert manager.activate_adapter(2)
487 assert manager.lora_index_to_id[0] == 2
488 assert manager.lora_index_to_id[1] == 1
489 assert manager.activate_adapter(3)
490 assert manager.lora_index_to_id[0] == 2
491 assert manager.lora_index_to_id[1] == 3
492 assert manager.pin_adapter(2)
493 assert manager.lora_index_to_id[0] == 2
494 assert manager.lora_index_to_id[1] == 3
495 assert manager.activate_adapter(1)
496 assert manager.lora_index_to_id[0] == 2
497 assert manager.lora_index_to_id[1] == 1
498 assert manager.deactivate_adapter(2)
499 assert manager.lora_index_to_id[0] is None
500 assert manager.lora_index_to_id[1] == 1
501 assert manager.activate_adapter(3)
502 assert manager.lora_index_to_id[0] == 3
503 assert manager.lora_index_to_id[1] == 1
504 assert manager.pin_adapter(3)
505 assert manager.pin_adapter(1)
506 with pytest.raises(RuntimeError):
507 assert manager.pin_adapter(2)
508 assert manager.lora_index_to_id[0] == 3
509 assert manager.lora_index_to_id[1] == 1
510 with pytest.raises(RuntimeError):
511 assert manager.activate_adapter(2)
512
513 assert manager.deactivate_adapter(3)
514 assert manager.pin_adapter(2)
515 assert manager.lora_index_to_id[0] == 2
516 assert manager.lora_index_to_id[1] == 1
517 assert manager.remove_adapter(3)
518 with pytest.raises(ValueError):
519 assert manager.pin_adapter(3)
520 assert (
521 manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
522 == device
523 )
524 assert manager.device == device
525
526
527 @pytest.mark.parametrize("device", DEVICES)
528 def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, device):
529 # This tests just the LRU cache functionality, everything else is
530 # tested in test_lora_model_manager
531 model = dummy_model
532 model_lora1 = create_lora(
533 1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
534 )
535 model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
536 model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
537 model_lora4 = create_lora(4, model, ["dense1", "dense2", "lm_head"], device=device)
538 manager = LRUCacheLoRAModelManager(
539 model,
540 2,
541 2,
542 2,
543 LoRAConfig(
544 max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
545 ),
546 device=device,
547 vllm_config=default_vllm_config,
548 )
549 assert all(x is None for x in manager.lora_index_to_id)
550
551 # Add up to capacity
552 assert manager.add_adapter(model_lora1)
553 assert manager.add_adapter(model_lora2)
554 assert manager.activate_adapter(1)
555 assert manager.activate_adapter(2)
556
557 assert set(manager.list_adapters()) == {1, 2}
558 assert manager.lora_index_to_id[0] == 1
559 assert manager.lora_index_to_id[1] == 2
560
561 # Add over capacity
562 assert manager.add_adapter(model_lora3)
563 assert manager.add_adapter(model_lora4)
564 assert manager.activate_adapter(3)
565 assert manager.activate_adapter(4)
566
567 assert set(manager.list_adapters()) == {3, 4}
568 assert manager.lora_index_to_id[0] == 3
569 assert manager.lora_index_to_id[1] == 4
570
571 # Add 3 again to move it to the top and then add 2
572 # should return false since it's in already
573 assert not manager.add_adapter(model_lora3)
574 assert not manager.activate_adapter(3)
575 assert manager.add_adapter(model_lora2)
576 assert manager.activate_adapter(2)
577
578 assert set(manager.list_adapters()) == {3, 2}
579 assert manager.lora_index_to_id[0] == 3
580 assert manager.lora_index_to_id[1] == 2
581
582 # Remove manually
583 assert manager.remove_adapter(3)
584 assert not manager.remove_adapter(3)
585
586 assert set(manager.list_adapters()) == {2}
587 assert manager.lora_index_to_id[0] is None
588 assert manager.lora_index_to_id[1] == 2
589
590 assert manager.add_adapter(model_lora3)
591 assert manager.activate_adapter(3)
592 assert manager.add_adapter(model_lora4)
593 assert manager.activate_adapter(4)
594
595 assert set(manager.list_adapters()) == {3, 4}
596 assert manager.lora_index_to_id[0] == 3
597 assert manager.lora_index_to_id[1] == 4
598
599 assert manager.remove_oldest_adapter()
600 assert set(manager.list_adapters()) == {4}
601 assert manager.lora_index_to_id[0] is None
602 assert manager.lora_index_to_id[1] == 4
603
604 assert manager.remove_oldest_adapter()
605 assert set(manager.list_adapters()) == set()
606 assert all(x is None for x in manager.lora_index_to_id)
607
608 assert not manager.remove_oldest_adapter()
609 assert set(manager.list_adapters()) == set()
610 assert all(x is None for x in manager.lora_index_to_id)
611
612 # pinning
613 assert manager.add_adapter(model_lora3)
614 assert manager.activate_adapter(3)
615 assert manager.add_adapter(model_lora4)
616 assert manager.activate_adapter(4)
617 assert set(manager.list_adapters()) == {3, 4}
618 with pytest.raises(ValueError):
619 assert manager.pin_adapter(1)
620 assert manager.pin_adapter(3)
621 # Remove manually
622 assert manager.remove_adapter(3)
623 assert not manager.remove_adapter(3)
624
625 assert set(manager.list_adapters()) == {4}
626 assert manager.lora_index_to_id[0] is None
627 assert manager.lora_index_to_id[1] == 4
628
629 assert manager.add_adapter(model_lora1)
630 assert manager.pin_adapter(1)
631 assert manager.add_adapter(model_lora2)
632 assert manager.activate_adapter(2)
633
634 assert set(manager.list_adapters()) == {1, 2}
635 assert manager.lora_index_to_id[0] == 1
636 assert manager.lora_index_to_id[1] == 2
637
638 assert manager.remove_oldest_adapter()
639 assert set(manager.list_adapters()) == {1}
640 assert manager.lora_index_to_id[0] == 1
641 assert manager.lora_index_to_id[1] is None
642
643 with pytest.raises(RuntimeError):
644 assert manager.remove_oldest_adapter()
645
646 assert set(manager.list_adapters()) == {1}
647 assert (
648 manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
649 == device
650 )
651 assert manager.device == device
652
653
654 @pytest.mark.parametrize("device", DEVICES)
655 def test_set_adapter_mapping_refreshes_after_slot_reassignment(
656 default_vllm_config, dist_init, dummy_model, device
657 ):
658 # An out-of-band add_lora() can LRU-evict and reassign GPU slots while the
659 # running batch -- and therefore its LoRAMapping -- is unchanged. The
660 # punica metadata must still be re-derived, otherwise in-flight requests
661 # are routed to the evicted layout and decode with the wrong adapter.
662 model = dummy_model
663 model_lora1 = create_lora(1, model, ["dense1", "dense2", "lm_head"], device=device)
664 model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
665 model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
666 manager = LRUCacheLoRAModelManager(
667 model,
668 2,
669 2,
670 2,
671 LoRAConfig(
672 max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
673 ),
674 device=device,
675 vllm_config=default_vllm_config,
676 )
677 punica_wrapper = manager.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY]
678
679 assert manager.add_adapter(model_lora1)
680 assert manager.activate_adapter(1)
681 assert manager.add_adapter(model_lora2)
682 assert manager.activate_adapter(2)
683 assert manager.lora_index_to_id == [1, 2]
684
685 # Two in-flight requests, one token each, on adapters 1 and 2.
686 manager.set_adapter_mapping(LoRAMapping((1, 2), (1, 2)))
687 assert punica_wrapper.token_lora_indices.tolist() == [0, 1]
688
689 # Out-of-band add_lora() with both slots held by the running batch:
690 # activating 3 evicts 1; re-activating the batch's adapters lands them
691 # in swapped slots while the batch itself is unchanged.
692 assert manager.add_adapter(model_lora3)
693 assert manager.activate_adapter(3)
694 assert manager.activate_adapter(1)
695 assert manager.activate_adapter(2)
696 assert manager.lora_index_to_id == [2, 1]
697
698 # Identical mapping, but the metadata must follow the new slot layout.
699 manager.set_adapter_mapping(LoRAMapping((1, 2), (1, 2)))
700 assert punica_wrapper.token_lora_indices.tolist() == [1, 0]
701
702 manager.remove_all_adapters()
703 assert manager._last_mapping is None
704 assert manager._last_slot_layout is None
705
706
707 @pytest.mark.parametrize("device", DEVICES)
708 def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path):
709 lora_config = LoRAConfig(
710 max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
711 )
712
713 dummy_lora_files = f"{tmp_path}/lora_adapter"
714 os.makedirs(dummy_lora_files, exist_ok=True)
715 create_peft_lora(
716 dummy_model,
717 save_dir=dummy_lora_files,
718 target_modules=["layer1.dense1", "dense2"],
719 lora_dtype=DEFAULT_DTYPE,
720 )
721
722 model_config = ModelConfig(max_model_len=16)
723 vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
724
725 vllm_config.scheduler_config.max_num_seqs = 4
726 vllm_config.scheduler_config.max_num_batched_tokens = 2
727 worker_adapter_manager = LRUCacheWorkerLoRAManager(
728 vllm_config, device, EMBEDDING_MODULES
729 )
730
731 worker_adapter_manager.max_num_seqs = 4
732 worker_adapter_manager.max_num_batched_tokens = 2
733
734 worker_adapter_manager.create_lora_manager(dummy_model, vllm_config)
735
736 mapping = LoRAMapping([], [])
737 worker_adapter_manager.set_active_adapters(
738 [LoRARequest("1", 1, dummy_lora_files), LoRARequest("2", 2, dummy_lora_files)],
739 mapping,
740 )
741 assert worker_adapter_manager.list_adapters() == {1, 2}
742 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
743 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
744
745 worker_adapter_manager.set_active_adapters(
746 [
747 LoRARequest("1", 1, dummy_lora_files),
748 LoRARequest("3", 3, dummy_lora_files),
749 LoRARequest("4", 4, dummy_lora_files),
750 ],
751 mapping,
752 )
753 assert worker_adapter_manager.list_adapters() == {1, 2, 3, 4}
754 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
755 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
756 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 3
757 assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
758
759 worker_adapter_manager.set_active_adapters(
760 [
761 LoRARequest("1", 1, dummy_lora_files),
762 LoRARequest("2", 2, dummy_lora_files),
763 LoRARequest("5", 5, dummy_lora_files),
764 ],
765 mapping,
766 )
767 assert worker_adapter_manager.list_adapters() == {1, 2, 4, 5}
768 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
769 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
770 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
771 assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
772
773 worker_adapter_manager.set_active_adapters(
774 [
775 LoRARequest("1", 1, dummy_lora_files),
776 LoRARequest("1", 1, dummy_lora_files),
777 LoRARequest("1", 1, dummy_lora_files),
778 ],
779 mapping,
780 )
781 assert worker_adapter_manager.list_adapters() == {1, 2, 4, 5}
782 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
783 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
784 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
785 assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
786
787 worker_adapter_manager.set_active_adapters(
788 [
789 LoRARequest("6", 6, dummy_lora_files),
790 LoRARequest("7", 7, dummy_lora_files),
791 LoRARequest("8", 8, dummy_lora_files),
792 ],
793 mapping,
794 )
795 assert worker_adapter_manager.list_adapters() == {1, 6, 7, 8}
796 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
797 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 7
798 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 8
799 assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 6
800
801 # Over capacity
802 with pytest.raises(RuntimeError):
803 worker_adapter_manager.set_active_adapters(
804 [
805 LoRARequest("10", 10, dummy_lora_files),
806 LoRARequest("11", 11, dummy_lora_files),
807 LoRARequest("12", 12, dummy_lora_files),
808 LoRARequest("13", 13, dummy_lora_files),
809 LoRARequest("14", 14, dummy_lora_files),
810 ],
811 mapping,
812 )
813
814 assert worker_adapter_manager.device == device
815 punica_wrapper = worker_adapter_manager._adapter_manager.punica_wrapper_mapping.get(
816 DEFAULT_LANGUAGE_WRAPPER_KEY
817 )
818 assert punica_wrapper.device == device
819
820
821 @pytest.mark.parametrize("device", DEVICES)
822 def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path):
823 # Should remove every LoRA not specified in the request.
824 lora_config = LoRAConfig(
825 max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
826 )
827
828 model_config = ModelConfig(max_model_len=16)
829 vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
830
831 vllm_config.scheduler_config.max_num_seqs = 4
832 vllm_config.scheduler_config.max_num_batched_tokens = 2
833
834 worker_adapter_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
835 worker_adapter_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
836 worker_adapter_manager.create_lora_manager(dummy_model_gate_up, vllm_config)
837
838 dummy_lora_files = f"{tmp_path}/lora_adapter"
839 os.makedirs(dummy_lora_files, exist_ok=True)
840 create_peft_lora(
841 dummy_model_gate_up,
842 save_dir=dummy_lora_files,
843 target_modules=["layer1.dense1", "dense2"],
844 lora_dtype=DEFAULT_DTYPE,
845 )
846
847 mapping = LoRAMapping([], [])
848 worker_adapter_manager.set_active_adapters(
849 [LoRARequest("1", 1, dummy_lora_files), LoRARequest("2", 2, dummy_lora_files)],
850 mapping,
851 )
852 assert worker_adapter_manager.list_adapters() == {1, 2}
853 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
854 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
855
856 worker_adapter_manager.set_active_adapters(
857 [
858 LoRARequest("1", 1, dummy_lora_files),
859 LoRARequest("3", 3, dummy_lora_files),
860 LoRARequest("4", 4, dummy_lora_files),
861 ],
862 mapping,
863 )
864 assert worker_adapter_manager.list_adapters() == {1, 3, 4}
865 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
866 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 3
867 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 4
868
869 worker_adapter_manager.set_active_adapters(
870 [
871 LoRARequest("1", 1, dummy_lora_files),
872 LoRARequest("2", 2, dummy_lora_files),
873 LoRARequest("5", 5, dummy_lora_files),
874 ],
875 mapping,
876 )
877 assert worker_adapter_manager.list_adapters() == {1, 2, 5}
878 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
879 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
880 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
881
882 worker_adapter_manager.set_active_adapters(
883 [
884 LoRARequest("1", 1, dummy_lora_files),
885 LoRARequest("1", 1, dummy_lora_files),
886 LoRARequest("1", 1, dummy_lora_files),
887 ],
888 mapping,
889 )
890 assert worker_adapter_manager.list_adapters() == {1}
891 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
892 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] is None
893 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] is None
894
895 worker_adapter_manager.set_active_adapters(
896 [
897 LoRARequest("6", 6, dummy_lora_files),
898 LoRARequest("7", 7, dummy_lora_files),
899 LoRARequest("8", 8, dummy_lora_files),
900 ],
901 mapping,
902 )
903 assert worker_adapter_manager.list_adapters() == {6, 7, 8}
904 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 8
905 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 6
906 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 7
907
908 # Over capacity
909 with pytest.raises(RuntimeError):
910 worker_adapter_manager.set_active_adapters(
911 [
912 LoRARequest("10", 10, dummy_lora_files),
913 LoRARequest("11", 11, dummy_lora_files),
914 LoRARequest("12", 12, dummy_lora_files),
915 LoRARequest("13", 13, dummy_lora_files),
916 LoRARequest("14", 14, dummy_lora_files),
917 ],
918 mapping,
919 )
920
921 assert worker_adapter_manager.device == device
922 punica_wrapper = worker_adapter_manager._adapter_manager.punica_wrapper_mapping.get(
923 DEFAULT_LANGUAGE_WRAPPER_KEY
924 )
925 assert punica_wrapper.device == device
926
927
928 @pytest.mark.parametrize("device", DEVICES)
929 def test_packed_loras(default_vllm_config, dist_init, dummy_model_gate_up, device):
930 model = dummy_model_gate_up
931 model_lora = create_packed_lora(
932 1,
933 model,
934 module_name="gate_up_proj",
935 replaced_module_names=["gate_proj", "up_proj"],
936 device=device,
937 )
938 model_lora1 = create_packed_lora(
939 2,
940 model,
941 module_name="gate_up_proj",
942 replaced_module_names=["gate_proj", "up_proj"],
943 device=device,
944 empty_replaced_module_name="gate_proj",
945 )
946
947 manager = LoRAModelManager(
948 model,
949 2,
950 2,
951 2,
952 LoRAConfig(
953 max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
954 ),
955 device=device,
956 vllm_config=default_vllm_config,
957 )
958 model = manager.model
959
960 assert isinstance(
961 model.get_submodule("gate_up_proj"), MergedColumnParallelLinearWithLoRA
962 )
963 # Verify packed lora is correct
964 model_lora_clone = model_lora.clone(1)
965 model_lora_clone1 = model_lora1.clone(1)
966 assert manager.add_adapter(model_lora)
967 assert manager.add_adapter(model_lora1)
968
969 assert model_lora.get_lora("gate_proj") is None
970 assert model_lora.get_lora("up_proj") is None
971 assert model_lora1.get_lora("up_proj") is None
972 packed_lora = model_lora.get_lora("gate_up_proj")
973 assert packed_lora and isinstance(packed_lora, PackedLoRALayerWeights)
974
975 torch.testing.assert_close(
976 packed_lora.lora_a[0], model_lora_clone.get_lora("gate_proj").lora_a
977 )
978 torch.testing.assert_close(
979 packed_lora.lora_b[0], model_lora_clone.get_lora("gate_proj").lora_b
980 )
981 torch.testing.assert_close(
982 packed_lora.lora_a[1], model_lora_clone.get_lora("up_proj").lora_a
983 )
984 torch.testing.assert_close(
985 packed_lora.lora_b[1], model_lora_clone.get_lora("up_proj").lora_b
986 )
987
988 packed_lora1 = model_lora1.get_lora("gate_up_proj")
989 assert packed_lora1 and isinstance(packed_lora1, PackedLoRALayerWeights)
990
991 assert packed_lora1.lora_a[0] is None
992 assert packed_lora1.lora_b[0] is None
993 torch.testing.assert_close(
994 packed_lora1.lora_a[1], model_lora_clone1.get_lora("up_proj").lora_a
995 )
996 torch.testing.assert_close(
997 packed_lora1.lora_b[1], model_lora_clone1.get_lora("up_proj").lora_b
998 )
999
1000
1001 def _test_target_modules(
1002 model,
1003 target_modules: list[str] | None,
1004 device: str,
1005 expected_lora: list[tuple[str, type]],
1006 expected_no_lora: list[tuple[str, type]],
1007 vllm_config,
1008 ):
1009 """Create a LoRAModelManager and assert which modules have LoRA applied."""
1010 LoRAModelManager(
1011 model,
1012 2,
1013 2,
1014 2,
1015 LoRAConfig(
1016 max_lora_rank=8,
1017 max_cpu_loras=2,
1018 max_loras=2,
1019 lora_dtype=DEFAULT_DTYPE,
1020 target_modules=target_modules,
1021 ),
1022 device=device,
1023 vllm_config=vllm_config,
1024 )
1025 for module_path, lora_cls in expected_lora:
1026 assert isinstance(model.get_submodule(module_path), lora_cls)
1027 for module_path, lora_cls in expected_no_lora:
1028 assert not isinstance(model.get_submodule(module_path), lora_cls)
1029
1030
1031 @pytest.mark.parametrize("device", DEVICES)
1032 def test_target_modules_config(default_vllm_config, dist_init, dummy_model, device):
1033 """Test that target_modules config restricts which modules get LoRA applied."""
1034 _test_target_modules(
1035 dummy_model,
1036 ["dense1"],
1037 device,
1038 expected_lora=[
1039 ("dense1", ColumnParallelLinearWithLoRA),
1040 ("layer1.dense1", ColumnParallelLinearWithLoRA),
1041 ],
1042 expected_no_lora=[
1043 ("dense2", RowParallelLinearWithLoRA),
1044 ("layer1.dense2", RowParallelLinearWithLoRA),
1045 ],
1046 vllm_config=default_vllm_config,
1047 )
1048
1049
1050 @pytest.mark.parametrize("device", DEVICES)
1051 def test_target_modules_multiple(default_vllm_config, dist_init, dummy_model, device):
1052 """Test that multiple target_modules work correctly."""
1053 _test_target_modules(
1054 dummy_model,
1055 ["dense1", "dense2"],
1056 device,
1057 expected_lora=[
1058 ("dense1", ColumnParallelLinearWithLoRA),
1059 ("layer1.dense1", ColumnParallelLinearWithLoRA),
1060 ("dense2", RowParallelLinearWithLoRA),
1061 ("layer1.dense2", RowParallelLinearWithLoRA),
1062 ],
1063 expected_no_lora=[],
1064 vllm_config=default_vllm_config,
1065 )
1066
1067
1068 @pytest.mark.parametrize("device", DEVICES)
1069 def test_target_modules_none_uses_all(
1070 default_vllm_config, dist_init, dummy_model, device
1071 ):
1072 """Test that target_modules=None uses all supported modules."""
1073 _test_target_modules(
1074 dummy_model,
1075 None,
1076 device,
1077 expected_lora=[
1078 ("dense1", ColumnParallelLinearWithLoRA),
1079 ("layer1.dense1", ColumnParallelLinearWithLoRA),
1080 ("dense2", RowParallelLinearWithLoRA),
1081 ("layer1.dense2", RowParallelLinearWithLoRA),
1082 ],
1083 expected_no_lora=[],
1084 vllm_config=default_vllm_config,
1085 )
1086
1087
1088 @pytest.mark.parametrize("device", DEVICES)
1089 def test_target_modules_match_packed_runtime_modules(
1090 default_vllm_config, dist_init, dummy_model_gate_up, device
1091 ):
1092 """Packed runtime modules should be selected by their adapter-visible names."""
1093 _test_target_modules(
1094 dummy_model_gate_up,
1095 ["gate_proj"],
1096 device,
1097 expected_lora=[("gate_up_proj", MergedColumnParallelLinearWithLoRA)],
1098 expected_no_lora=[
1099 ("dense1", ColumnParallelLinearWithLoRA),
1100 ("dense2", RowParallelLinearWithLoRA),
1101 ("layer1.dense1", ColumnParallelLinearWithLoRA),
1102 ("layer1.dense2", RowParallelLinearWithLoRA),
1103 ],
1104 vllm_config=default_vllm_config,
1105 )