@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 285178b3b8
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 14, 05:30 UTC

[V0 Deprecation] Remove V0 LoRA test (#23418)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Aug 22, 2025 at 17:56 UTC 285178b3b824d70b46b351daa8f8942d23da264a
7 files changed +158 -116
tests/lora/conftest.py
+4 -27
@@ -3,15 +3,13 @@
3
4 import tempfile
5 from collections import OrderedDict
6 -from unittest.mock import MagicMock, patch
6 +from unittest.mock import MagicMock
7
8 import pytest
9 import torch
10 import torch.nn as nn
11 from huggingface_hub import snapshot_download
12
13 -import vllm
14 -from vllm.config import LoRAConfig
13 from vllm.distributed import (cleanup_dist_env_and_memory,
14 init_distributed_environment,
15 initialize_model_parallel)
@@ -21,7 +19,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
19 from vllm.model_executor.layers.logits_processor import LogitsProcessor
20 from vllm.model_executor.layers.sampler import Sampler
21 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
24 -from vllm.model_executor.model_loader import get_model
22 from vllm.model_executor.models.interfaces import SupportsLoRA
23 from vllm.platforms import current_platform
24
@@ -104,6 +101,7 @@ def dummy_model() -> nn.Module:
101 ]))
102 model.config = MagicMock()
103 model.embedding_modules = {"lm_head": "lm_head"}
104 + model.unpadded_vocab_size = 32000
105 return model
106
107
@@ -137,6 +135,8 @@ def dummy_model_gate_up() -> nn.Module:
135 ],
136 }
137 model.embedding_modules = {"lm_head": "lm_head"}
138 + model.unpadded_vocab_size = 32000
139 +
140 return model
141
142
@@ -221,29 +221,6 @@ def phi2_lora_files():
221 return snapshot_download(repo_id="isotr0py/phi-2-test-sql-lora")
222
223
224 -@pytest.fixture
225 -def llama_2_7b_engine_extra_embeddings():
226 - cleanup_dist_env_and_memory(shutdown_ray=True)
227 - get_model_old = get_model
228 -
229 - def get_model_patched(**kwargs):
230 - kwargs["vllm_config"].lora_config = LoRAConfig(max_loras=4,
231 - max_lora_rank=8)
232 - return get_model_old(**kwargs)
233 -
234 - with patch("vllm.worker.model_runner.get_model", get_model_patched):
235 - engine = vllm.LLM("meta-llama/Llama-2-7b-hf", enable_lora=False)
236 - yield engine.llm_engine
237 - del engine
238 - cleanup_dist_env_and_memory(shutdown_ray=True)
239 -
240 -
241 -@pytest.fixture
242 -def llama_2_7b_model_extra_embeddings(llama_2_7b_engine_extra_embeddings):
243 - yield (llama_2_7b_engine_extra_embeddings.model_executor.driver_worker.
244 - model_runner.model)
245 -
246 -
224 @pytest.fixture
225 def reset_default_device():
226 """
tests/lora/test_add_lora.py
+4 -7
@@ -5,7 +5,6 @@ import time
5
6 import pytest
7
8 -import vllm.envs as env
8 from vllm.engine.arg_utils import AsyncEngineArgs
9 from vllm.entrypoints.openai.api_server import (
10 build_async_engine_client_from_engine_args)
@@ -98,12 +97,10 @@ async def test_add_lora(chatglm3_lora_files):
97 # Run with warmup
98 add_lora_tasks = [llm.add_lora(lr) for lr in warmup_run_requests]
99 add_lora_results = await asyncio.gather(*add_lora_tasks)
101 - if env.VLLM_USE_V1:
102 - # Test that all all_lora calls are successful.
103 - assert all(add_lora_results)
104 - else:
105 - # No way to check V0 engine results as the calls just return None.
106 - pass
100 +
101 + # Test that all all_lora calls are successful.
102 + assert all(add_lora_results)
103 +
104 time_with_add_lora = await requests_processing_time(
105 llm, warmup_run_requests)
106
tests/lora/test_llama_tp.py
+1 -4
@@ -113,8 +113,7 @@ def test_llama_lora(sql_lora_files):
113 enable_lora=True,
114 # also test odd max_num_seqs
115 max_num_seqs=13,
116 - max_loras=4,
117 - enable_chunked_prefill=True)
116 + max_loras=4)
117 generate_and_test(llm, sql_lora_files)
118
119
@@ -128,7 +127,6 @@ def test_llama_lora_tp4(sql_lora_files):
127 max_num_seqs=16,
128 max_loras=4,
129 tensor_parallel_size=4,
131 - enable_chunked_prefill=True,
130 )
131 generate_and_test(llm, sql_lora_files)
132
@@ -144,7 +142,6 @@ def test_llama_lora_tp4_fully_sharded_loras(sql_lora_files):
142 max_loras=4,
143 tensor_parallel_size=4,
144 fully_sharded_loras=True,
147 - enable_chunked_prefill=True,
145 )
146 generate_and_test(llm, sql_lora_files)
147
tests/lora/test_lora_manager.py
+68 -62
@@ -21,6 +21,8 @@ from vllm.lora.worker_manager import (LRUCacheWorkerLoRAManager,
21 WorkerLoRAManager)
22 from vllm.platforms import current_platform
23
24 +from .utils import create_peft_lora
25 +
26 EMBEDDING_MODULES = {
27 "embed_tokens": "input_embeddings",
28 "lm_head": "output_embeddings",
@@ -35,17 +37,6 @@ DEVICES = ([
37 DEFAULT_DTYPE = torch.get_default_dtype()
38
39
38 -@pytest.fixture(scope="function", autouse=True)
39 -def use_v0_only(monkeypatch: pytest.MonkeyPatch):
40 - """
41 - Some tests depend on V0 internals. Since both V0 and V1 use the same
42 - LoRAModelManager it is okay to just test V0.
43 - """
44 - with monkeypatch.context() as m:
45 - m.setenv('VLLM_USE_V1', '0')
46 - yield
47 -
48 -
40 @pytest.mark.parametrize("device", DEVICES)
41 def test_from_lora_tensors(sql_lora_files, device):
42 tensors = load_file(
@@ -326,7 +317,6 @@ def test_lru_lora_model_manager(dist_init, dummy_model, device):
317 max_loras=2,
318 lora_dtype=DEFAULT_DTYPE),
319 device=device)
329 -
320 assert all(x is None for x in manager.lora_index_to_id)
321
322 # Add up to capacity
@@ -430,32 +420,40 @@ def test_lru_lora_model_manager(dist_init, dummy_model, device):
420
421
422 @pytest.mark.parametrize("device", DEVICES)
433 -def test_lru_cache_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
434 - sql_lora_files, device):
423 +def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device,
424 + tmp_path):
425 lora_config = LoRAConfig(max_lora_rank=8,
426 max_cpu_loras=4,
427 max_loras=4,
428 lora_dtype=DEFAULT_DTYPE)
429 +
430 + dummy_lora_files = f"{tmp_path}/lora_adapter"
431 + os.makedirs(dummy_lora_files, exist_ok=True)
432 + create_peft_lora(
433 + dummy_model,
434 + save_dir=dummy_lora_files,
435 + target_modules=["layer1.dense1", "dense2"],
436 + lora_dtype=DEFAULT_DTYPE,
437 + )
438 worker_adapter_manager = LRUCacheWorkerLoRAManager(
440 - 4, 2, llama_2_7b_model_extra_embeddings.unpadded_vocab_size -
441 - lora_config.lora_extra_vocab_size, lora_config, device,
442 - EMBEDDING_MODULES, EMBEDDING_PADDING_MODULES)
443 - worker_adapter_manager.create_lora_manager(
444 - llama_2_7b_model_extra_embeddings)
439 + 4, 2,
440 + dummy_model.unpadded_vocab_size - lora_config.lora_extra_vocab_size,
441 + lora_config, device, EMBEDDING_MODULES, EMBEDDING_PADDING_MODULES)
442 + worker_adapter_manager.create_lora_manager(dummy_model)
443
444 mapping = LoRAMapping([], [])
445 worker_adapter_manager.set_active_adapters([
448 - LoRARequest("1", 1, sql_lora_files),
449 - LoRARequest("2", 2, sql_lora_files)
446 + LoRARequest("1", 1, dummy_lora_files),
447 + LoRARequest("2", 2, dummy_lora_files)
448 ], mapping)
449 assert worker_adapter_manager.list_adapters() == {1, 2}
450 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
451 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
452
453 worker_adapter_manager.set_active_adapters([
456 - LoRARequest("1", 1, sql_lora_files),
457 - LoRARequest("3", 3, sql_lora_files),
458 - LoRARequest("4", 4, sql_lora_files)
454 + LoRARequest("1", 1, dummy_lora_files),
455 + LoRARequest("3", 3, dummy_lora_files),
456 + LoRARequest("4", 4, dummy_lora_files)
457 ], mapping)
458 assert worker_adapter_manager.list_adapters() == {1, 2, 3, 4}
459 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
@@ -464,9 +462,9 @@ def test_lru_cache_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
462 assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
463
464 worker_adapter_manager.set_active_adapters([
467 - LoRARequest("1", 1, sql_lora_files),
468 - LoRARequest("2", 2, sql_lora_files),
469 - LoRARequest("5", 5, sql_lora_files)
465 + LoRARequest("1", 1, dummy_lora_files),
466 + LoRARequest("2", 2, dummy_lora_files),
467 + LoRARequest("5", 5, dummy_lora_files)
468 ], mapping)
469 assert worker_adapter_manager.list_adapters() == {1, 2, 4, 5}
470 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
@@ -475,9 +473,9 @@ def test_lru_cache_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
473 assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
474
475 worker_adapter_manager.set_active_adapters([
478 - LoRARequest("1", 1, sql_lora_files),
479 - LoRARequest("1", 1, sql_lora_files),
480 - LoRARequest("1", 1, sql_lora_files)
476 + LoRARequest("1", 1, dummy_lora_files),
477 + LoRARequest("1", 1, dummy_lora_files),
478 + LoRARequest("1", 1, dummy_lora_files)
479 ], mapping)
480 assert worker_adapter_manager.list_adapters() == {1, 2, 4, 5}
481 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
@@ -486,9 +484,9 @@ def test_lru_cache_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
484 assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4
485
486 worker_adapter_manager.set_active_adapters([
489 - LoRARequest("6", 6, sql_lora_files),
490 - LoRARequest("7", 7, sql_lora_files),
491 - LoRARequest("8", 8, sql_lora_files)
487 + LoRARequest("6", 6, dummy_lora_files),
488 + LoRARequest("7", 7, dummy_lora_files),
489 + LoRARequest("8", 8, dummy_lora_files)
490 ], mapping)
491 assert worker_adapter_manager.list_adapters() == {1, 6, 7, 8}
492 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
@@ -499,11 +497,11 @@ def test_lru_cache_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
497 # Over capacity
498 with pytest.raises(RuntimeError):
499 worker_adapter_manager.set_active_adapters([
502 - LoRARequest("10", 10, sql_lora_files),
503 - LoRARequest("11", 11, sql_lora_files),
504 - LoRARequest("12", 12, sql_lora_files),
505 - LoRARequest("13", 13, sql_lora_files),
506 - LoRARequest("14", 14, sql_lora_files)
500 + LoRARequest("10", 10, dummy_lora_files),
501 + LoRARequest("11", 11, dummy_lora_files),
502 + LoRARequest("12", 12, dummy_lora_files),
503 + LoRARequest("13", 13, dummy_lora_files),
504 + LoRARequest("14", 14, dummy_lora_files)
505 ], mapping)
506
507 assert worker_adapter_manager.device == device
@@ -512,33 +510,41 @@ def test_lru_cache_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
510
511
512 @pytest.mark.parametrize("device", DEVICES)
515 -def test_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
516 - sql_lora_files, device):
513 +def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device,
514 + tmp_path):
515 # Should remove every LoRA not specified in the request.
516 lora_config = LoRAConfig(max_lora_rank=8,
517 max_cpu_loras=4,
518 max_loras=4,
519 lora_dtype=DEFAULT_DTYPE)
520 worker_adapter_manager = WorkerLoRAManager(
523 - 4, 2, llama_2_7b_model_extra_embeddings.unpadded_vocab_size -
521 + 4, 2, dummy_model_gate_up.unpadded_vocab_size -
522 lora_config.lora_extra_vocab_size, lora_config, device,
523 EMBEDDING_MODULES, EMBEDDING_PADDING_MODULES)
526 - worker_adapter_manager.create_lora_manager(
527 - llama_2_7b_model_extra_embeddings)
524 + worker_adapter_manager.create_lora_manager(dummy_model_gate_up)
525 +
526 + dummy_lora_files = f"{tmp_path}/lora_adapter"
527 + os.makedirs(dummy_lora_files, exist_ok=True)
528 + create_peft_lora(
529 + dummy_model_gate_up,
530 + save_dir=dummy_lora_files,
531 + target_modules=["layer1.dense1", "dense2"],
532 + lora_dtype=DEFAULT_DTYPE,
533 + )
534
535 mapping = LoRAMapping([], [])
536 worker_adapter_manager.set_active_adapters([
531 - LoRARequest("1", 1, sql_lora_files),
532 - LoRARequest("2", 2, sql_lora_files)
537 + LoRARequest("1", 1, dummy_lora_files),
538 + LoRARequest("2", 2, dummy_lora_files)
539 ], mapping)
540 assert worker_adapter_manager.list_adapters() == {1, 2}
541 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
542 assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
543
544 worker_adapter_manager.set_active_adapters([
539 - LoRARequest("1", 1, sql_lora_files),
540 - LoRARequest("3", 3, sql_lora_files),
541 - LoRARequest("4", 4, sql_lora_files)
545 + LoRARequest("1", 1, dummy_lora_files),
546 + LoRARequest("3", 3, dummy_lora_files),
547 + LoRARequest("4", 4, dummy_lora_files)
548 ], mapping)
549 assert worker_adapter_manager.list_adapters() == {1, 3, 4}
550 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
@@ -546,9 +552,9 @@ def test_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
552 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 4
553
554 worker_adapter_manager.set_active_adapters([
549 - LoRARequest("1", 1, sql_lora_files),
550 - LoRARequest("2", 2, sql_lora_files),
551 - LoRARequest("5", 5, sql_lora_files)
555 + LoRARequest("1", 1, dummy_lora_files),
556 + LoRARequest("2", 2, dummy_lora_files),
557 + LoRARequest("5", 5, dummy_lora_files)
558 ], mapping)
559 assert worker_adapter_manager.list_adapters() == {1, 2, 5}
560 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
@@ -556,9 +562,9 @@ def test_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
562 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
563
564 worker_adapter_manager.set_active_adapters([
559 - LoRARequest("1", 1, sql_lora_files),
560 - LoRARequest("1", 1, sql_lora_files),
561 - LoRARequest("1", 1, sql_lora_files)
565 + LoRARequest("1", 1, dummy_lora_files),
566 + LoRARequest("1", 1, dummy_lora_files),
567 + LoRARequest("1", 1, dummy_lora_files)
568 ], mapping)
569 assert worker_adapter_manager.list_adapters() == {1}
570 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
@@ -566,9 +572,9 @@ def test_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
572 assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] is None
573
574 worker_adapter_manager.set_active_adapters([
569 - LoRARequest("6", 6, sql_lora_files),
570 - LoRARequest("7", 7, sql_lora_files),
571 - LoRARequest("8", 8, sql_lora_files)
575 + LoRARequest("6", 6, dummy_lora_files),
576 + LoRARequest("7", 7, dummy_lora_files),
577 + LoRARequest("8", 8, dummy_lora_files)
578 ], mapping)
579 assert worker_adapter_manager.list_adapters() == {6, 7, 8}
580 assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 8
@@ -578,11 +584,11 @@ def test_worker_adapter_manager(llama_2_7b_model_extra_embeddings,
584 # Over capacity
585 with pytest.raises(RuntimeError):
586 worker_adapter_manager.set_active_adapters([
581 - LoRARequest("10", 10, sql_lora_files),
582 - LoRARequest("11", 11, sql_lora_files),
583 - LoRARequest("12", 12, sql_lora_files),
584 - LoRARequest("13", 13, sql_lora_files),
585 - LoRARequest("14", 14, sql_lora_files)
587 + LoRARequest("10", 10, dummy_lora_files),
588 + LoRARequest("11", 11, dummy_lora_files),
589 + LoRARequest("12", 12, dummy_lora_files),
590 + LoRARequest("13", 13, dummy_lora_files),
591 + LoRARequest("14", 14, dummy_lora_files)
592 ], mapping)
593
594 assert worker_adapter_manager.device == device
tests/lora/test_mixtral.py
-1
@@ -50,7 +50,6 @@ def test_mixtral_lora(mixtral_lora_files, tp_size):
50 max_loras=4,
51 distributed_executor_backend="ray",
52 tensor_parallel_size=tp_size,
53 - enable_chunked_prefill=True,
53 )
54
55 expected_lora_output = [
tests/lora/test_worker.py
+5 -15
@@ -4,17 +4,14 @@
4 import os
5 import random
6 import tempfile
7 -from typing import Union
7 from unittest.mock import patch
8
10 -import vllm.envs as envs
9 from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig,
10 ModelConfig, ParallelConfig, SchedulerConfig,
11 VllmConfig)
12 from vllm.lora.models import LoRAMapping
13 from vllm.lora.request import LoRARequest
16 -from vllm.v1.worker.gpu_worker import Worker as V1Worker
17 -from vllm.worker.worker import Worker
14 +from vllm.v1.worker.gpu_worker import Worker
15
16 NUM_LORAS = 16
17
@@ -22,18 +19,11 @@ NUM_LORAS = 16
19 @patch.dict(os.environ, {"RANK": "0"})
20 def test_worker_apply_lora(sql_lora_files):
21
25 - def set_active_loras(worker: Union[Worker, V1Worker],
26 - lora_requests: list[LoRARequest]):
22 + def set_active_loras(worker: Worker, lora_requests: list[LoRARequest]):
23 lora_mapping = LoRAMapping([], [])
28 - if isinstance(worker, Worker):
29 - # v0 case
30 - worker.model_runner.set_active_loras(lora_requests, lora_mapping)
31 - else:
32 - # v1 case
33 - worker.model_runner.lora_manager.set_active_adapters(
34 - lora_requests, lora_mapping)
24
36 - worker_cls = V1Worker if envs.VLLM_USE_V1 else Worker
25 + worker.model_runner.lora_manager.set_active_adapters(
26 + lora_requests, lora_mapping)
27
28 vllm_config = VllmConfig(
29 model_config=ModelConfig(
@@ -62,7 +52,7 @@ def test_worker_apply_lora(sql_lora_files):
52 max_cpu_loras=NUM_LORAS,
53 max_loras=NUM_LORAS),
54 )
65 - worker = worker_cls(
55 + worker = Worker(
56 vllm_config=vllm_config,
57 local_rank=0,
58 rank=0,
tests/lora/utils.py
+76
@@ -1,10 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import json
5 +import os
6 from dataclasses import dataclass
7 from typing import Optional, Union
8
9 import torch
10 +from safetensors.torch import save_file
11
12 from vllm.lora.lora import LoRALayerWeights, PackedLoRALayerWeights
13
@@ -340,3 +343,76 @@ def generate_data_for_nslices(
343 seq_len_tensor,
344 indices,
345 )
346 +
347 +
348 +def create_peft_lora(
349 + model: torch.nn.Module,
350 + save_dir: str,
351 + target_modules: list[str],
352 + rank: int = 8,
353 + alpha: int = 16,
354 + dropout: float = 0.1,
355 + lora_dtype: torch.dtype = torch.float16,
356 +) -> dict[str, torch.Tensor]:
357 + lora_weights = {}
358 + adapter_config = {
359 + "peft_type": "LORA",
360 + "auto_mapping": None,
361 + "base_model_name_or_path": "dummy_model",
362 + "revision": None,
363 + "task_type": "CAUSAL_LM",
364 + "inference_mode": False,
365 + "r": rank,
366 + "lora_alpha": alpha,
367 + "lora_dropout": dropout,
368 + "fan_in_fan_out": False,
369 + "bias": "none",
370 + "modules_to_save": None,
371 + "init_lora_weights": True,
372 + "layers_to_transform": None,
373 + "layers_pattern": None,
374 + "target_modules": target_modules,
375 + "exclude_modules": None,
376 + "use_rslora": False,
377 + "use_dora": False,
378 + "loftq_config": None,
379 + }
380 +
381 + for module_name in target_modules:
382 +
383 + module = model
384 + for attr in module_name.split("."):
385 + module = getattr(module, attr)
386 +
387 + if hasattr(module, "input_size") and hasattr(module, "output_size"):
388 +
389 + in_features = module.input_size
390 + out_features = module.output_size
391 +
392 + elif hasattr(module, "embedding_dim") and hasattr(
393 + module, "num_embeddings"):
394 + # ParallelLMHead
395 + in_features = module.embedding_dim
396 + out_features = module.num_embeddings
397 + else:
398 + raise ValueError(
399 + f"Unable to determine dimensions for module {module_name}")
400 +
401 + lora_A = torch.randn(rank, in_features, dtype=lora_dtype)
402 +
403 + torch.nn.init.kaiming_uniform_(lora_A, a=5**0.5)
404 +
405 + lora_B = torch.zeros(out_features, rank, dtype=lora_dtype)
406 +
407 + # PEFT style
408 + lora_weights[f"base_model.model.{module_name}.lora_A.weight"] = lora_A
409 + lora_weights[f"base_model.model.{module_name}.lora_B.weight"] = lora_B
410 +
411 + config_path = os.path.join(save_dir, "adapter_config.json")
412 + with open(config_path, "w", encoding="utf-8") as f:
413 + json.dump(adapter_config, f, indent=2, ensure_ascii=False)
414 +
415 + weights_path = os.path.join(save_dir, "adapter_model.safetensors")
416 + save_file(lora_weights, weights_path)
417 +
418 + return lora_weights