test_lora_manager.py 24.8 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
6
7
8
9
10
import os

import pytest
import torch
from safetensors.torch import load_file
from torch import nn

11
from vllm.config import ModelConfig, VllmConfig
12
from vllm.config.lora import LoRAConfig
13
14
15
16
17
from vllm.lora.layers import (
    ColumnParallelLinearWithLoRA,
    MergedColumnParallelLinearWithLoRA,
    RowParallelLinearWithLoRA,
)
18
from vllm.lora.lora_model import LoRAModel
19
from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
20
from vllm.lora.model_manager import (
21
22
23
24
    LoRAMapping,
    LoRAModelManager,
    LRUCacheLoRAModelManager,
)
25
from vllm.lora.peft_helper import PEFTHelper
26
from vllm.lora.request import LoRARequest
27
from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager, WorkerLoRAManager
28
from vllm.platforms import current_platform
29

30
31
from .utils import create_peft_lora

Terry's avatar
Terry committed
32
33
34
35
36
37
EMBEDDING_MODULES = {
    "embed_tokens": "input_embeddings",
    "lm_head": "output_embeddings",
}


38
39
40
41
42
DEVICES = (
    [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
    if current_platform.is_cuda_alike()
    else ["cpu"]
)
43

44
45
DEFAULT_DTYPE = torch.get_default_dtype()

46

47
@pytest.mark.parametrize("device", DEVICES)
48
49
def test_from_lora_tensors(qwen3_lora_files, device):
    tensors = load_file(os.path.join(qwen3_lora_files, "adapter_model.safetensors"))
50

51
    peft_helper = PEFTHelper.from_local_dir(
52
        qwen3_lora_files, max_position_embeddings=4096
53
    )
Terry's avatar
Terry committed
54
55
56
    lora_model = LoRAModel.from_lora_tensors(
        1,
        tensors,
57
58
        peft_helper=peft_helper,
        device=device,
59
    )
60
61
62
    for module_name, lora in lora_model.loras.items():
        assert lora.module_name == module_name
        assert lora.rank == 8
63
        assert lora.lora_alpha == 32
64
65
        assert lora.lora_a is not None
        assert lora.lora_b is not None
66
67
        assert lora.lora_a.device == torch.device(device)
        assert lora.lora_b.device == torch.device(device)
68
69
70
        assert lora.lora_a.shape[0] == lora.lora_b.shape[1], (
            f"{lora.lora_a.shape=}, {lora.lora_b.shape=}"
        )
71
        assert lora.lora_a.shape[0] == 8
72
73


74
75
76
def create_lora(
    lora_id: int, model: nn.Module, sub_modules: list[str], device: torch.device
) -> LoRAModel:
77
    loras: dict[str, LoRALayerWeights] = {}
78
79
80
81
82
83
    for name in sub_modules:
        w = model.get_submodule(name).weight
        loras[name] = LoRALayerWeights(
            name,
            8,
            16,
84
85
            torch.rand([8, w.shape[1]], device=device),
            torch.rand([w.shape[0], 8], device=device),
86
87
88
89
90
91
92
93
94
        )
    return LoRAModel(lora_id, 8, loras)


def create_packed_lora(
    lora_id: int,
    model: nn.Module,
    module_name,
    replaced_module_names,
95
    device: torch.device,
96
97
98
    empty_replaced_module_name=None,
) -> LoRAModel:
    w = model.get_submodule(module_name).weight
99
    loras: dict[str, LoRALayerWeights] = {}
100
101
102
103
104
105
106
    for replaced_module_name in replaced_module_names:
        if replaced_module_name == empty_replaced_module_name:
            continue
        loras[replaced_module_name] = LoRALayerWeights(
            replaced_module_name,
            8,
            16,
107
            torch.rand([8, w.shape[1]], device=device),
108
            torch.rand([w.shape[0] // len(replaced_module_names), 8], device=device),
109
110
111
112
113
114
        )
    return LoRAModel(lora_id, 8, loras)


def test_replace_submodules(dist_init, dummy_model):
    model = dummy_model
Terry's avatar
Terry committed
115
    manager = LoRAModelManager(
116
117
118
119
120
121
122
123
124
        model,
        1,
        1,
        1,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
        ),
        torch.device(DEVICES[0]),
    )
125
    model = manager.model
126
127
128
129
    assert isinstance(model.get_submodule("dense1"), ColumnParallelLinearWithLoRA)
    assert isinstance(
        model.get_submodule("layer1.dense1"), ColumnParallelLinearWithLoRA
    )
130
    assert isinstance(model.get_submodule("dense2"), RowParallelLinearWithLoRA)
131
    assert isinstance(model.get_submodule("layer1.dense2"), RowParallelLinearWithLoRA)
132
133


134
@pytest.mark.parametrize("device", DEVICES)
135
def test_lora_model_manager(dist_init, dummy_model, device):
136
    model = dummy_model
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
    model_lora1 = create_lora(
        1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
    )
    model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
    model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
    manager = LoRAModelManager(
        model,
        2,
        2,
        2,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
        ),
        device=device,
    )
152
    assert all(x is None for x in manager.lora_index_to_id)
153
154
    assert manager.add_adapter(model_lora1)
    assert manager.activate_adapter(1)
155
    assert manager.lora_index_to_id[0] == 1
156
157
158
159
    assert not manager.add_adapter(model_lora1)
    assert not manager.activate_adapter(1)
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(2)
160
161
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2
162
163
164
    assert not manager.add_adapter(model_lora2)
    assert not manager.activate_adapter(2)
    assert manager.add_adapter(model_lora3)
165
166
167
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2
    with pytest.raises(ValueError):
168
        assert manager.activate_adapter(3)
169
170
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2
171
    assert manager.remove_adapter(model_lora2.id)
172
    assert manager.lora_index_to_id[1] is None
173
174
175
176
    assert not manager.remove_adapter(model_lora2.id)
    assert manager.remove_adapter(model_lora1.id)
    assert not manager.remove_adapter(model_lora1.id)
    assert manager.add_adapter(model_lora1)
177
178
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] is None
179
180
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(3)
181
182
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] is None
183
    assert manager.activate_adapter(2)
184
185
186
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 2

187
188
    assert manager.device == device
    assert manager.punica_wrapper.device == device
189
190
191
192
193
194
195
    assert hasattr(manager, "supported_lora_modules")
    assert sorted(manager.supported_lora_modules) == [
        "dense1",
        "dense2",
        "lm_head",
        "output",
    ]
196

197

198
@pytest.mark.parametrize("device", DEVICES)
199
def test_lora_lru_cache_model_manager(dist_init, dummy_model, device):
200
    model = dummy_model
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
    model_lora1 = create_lora(
        1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
    )
    model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
    model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
    manager = LRUCacheLoRAModelManager(
        model,
        2,
        2,
        2,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
        ),
        device=device,
    )
216
    assert all(x is None for x in manager.lora_index_to_id)
217
218
    assert manager.add_adapter(model_lora1)
    assert manager.activate_adapter(1)
219
    assert manager.lora_index_to_id[0] == 1
220
221
222
223
    assert not manager.add_adapter(model_lora1)
    assert not manager.activate_adapter(1)
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(2)
224
225
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2
226
227
228
    assert not manager.add_adapter(model_lora2)
    assert not manager.activate_adapter(2)
    assert manager.add_adapter(model_lora3)
229
230
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2
231
    assert manager.activate_adapter(3)
232
233
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 2
234
    assert manager.remove_adapter(model_lora2.id)
235
    assert manager.lora_index_to_id[1] is None
236
237
238
239
240
    assert not manager.remove_adapter(model_lora2.id)
    assert manager.remove_adapter(model_lora1.id)
    assert not manager.remove_adapter(model_lora1.id)
    assert manager.add_adapter(model_lora1)
    assert manager.activate_adapter(1)
241
242
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 1
243
244
    assert manager.add_adapter(model_lora2)
    assert manager.deactivate_adapter(3)
245
246
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] == 1
247
    assert manager.activate_adapter(2)
248
249
    assert manager.lora_index_to_id[0] == 2
    assert manager.lora_index_to_id[1] == 1
250
    assert manager.activate_adapter(3)
251
252
    assert manager.lora_index_to_id[0] == 2
    assert manager.lora_index_to_id[1] == 3
253
    assert manager.pin_adapter(2)
254
255
    assert manager.lora_index_to_id[0] == 2
    assert manager.lora_index_to_id[1] == 3
256
    assert manager.activate_adapter(1)
257
258
    assert manager.lora_index_to_id[0] == 2
    assert manager.lora_index_to_id[1] == 1
259
    assert manager.deactivate_adapter(2)
260
261
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] == 1
262
    assert manager.activate_adapter(3)
263
264
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 1
265
266
    assert manager.pin_adapter(3)
    assert manager.pin_adapter(1)
267
    with pytest.raises(RuntimeError):
268
        assert manager.pin_adapter(2)
269
270
271
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 1
    with pytest.raises(RuntimeError):
272
        assert manager.activate_adapter(2)
273

274
275
    assert manager.deactivate_adapter(3)
    assert manager.pin_adapter(2)
276
277
    assert manager.lora_index_to_id[0] == 2
    assert manager.lora_index_to_id[1] == 1
278
    assert manager.remove_adapter(3)
279
    with pytest.raises(ValueError):
280
        assert manager.pin_adapter(3)
281

282
283
284
    assert manager.punica_wrapper.device == device
    assert manager.device == device

285

286
@pytest.mark.parametrize("device", DEVICES)
287
def test_lru_lora_model_manager(dist_init, dummy_model, device):
288
289
290
    # This tests just the LRU cache functionality, everything else is
    # tested in test_lora_model_manager
    model = dummy_model
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
    model_lora1 = create_lora(
        1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
    )
    model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
    model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
    model_lora4 = create_lora(4, model, ["dense1", "dense2", "lm_head"], device=device)
    manager = LRUCacheLoRAModelManager(
        model,
        2,
        2,
        2,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
        ),
        device=device,
    )
307
308
309
    assert all(x is None for x in manager.lora_index_to_id)

    # Add up to capacity
310
311
312
313
    assert manager.add_adapter(model_lora1)
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(1)
    assert manager.activate_adapter(2)
314

315
    assert set(manager.list_adapters()) == {1, 2}
316
317
318
319
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2

    # Add over capacity
320
321
322
323
    assert manager.add_adapter(model_lora3)
    assert manager.add_adapter(model_lora4)
    assert manager.activate_adapter(3)
    assert manager.activate_adapter(4)
324

325
    assert set(manager.list_adapters()) == {3, 4}
326
327
328
329
330
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 4

    # Add 3 again to move it to the top and then add 2
    # should return false since it's in already
331
332
333
334
    assert not manager.add_adapter(model_lora3)
    assert not manager.activate_adapter(3)
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(2)
335

336
    assert set(manager.list_adapters()) == {3, 2}
337
338
339
340
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 2

    # Remove manually
341
342
    assert manager.remove_adapter(3)
    assert not manager.remove_adapter(3)
343

344
    assert set(manager.list_adapters()) == {2}
345
346
347
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] == 2

348
349
350
351
    assert manager.add_adapter(model_lora3)
    assert manager.activate_adapter(3)
    assert manager.add_adapter(model_lora4)
    assert manager.activate_adapter(4)
352

353
    assert set(manager.list_adapters()) == {3, 4}
354
355
356
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 4

357
358
    assert manager.remove_oldest_adapter()
    assert set(manager.list_adapters()) == {4}
359
360
361
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] == 4

362
363
    assert manager.remove_oldest_adapter()
    assert set(manager.list_adapters()) == set()
364
365
    assert all(x is None for x in manager.lora_index_to_id)

366
367
    assert not manager.remove_oldest_adapter()
    assert set(manager.list_adapters()) == set()
368
369
    assert all(x is None for x in manager.lora_index_to_id)

370
    # pinning
371
372
373
374
375
    assert manager.add_adapter(model_lora3)
    assert manager.activate_adapter(3)
    assert manager.add_adapter(model_lora4)
    assert manager.activate_adapter(4)
    assert set(manager.list_adapters()) == {3, 4}
376
    with pytest.raises(ValueError):
377
378
        assert manager.pin_adapter(1)
    assert manager.pin_adapter(3)
379
    # Remove manually
380
381
    assert manager.remove_adapter(3)
    assert not manager.remove_adapter(3)
382

383
    assert set(manager.list_adapters()) == {4}
384
385
386
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] == 4

387
388
389
390
    assert manager.add_adapter(model_lora1)
    assert manager.pin_adapter(1)
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(2)
391

392
    assert set(manager.list_adapters()) == {1, 2}
393
394
395
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2

396
397
    assert manager.remove_oldest_adapter()
    assert set(manager.list_adapters()) == {1}
398
399
400
401
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] is None

    with pytest.raises(RuntimeError):
402
        assert manager.remove_oldest_adapter()
403

404
    assert set(manager.list_adapters()) == {1}
405
406
    assert manager.punica_wrapper.device == device
    assert manager.device == device
407

408

409
@pytest.mark.parametrize("device", DEVICES)
410
411
412
413
def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path):
    lora_config = LoRAConfig(
        max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
    )
414
415
416
417
418
419
420
421
422

    dummy_lora_files = f"{tmp_path}/lora_adapter"
    os.makedirs(dummy_lora_files, exist_ok=True)
    create_peft_lora(
        dummy_model,
        save_dir=dummy_lora_files,
        target_modules=["layer1.dense1", "dense2"],
        lora_dtype=DEFAULT_DTYPE,
    )
423
424

    model_config = ModelConfig(max_model_len=16)
425
    vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
426
427
428

    vllm_config.scheduler_config.max_num_seqs = 4
    vllm_config.scheduler_config.max_num_batched_tokens = 2
429
    worker_adapter_manager = LRUCacheWorkerLoRAManager(
430
        vllm_config, device, EMBEDDING_MODULES
431
    )
432
433
434
435

    worker_adapter_manager.max_num_seqs = 4
    worker_adapter_manager.max_num_batched_tokens = 2

436
    worker_adapter_manager.create_lora_manager(dummy_model)
437
438

    mapping = LoRAMapping([], [])
439
440
441
442
    worker_adapter_manager.set_active_adapters(
        [LoRARequest("1", 1, dummy_lora_files), LoRARequest("2", 2, dummy_lora_files)],
        mapping,
    )
443
444
445
    assert worker_adapter_manager.list_adapters() == {1, 2}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
446

447
448
449
450
451
452
453
454
    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("3", 3, dummy_lora_files),
            LoRARequest("4", 4, dummy_lora_files),
        ],
        mapping,
    )
455
456
457
458
459
    assert worker_adapter_manager.list_adapters() == {1, 2, 3, 4}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 3
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
460

461
462
463
464
465
466
467
468
    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("2", 2, dummy_lora_files),
            LoRARequest("5", 5, dummy_lora_files),
        ],
        mapping,
    )
469
470
471
472
473
    assert worker_adapter_manager.list_adapters() == {1, 2, 4, 5}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
474

475
476
477
478
479
480
481
482
    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("1", 1, dummy_lora_files),
        ],
        mapping,
    )
483
484
485
486
487
    assert worker_adapter_manager.list_adapters() == {1, 2, 4, 5}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
488

489
490
491
492
493
494
495
496
    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("6", 6, dummy_lora_files),
            LoRARequest("7", 7, dummy_lora_files),
            LoRARequest("8", 8, dummy_lora_files),
        ],
        mapping,
    )
497
498
499
500
501
    assert worker_adapter_manager.list_adapters() == {1, 6, 7, 8}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 7
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 8
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 6
502
503
504

    # Over capacity
    with pytest.raises(RuntimeError):
505
506
507
508
509
510
511
512
513
514
        worker_adapter_manager.set_active_adapters(
            [
                LoRARequest("10", 10, dummy_lora_files),
                LoRARequest("11", 11, dummy_lora_files),
                LoRARequest("12", 12, dummy_lora_files),
                LoRARequest("13", 13, dummy_lora_files),
                LoRARequest("14", 14, dummy_lora_files),
            ],
            mapping,
        )
515

516
    assert worker_adapter_manager.device == device
517
    assert worker_adapter_manager._adapter_manager.punica_wrapper.device == device
518

519

520
@pytest.mark.parametrize("device", DEVICES)
521
def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path):
522
    # Should remove every LoRA not specified in the request.
523
524
525
    lora_config = LoRAConfig(
        max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
    )
526
527

    model_config = ModelConfig(max_model_len=16)
528
    vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
529
530
531
532

    vllm_config.scheduler_config.max_num_seqs = 4
    vllm_config.scheduler_config.max_num_batched_tokens = 2

533
    worker_adapter_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
534
    worker_adapter_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
535
536
537
538
539
540
541
542
543
544
    worker_adapter_manager.create_lora_manager(dummy_model_gate_up)

    dummy_lora_files = f"{tmp_path}/lora_adapter"
    os.makedirs(dummy_lora_files, exist_ok=True)
    create_peft_lora(
        dummy_model_gate_up,
        save_dir=dummy_lora_files,
        target_modules=["layer1.dense1", "dense2"],
        lora_dtype=DEFAULT_DTYPE,
    )
545
546

    mapping = LoRAMapping([], [])
547
548
549
550
    worker_adapter_manager.set_active_adapters(
        [LoRARequest("1", 1, dummy_lora_files), LoRARequest("2", 2, dummy_lora_files)],
        mapping,
    )
551
552
553
    assert worker_adapter_manager.list_adapters() == {1, 2}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
554

555
556
557
558
559
560
561
562
    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("3", 3, dummy_lora_files),
            LoRARequest("4", 4, dummy_lora_files),
        ],
        mapping,
    )
563
564
565
566
    assert worker_adapter_manager.list_adapters() == {1, 3, 4}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 3
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 4
567

568
569
570
571
572
573
574
575
    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("2", 2, dummy_lora_files),
            LoRARequest("5", 5, dummy_lora_files),
        ],
        mapping,
    )
576
577
578
579
    assert worker_adapter_manager.list_adapters() == {1, 2, 5}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
580

581
582
583
584
585
586
587
588
    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("1", 1, dummy_lora_files),
        ],
        mapping,
    )
589
590
591
592
    assert worker_adapter_manager.list_adapters() == {1}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] is None
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] is None
593

594
595
596
597
598
599
600
601
    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("6", 6, dummy_lora_files),
            LoRARequest("7", 7, dummy_lora_files),
            LoRARequest("8", 8, dummy_lora_files),
        ],
        mapping,
    )
602
603
604
605
    assert worker_adapter_manager.list_adapters() == {6, 7, 8}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 8
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 6
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 7
606
607
608

    # Over capacity
    with pytest.raises(RuntimeError):
609
610
611
612
613
614
615
616
617
618
        worker_adapter_manager.set_active_adapters(
            [
                LoRARequest("10", 10, dummy_lora_files),
                LoRARequest("11", 11, dummy_lora_files),
                LoRARequest("12", 12, dummy_lora_files),
                LoRARequest("13", 13, dummy_lora_files),
                LoRARequest("14", 14, dummy_lora_files),
            ],
            mapping,
        )
619

620
    assert worker_adapter_manager.device == device
621
    assert worker_adapter_manager._adapter_manager.punica_wrapper.device == device
622

623

624
@pytest.mark.parametrize("device", DEVICES)
625
def test_packed_loras(dist_init, dummy_model_gate_up, device):
626
627
628
629
630
    model = dummy_model_gate_up
    model_lora = create_packed_lora(
        1,
        model,
        module_name="gate_up_proj",
631
        replaced_module_names=["gate_proj", "up_proj"],
632
633
        device=device,
    )
634
635
636
637
638
    model_lora1 = create_packed_lora(
        2,
        model,
        module_name="gate_up_proj",
        replaced_module_names=["gate_proj", "up_proj"],
639
        device=device,
640
641
642
        empty_replaced_module_name="gate_proj",
    )

643
644
645
646
647
648
649
650
651
652
    manager = LoRAModelManager(
        model,
        2,
        2,
        2,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
        ),
        device=device,
    )
653
654
    model = manager.model

655
656
657
    assert isinstance(
        model.get_submodule("gate_up_proj"), MergedColumnParallelLinearWithLoRA
    )
658
659
660
    # Verify packed lora is correct
    model_lora_clone = model_lora.clone(1)
    model_lora_clone1 = model_lora1.clone(1)
661
662
    assert manager.add_adapter(model_lora)
    assert manager.add_adapter(model_lora1)
663

664
665
666
    assert model_lora.get_lora("gate_proj") is None
    assert model_lora.get_lora("up_proj") is None
    assert model_lora1.get_lora("up_proj") is None
667
668
669
    packed_lora = model_lora.get_lora("gate_up_proj")
    assert packed_lora and isinstance(packed_lora, PackedLoRALayerWeights)

670
671
672
673
674
675
676
677
678
679
680
681
    torch.testing.assert_close(
        packed_lora.lora_a[0], model_lora_clone.get_lora("gate_proj").lora_a
    )
    torch.testing.assert_close(
        packed_lora.lora_b[0], model_lora_clone.get_lora("gate_proj").lora_b
    )
    torch.testing.assert_close(
        packed_lora.lora_a[1], model_lora_clone.get_lora("up_proj").lora_a
    )
    torch.testing.assert_close(
        packed_lora.lora_b[1], model_lora_clone.get_lora("up_proj").lora_b
    )
682
683
684
685
686
687

    packed_lora1 = model_lora1.get_lora("gate_up_proj")
    assert packed_lora1 and isinstance(packed_lora1, PackedLoRALayerWeights)

    assert packed_lora1.lora_a[0] is None
    assert packed_lora1.lora_b[0] is None
688
689
690
691
692
693
    torch.testing.assert_close(
        packed_lora1.lora_a[1], model_lora_clone1.get_lora("up_proj").lora_a
    )
    torch.testing.assert_close(
        packed_lora1.lora_b[1], model_lora_clone1.get_lora("up_proj").lora_b
    )