@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2821986450

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

[Core] Modify the initialization parameters of the lora manager (#25249)

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

Jee Jee Li committed Sep 20, 2025 at 02:01 UTC 2821986450bc31869714885ed4203650a42f3cb0

10 files changed +51 -52

tests/lora/test_lora_manager.py

+28 -8

@@ -8,11 +8,12 @@ import torch

8 from safetensors.torch import load_file

9 from torch import nn

11 +from vllm.config import ModelConfig, VllmConfig

12 from vllm.config.lora import LoRAConfig

13 from vllm.lora.layers import (ColumnParallelLinearWithLoRA,

14 MergedColumnParallelLinearWithLoRA,

15 RowParallelLinearWithLoRA)

15 -from vllm.lora.lora import LoRALayerWeights, PackedLoRALayerWeights

16 +from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights

17 from vllm.lora.models import (LoRAMapping, LoRAModel, LoRAModelManager,

18 LRUCacheLoRAModelManager)

19 from vllm.lora.peft_helper import PEFTHelper

@@ -435,10 +436,19 @@ def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device,

436 target_modules=["layer1.dense1", "dense2"],

437 lora_dtype=DEFAULT_DTYPE,

438 )

439 +

440 + model_config = ModelConfig(max_model_len=16)

441 + vllm_config = VllmConfig(model_config=model_config,

442 + lora_config=lora_config)

443 +

444 + vllm_config.scheduler_config.max_num_seqs = 4

445 + vllm_config.scheduler_config.max_num_batched_tokens = 2

446 worker_adapter_manager = LRUCacheWorkerLoRAManager(

439 - 4, 2,

440 - dummy_model.unpadded_vocab_size - lora_config.lora_extra_vocab_size,

441 - lora_config, device, EMBEDDING_MODULES, EMBEDDING_PADDING_MODULES)

447 + vllm_config, device, EMBEDDING_MODULES, EMBEDDING_PADDING_MODULES)

448 +

449 + worker_adapter_manager.max_num_seqs = 4

450 + worker_adapter_manager.max_num_batched_tokens = 2

451 +

452 worker_adapter_manager.create_lora_manager(dummy_model)

453

454 mapping = LoRAMapping([], [])

@@ -517,10 +527,20 @@ def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device,

527 max_cpu_loras=4,

528 max_loras=4,

529 lora_dtype=DEFAULT_DTYPE)

520 - worker_adapter_manager = WorkerLoRAManager(

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)

530 +

531 + model_config = ModelConfig(max_model_len=16)

532 + vllm_config = VllmConfig(model_config=model_config,

533 + lora_config=lora_config)

534 +

535 + vllm_config.scheduler_config.max_num_seqs = 4

536 + vllm_config.scheduler_config.max_num_batched_tokens = 2

537 +

538 + worker_adapter_manager = WorkerLoRAManager(vllm_config, device,

539 + EMBEDDING_MODULES,

540 + EMBEDDING_PADDING_MODULES)

541 + worker_adapter_manager.vocab_size = (

542 + dummy_model_gate_up.unpadded_vocab_size -

543 + lora_config.lora_extra_vocab_size)

544 worker_adapter_manager.create_lora_manager(dummy_model_gate_up)

545

546 dummy_lora_files = f"{tmp_path}/lora_adapter"

tests/lora/utils.py

+1 -1

@@ -9,7 +9,7 @@ from typing import Optional, Union

9 import torch

10 from safetensors.torch import save_file

12 -from vllm.lora.lora import LoRALayerWeights, PackedLoRALayerWeights

12 +from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights

15 class DummyLoRAManager:

vllm/lora/lora_weights.py renamed

vllm/lora/models.py

+1 -1

@@ -14,7 +14,7 @@ from torch import nn

14 from vllm.config.lora import LoRAConfig

15 from vllm.logger import init_logger

16 from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping

17 -from vllm.lora.lora import LoRALayerWeights, PackedLoRALayerWeights

17 +from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights

18 from vllm.lora.peft_helper import PEFTHelper

19 from vllm.lora.punica_wrapper import get_punica_wrapper

20 from vllm.lora.utils import (from_layer, from_layer_logits_processor,

vllm/lora/worker_manager.py

+12 -11

@@ -6,7 +6,7 @@ from typing import Any, Literal, Optional, Union

7 import torch

9 -from vllm.config.lora import LoRAConfig

9 +from vllm.config import VllmConfig

10 from vllm.logger import init_logger

11 from vllm.lora.models import (LoRAModel, LoRAModelManager,

12 LRUCacheLoRAModelManager, create_lora_manager)

@@ -27,25 +27,26 @@ class WorkerLoRAManager:

28 def __init__(

29 self,

30 - max_num_seqs: int,

31 - max_num_batched_tokens: int,

32 - vocab_size: int,

33 - lora_config: LoRAConfig,

30 + vllm_config: VllmConfig,

31 device: torch.device,

32 embedding_modules: dict[str, str],

33 embedding_padding_modules: list[str],

34 lora_model_cls: type[LoRAModel] = LoRAModel,

38 - max_position_embeddings: Optional[int] = None,

35 ):

36 self._lora_model_cls = lora_model_cls

37 self.embedding_modules = embedding_modules

38 self.embedding_padding_modules = embedding_padding_modules

39 self._cached_dummy_lora: Union[None, Literal[False], LoRAModel] = False

44 - self.max_num_seqs = max_num_seqs

45 - self.max_num_batched_tokens = max_num_batched_tokens

46 - self.vocab_size = vocab_size

47 - self.lora_config = lora_config

48 - self.max_position_embeddings = max_position_embeddings

40 + self.max_num_seqs = vllm_config.scheduler_config.max_num_seqs

41 + self.max_num_batched_tokens = (

42 + vllm_config.scheduler_config.max_num_batched_tokens)

43 + self.vocab_size = vllm_config.model_config.get_vocab_size()

44 + self.lora_config = vllm_config.lora_config

45 +

46 + # Use get_text_config() in case of multimodal models

47 + text_config = vllm_config.model_config.hf_config.get_text_config()

48 +

49 + self.max_position_embeddings = text_config.max_position_embeddings

50 self.device = device

51 # Lazily initialized by create_lora_manager.

52 self._adapter_manager: LoRAModelManager

vllm/v1/worker/cpu_model_runner.py

+2 -3

@@ -107,9 +107,8 @@ class CPUModelRunner(GPUModelRunner):

107 self.model = get_model(vllm_config=self.vllm_config)

108

109 if self.lora_config:

110 - self.model = self.load_lora_model(self.model, self.model_config,

111 - self.scheduler_config,

112 - self.lora_config, self.device)

110 + self.model = self.load_lora_model(self.model, self.vllm_config,

111 + self.device)

112

113 def get_model(self) -> nn.Module:

114 return self.model

vllm/v1/worker/gpu_model_runner.py

+1 -4

@@ -2552,10 +2552,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2552 self.model = model_loader.load_model(

2553 vllm_config=self.vllm_config, model_config=self.model_config)

2554 if self.lora_config:

2555 - self.model = self.load_lora_model(self.model,

2556 - self.model_config,

2557 - self.scheduler_config,

2558 - self.lora_config,

2555 + self.model = self.load_lora_model(self.model, self.vllm_config,

2556 self.device)

2557 if hasattr(self, "drafter"):

2558 logger.info("Loading drafter model...")

vllm/v1/worker/lora_model_runner_mixin.py

+3 -12

@@ -11,7 +11,7 @@ import numpy as np

11 import torch

12 import torch.nn as nn

14 -from vllm.config import ModelConfig, SchedulerConfig

14 +from vllm.config import VllmConfig

15 from vllm.config.lora import LoRAConfig

16 from vllm.logger import init_logger

17 from vllm.lora.layers import LoRAMapping

@@ -31,9 +31,7 @@ class LoRAModelRunnerMixin:

32 LORA_WARMUP_RANK = 8

34 - def load_lora_model(self, model: nn.Module, model_config: ModelConfig,

35 - scheduler_config: SchedulerConfig,

36 - lora_config: LoRAConfig,

34 + def load_lora_model(self, model: nn.Module, vllm_config: VllmConfig,

35 device: torch.device) -> nn.Module:

37 if not supports_lora(model):

@@ -44,19 +42,12 @@ class LoRAModelRunnerMixin:

42 logger.warning("Regarding multimodal models, vLLM currently "

43 "only supports adding LoRA to language model.")

47 - # Use get_text_config() in case of multimodal models

48 - text_config = model_config.hf_config.get_text_config()

49 -

45 # Add LoRA Manager to the Model Runner

46 self.lora_manager = LRUCacheWorkerLoRAManager(

52 - scheduler_config.max_num_seqs,

53 - scheduler_config.max_num_batched_tokens,

54 - model_config.get_vocab_size(),

55 - lora_config,

47 + vllm_config,

48 device,

49 model.embedding_modules,

50 model.embedding_padding_modules,

59 - max_position_embeddings=text_config.max_position_embeddings,

51 )

52 return self.lora_manager.create_lora_manager(model)

vllm/v1/worker/tpu_model_runner.py

+1 -3

@@ -1178,9 +1178,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1178 "or sharding the weights on more chips. "

1179 f"See the detailed error: {e}") from e

1180 if self.lora_config is not None:

1181 - model = self.load_lora_model(model, self.model_config,

1182 - self.scheduler_config,

1183 - self.lora_config, self.device)

1181 + model = self.load_lora_model(model, self.vllm_config, self.device)

1182 replace_set_lora(model)

1183

1184 # Sync all pending XLA execution during model initialization and weight

vllm/worker/model_runner.py

+2 -9

@@ -1078,20 +1078,13 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):

1078 "Regarding multimodal models, vLLM currently "

1079 "only supports adding LoRA to language model.")

1080

1081 - # Use get_text_config() in case of multimodal models

1082 - text_config = self.model_config.hf_config.get_text_config()

1083 -

1081 self.lora_manager = LRUCacheWorkerLoRAManager(

1085 - self.scheduler_config.max_num_seqs,

1086 - self.scheduler_config.max_num_batched_tokens,

1087 - self.vocab_size,

1088 - self.lora_config,

1082 + self.vllm_config,

1083 self.device,

1084 self.model.embedding_modules,

1085 self.model.embedding_padding_modules,

1092 - max_position_embeddings=text_config.

1093 - max_position_embeddings,

1086 )

1087 +

1088 self.model = self.lora_manager.create_lora_manager(self.model)

1089 time_after_load = time.perf_counter()

1090