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

Fork 0

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

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

Synced Nov 15, 14:45 UTC

[Core][LoRA][1/N] Add LoRA for EncoderDecoderModelRunner (#15990)

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

Jee Jee Li committed Apr 11, 2025 at 15:32 UTC f7030df3be651bbce42932be736129d37caa856b

3 files changed +41 -1

vllm/lora/layers.py

@@ -866,6 +866,11 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):

866 and len(packed_modules_list) == 3)

867

868

869 +#TODO: Implement this

870 +class QKVCrossParallelLinearWithLoRA(BaseLayerWithLoRA):

871 + pass

872 +

873 +

874 class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

875

876 def __init__(self, base_layer: RowParallelLinear) -> None:

vllm/model_executor/models/mllama.py

+11

@@ -52,6 +52,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

52 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

53 from vllm.model_executor.model_loader.weight_utils import (

54 default_weight_loader, maybe_remap_kv_scale_name)

55 +from vllm.model_executor.models.module_mapping import MultiModelKeys

56 from vllm.model_executor.sampling_metadata import SamplingMetadata

57 from vllm.multimodal import MULTIMODAL_REGISTRY

58 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalEncDecInputs,

@@ -1181,6 +1182,7 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1182 super().__init__()

1183 config: MllamaConfig = vllm_config.model_config.hf_config

1184 quant_config = vllm_config.quant_config

1185 + self.config = config

1186 self.quant_config = quant_config

1187 self.vocab_size = config.text_config.vocab_size

1188 self.hidden_size = config.text_config.hidden_size

@@ -1517,6 +1519,15 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal,

1519 updated_params.add(name)

1520 return updated_params

1521

1522 + def get_mm_mapping(self) -> MultiModelKeys:

1523 + """

1524 + Get the module prefix in multimodal models

1525 + """

1526 + return MultiModelKeys.from_string_field(

1527 + language_model="language_model",

1528 + connector="multi_modal_projector",

1529 + tower_model="vision_model")

1530 +

1531

1532 def skip_attention_mask(sparse_mask: List[List[int]]) -> bool:

1533 for mask in sparse_mask:

vllm/worker/enc_dec_model_runner.py

+25 -1

@@ -16,6 +16,7 @@ from vllm.config import VllmConfig

16 from vllm.forward_context import set_forward_context

17 from vllm.inputs import INPUT_REGISTRY, InputRegistry

18 from vllm.logger import init_logger

19 +from vllm.lora.request import LoRARequest

20 from vllm.model_executor import SamplingMetadata

21 from vllm.model_executor.layers.sampler import SamplerOutput

22 from vllm.multimodal import (MULTIMODAL_REGISTRY, MultiModalKwargs,

@@ -34,6 +35,7 @@ from vllm.worker.model_runner_base import (

35 from vllm.worker.utils import assert_enc_dec_mr_supported_scenario

37 logger = init_logger(__name__)

38 +LORA_WARMUP_RANK = 8

41 @dataclasses.dataclass(frozen=True)

@@ -160,7 +162,11 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):

162 if num_steps > 1:

163 raise ValueError("num_steps > 1 is not supported in "

164 "EncoderDecoderModelRunner")

163 -

165 + if self.lora_config:

166 + assert model_input.lora_requests is not None

167 + assert model_input.lora_mapping is not None

168 + self.set_active_loras(model_input.lora_requests,

169 + model_input.lora_mapping)

170 if (model_input.attn_metadata is not None

171 and model_input.attn_metadata.prefill_metadata is None

172 and model_input.attn_metadata.decode_metadata.use_cuda_graph):

@@ -268,6 +274,22 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):

274 max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens

275 max_num_seqs = self.scheduler_config.max_num_seqs

276

277 + # This represents the maximum number of different requests

278 + # that will have unique loras, and therefore the max amount of

279 + # memory consumption. Create dummy lora request copies from the

280 + # lora request passed in, which contains a lora from the lora

281 + # warmup path.

282 + dummy_lora_requests: List[LoRARequest] = []

283 + dummy_lora_requests_per_seq: List[LoRARequest] = []

284 + if self.lora_config:

285 + dummy_lora_requests = self._add_dummy_loras(

286 + self.lora_config.max_loras)

287 + assert len(dummy_lora_requests) == self.lora_config.max_loras

288 + dummy_lora_requests_per_seq = [

289 + dummy_lora_requests[idx % len(dummy_lora_requests)]

290 + for idx in range(max_num_seqs)

291 + ]

292 +

293 # Profile memory usage with max_num_sequences sequences and the total

294 # number of tokens equal to max_num_batched_tokens.

295 seqs: List[SequenceGroupMetadata] = []

@@ -315,6 +337,8 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):

337 block_tables=None,

338 encoder_seq_data=encoder_dummy_data.seq_data,

339 cross_block_table=None,

340 + lora_request=dummy_lora_requests_per_seq[group_id]

341 + if dummy_lora_requests_per_seq else None,

342 multi_modal_data=decoder_dummy_data.multi_modal_data

343 or encoder_dummy_data.multi_modal_data,

344 multi_modal_placeholders=decoder_dummy_data.