@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2f5f9acd55
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 12, 11:00 UTC

[LoRA] Continue optimizing MoE LoRA weight loading (#29322)

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

Jee Jee Li committed Nov 27, 2025 at 21:56 UTC 2f5f9acd551cfb737997a1f7f86982ec74aabf79
15 files changed +236 -165
tests/lora/test_lora_checkpoints.py
+8 -7
@@ -28,12 +28,13 @@ def test_load_checkpoints(
28 packed_modules_mapping = BaiChuanBaseForCausalLM.packed_modules_mapping
29 embedding_modules = BaiChuanBaseForCausalLM.embedding_modules
30 embed_padding_modules = BaiChuanBaseForCausalLM.embedding_padding_modules
31 - expected_lora_modules: list[str] = []
31 + expected_lora_lst: list[str] = []
32 for module in BAICHUAN_LORA_MODULES:
33 if module in packed_modules_mapping:
34 - expected_lora_modules.extend(packed_modules_mapping[module])
34 + expected_lora_lst.extend(packed_modules_mapping[module])
35 else:
36 - expected_lora_modules.append(module)
36 + expected_lora_lst.append(module)
37 + expected_lora_modules = set(expected_lora_lst)
38 if lora_name == "baichuan7B":
39 peft_helper = PEFTHelper.from_local_dir(
40 baichuan_lora_files, max_position_embeddings=4096
@@ -103,13 +104,13 @@ def test_lora_weights_mapping(baichuan_lora_files):
104 packed_modules_mapping = BaiChuanBaseForCausalLM.packed_modules_mapping
105 embedding_modules = BaiChuanBaseForCausalLM.embedding_modules
106 embed_padding_modules = BaiChuanBaseForCausalLM.embedding_padding_modules
106 - expected_lora_modules: list[str] = []
107 + expected_lora_lst: list[str] = []
108 for module in BAICHUAN_LORA_MODULES:
109 if module in packed_modules_mapping:
109 - expected_lora_modules.extend(packed_modules_mapping[module])
110 + expected_lora_lst.extend(packed_modules_mapping[module])
111 else:
111 - expected_lora_modules.append(module)
112 -
112 + expected_lora_lst.append(module)
113 + expected_lora_modules = set(expected_lora_lst)
114 hf_to_vllm_mapper = WeightsMapper(
115 orig_to_new_prefix={
116 "model.": "language_model.model.",
tests/lora/test_lora_huggingface.py
+4 -4
@@ -26,13 +26,13 @@ def test_load_checkpoints_from_huggingface(lora_fixture_name, request):
26 packed_modules_mapping = LlamaForCausalLM.packed_modules_mapping
27 embedding_modules = LlamaForCausalLM.embedding_modules
28 embed_padding_modules = LlamaForCausalLM.embedding_padding_modules
29 - expected_lora_modules: list[str] = []
29 + expected_lora_lst: list[str] = []
30 for module in LLAMA_LORA_MODULES:
31 if module in packed_modules_mapping:
32 - expected_lora_modules.extend(packed_modules_mapping[module])
32 + expected_lora_lst.extend(packed_modules_mapping[module])
33 else:
34 - expected_lora_modules.append(module)
35 -
34 + expected_lora_lst.append(module)
35 + expected_lora_modules = set(expected_lora_lst)
36 lora_path = get_adapter_absolute_path(lora_name)
37
38 # lora loading should work for either absolute path and huggingface id.
vllm/lora/layers/base.py
+1 -1
@@ -60,7 +60,7 @@ class BaseLayerWithLoRA(nn.Module):
60 source_layer: nn.Module,
61 lora_config: LoRAConfig,
62 packed_modules_list: list,
63 - model_config: PretrainedConfig | None,
63 + model_config: PretrainedConfig | None = None,
64 ) -> bool:
65 """Returns True if the layer can be replaced by this LoRA layer."""
66 raise NotImplementedError
vllm/lora/layers/column_parallel_linear.py
+8 -8
@@ -153,7 +153,7 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
153 source_layer: nn.Module,
154 lora_config: LoRAConfig,
155 packed_modules_list: list,
156 - model_config: PretrainedConfig | None,
156 + model_config: PretrainedConfig | None = None,
157 ) -> bool:
158 return type(source_layer) is ColumnParallelLinear or (
159 type(source_layer) is MergedColumnParallelLinear
@@ -272,7 +272,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
272 source_layer: nn.Module,
273 lora_config: LoRAConfig,
274 packed_modules_list: list,
275 - model_config: PretrainedConfig | None,
275 + model_config: PretrainedConfig | None = None,
276 ) -> bool:
277 return (
278 type(source_layer) is MergedColumnParallelLinear
@@ -338,7 +338,7 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
338 source_layer: nn.Module,
339 lora_config: LoRAConfig,
340 packed_modules_list: list,
341 - model_config: PretrainedConfig | None,
341 + model_config: PretrainedConfig | None = None,
342 ) -> bool:
343 return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 1
344
@@ -396,7 +396,7 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):
396 source_layer: nn.Module,
397 lora_config: LoRAConfig,
398 packed_modules_list: list,
399 - model_config: PretrainedConfig | None,
399 + model_config: PretrainedConfig | None = None,
400 ) -> bool:
401 return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 3
402
@@ -434,7 +434,7 @@ class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):
434 source_layer: nn.Module,
435 lora_config: LoRAConfig,
436 packed_modules_list: list,
437 - model_config: PretrainedConfig | None,
437 + model_config: PretrainedConfig | None = None,
438 ) -> bool:
439 # specifying kwargs so they can be easily accessed in decorator
440 return super().can_replace_layer(
@@ -480,7 +480,7 @@ class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLo
480 source_layer: nn.Module,
481 lora_config: LoRAConfig,
482 packed_modules_list: list,
483 - model_config: PretrainedConfig | None,
483 + model_config: PretrainedConfig | None = None,
484 ) -> bool:
485 # specifying kwargs so they can be easily accessed in decorator
486 return super().can_replace_layer(
@@ -516,7 +516,7 @@ class QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA):
516 source_layer: nn.Module,
517 lora_config: LoRAConfig,
518 packed_modules_list: list,
519 - model_config: PretrainedConfig | None,
519 + model_config: PretrainedConfig | None = None,
520 ) -> bool:
521 # specifying kwargs so they can be easily accessed in decorator
522 return super().can_replace_layer(
@@ -565,7 +565,7 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):
565 source_layer: nn.Module,
566 lora_config: LoRAConfig,
567 packed_modules_list: list,
568 - model_config: PretrainedConfig | None,
568 + model_config: PretrainedConfig | None = None,
569 ) -> bool:
570 # specifying kwargs so they can be easily accessed in decorator
571 return super().can_replace_layer(
vllm/lora/layers/fused_moe.py
+115 -103
@@ -401,6 +401,61 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
401 self.w13_lora_b_stacked[1][lora_id][experts_id]
402 )
403
404 + def _slice_w13_a(self, w13_lora_a: torch.Tensor) -> torch.Tensor:
405 + """
406 + Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA
407 + """
408 + if self.tp_size == 1 or not self.fully_sharded:
409 + return w13_lora_a
410 +
411 + # w13_lora_a shape (num_experts,rank,input_size)
412 + current_lora_rank = w13_lora_a.shape[1]
413 + assert current_lora_rank % self.tp_size == 0
414 + # Based on S-LoRA, we slice W13/W1/W3 A along the rank dim.
415 + sliced_rank = current_lora_rank // self.tp_size
416 + start_idx = self.tp_rank * sliced_rank
417 + end_idx = (self.tp_rank + 1) * sliced_rank
418 + return w13_lora_a[:, start_idx:end_idx, :]
419 +
420 + def _slice_w13_b(self, w13_lora_b: torch.Tensor):
421 + if self.tp_size == 1:
422 + return w13_lora_b
423 +
424 + # w13_lora_b shape (num_experts,output_size,rank)
425 + shard_size = self.base_layer.intermediate_size_per_partition
426 + start_idx = self.tp_rank * shard_size
427 + end_idx = (self.tp_rank + 1) * shard_size
428 +
429 + return w13_lora_b[:, start_idx:end_idx, :]
430 +
431 + def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor:
432 + """
433 + Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA
434 + """
435 + if self.tp_size == 1:
436 + return w2_lora_a
437 + # w2_lora_a shape (num_experts,rank,input_size)
438 + shard_size = self.base_layer.intermediate_size_per_partition
439 + start_idx = self.tp_rank * shard_size
440 + end_idx = (self.tp_rank + 1) * shard_size
441 +
442 + return w2_lora_a[:, :, start_idx:end_idx]
443 +
444 + def _slice_w2_b(self, w2_lora_b: torch.Tensor) -> torch.Tensor:
445 + """
446 + Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA
447 + """
448 + if self.tp_size == 1 or not self.fully_sharded:
449 + return w2_lora_b
450 + # Based on S-LoRA, we slice W2 B along the hidden_size dim.
451 + # w2_lora_b shape (num_experts,output_size,rank)
452 + current_lora_size = w2_lora_b.shape[1]
453 +
454 + sliced_size = current_lora_size // self.tp_size
455 + start_idx = self.tp_rank * sliced_size
456 + end_idx = (self.tp_rank + 1) * sliced_size
457 + return w2_lora_b[:, start_idx:end_idx, :]
458 +
459 def reset_lora(self, index: int):
460 """Resets the lora weights at index back to 0."""
461 for pos in range(self._w13_slices):
@@ -411,6 +466,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
466 self.w2_lora_b_stacked[0][index] = 0
467 self.adapter_enabled[index] = 0
468
469 + #
470 +
471 def set_lora(
472 self,
473 index: int,
@@ -418,69 +475,55 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
475 lora_b: torch.Tensor | list[torch.Tensor],
476 ):
477 """Overwrites lora tensors at index."""
478 + # Make mypy happy
479 assert isinstance(lora_a, list)
480 assert isinstance(lora_b, list)
481 +
482 self.reset_lora(index)
483 self.adapter_enabled[index] = 1
425 - for eid in range(len(lora_a) // 3):
426 - w1_lora_a = lora_a[eid * 3]
427 - w2_lora_a = lora_a[eid * 3 + 1]
428 - w3_lora_a = lora_a[eid * 3 + 2]
429 - w1_lora_b = lora_b[eid * 3]
430 - w2_lora_b = lora_b[eid * 3 + 1]
431 - w3_lora_b = lora_b[eid * 3 + 2]
432 -
433 - # Handle the case of adding LoRA to only a subset of experts
434 - if w1_lora_a is None or w2_lora_a is None or w3_lora_a is None:
435 - continue
436 -
437 - if self.tp_size > 1:
438 - shard_size = self.base_layer.intermediate_size_per_partition
439 - start_idx = self.tp_rank * shard_size
440 - end_idx = (self.tp_rank + 1) * shard_size
441 -
442 - w1_lora_b = w1_lora_b[start_idx:end_idx, :]
443 - w3_lora_b = w3_lora_b[start_idx:end_idx, :]
444 - w2_lora_a = w2_lora_a[:, start_idx:end_idx]
445 -
446 - if self.fully_sharded:
447 - # Based on S-LoRA, we slice W1 and W3 A along the rank dim,
448 - # and W2 B along the hidden_size dim.
449 - w13_shard_size = self.w13_lora_a_stacked[0][index, eid].shape[0]
450 - w13_start_idx = self.tp_rank * w13_shard_size
451 - w13_end_idx = (self.tp_rank + 1) * w13_shard_size
452 - w1_lora_a = w1_lora_a[w13_start_idx:w13_end_idx, :]
453 - w3_lora_a = w3_lora_a[w13_start_idx:w13_end_idx, :]
454 -
455 - w2_shard_size = self.w2_lora_b_stacked[0][index, eid].shape[0]
456 - w2_start_idx = self.tp_rank * w2_shard_size
457 - w2_end_idx = (self.tp_rank + 1) * w2_shard_size
458 - w2_lora_b = w2_lora_b[w2_start_idx:w2_end_idx, :]
459 - # w1 lora_a
460 - self.w13_lora_a_stacked[0][
461 - index, eid, : w1_lora_a.shape[0], : w1_lora_a.shape[1]
462 - ].copy_(w1_lora_a, non_blocking=True)
463 - # w3 lora_a
464 - self.w13_lora_a_stacked[1][
465 - index, eid, : w3_lora_a.shape[0], : w3_lora_a.shape[1]
466 - ].copy_(w3_lora_a, non_blocking=True)
467 -
468 - # w1 lora_b
469 - self.w13_lora_b_stacked[0][
470 - index, eid, : w1_lora_b.shape[0], : w1_lora_b.shape[1]
471 - ].copy_(w1_lora_b, non_blocking=True)
472 - # w3 lora_b
473 - self.w13_lora_b_stacked[1][
474 - index, eid, : w3_lora_b.shape[0], : w3_lora_b.shape[1]
475 - ].copy_(w3_lora_b, non_blocking=True)
476 -
477 - self.w2_lora_a_stacked[0][
478 - index, eid, : w2_lora_a.shape[0], : w2_lora_a.shape[1]
479 - ].copy_(w2_lora_a, non_blocking=True)
480 -
481 - self.w2_lora_b_stacked[0][
482 - index, eid, : w2_lora_b.shape[0], : w2_lora_b.shape[1]
483 - ].copy_(w2_lora_b, non_blocking=True)
484 +
485 + num_experts = self.w13_lora_a_stacked[0].shape[1]
486 +
487 + w1_lora_a, w2_lora_a, w3_lora_a = lora_a
488 + w1_lora_b, w2_lora_b, w3_lora_b = lora_b
489 + assert (
490 + num_experts
491 + == w1_lora_a.shape[0]
492 + == w2_lora_a.shape[0]
493 + == w3_lora_a.shape[0]
494 + )
495 +
496 + slliced_w1_lora_a = self._slice_w13_a(w1_lora_a)
497 + slliced_w1_lora_b = self._slice_w13_b(w1_lora_b)
498 + slliced_w3_lora_a = self._slice_w13_a(w3_lora_a)
499 + slliced_w3_lora_b = self._slice_w13_b(w3_lora_b)
500 +
501 + sliced_w2_lora_a = self._slice_w2_a(w2_lora_a)
502 + sliced_w2_lora_b = self._slice_w2_b(w2_lora_b)
503 +
504 + self.w13_lora_a_stacked[0][
505 + index, :, : slliced_w1_lora_a.shape[1], : slliced_w1_lora_a.shape[2]
506 + ].copy_(slliced_w1_lora_a, non_blocking=True)
507 +
508 + self.w13_lora_a_stacked[1][
509 + index, :, : slliced_w3_lora_a.shape[1], : slliced_w3_lora_a.shape[2]
510 + ].copy_(slliced_w3_lora_a, non_blocking=True)
511 +
512 + self.w13_lora_b_stacked[0][
513 + index, :, : slliced_w1_lora_b.shape[1], : slliced_w1_lora_b.shape[2]
514 + ].copy_(slliced_w1_lora_b, non_blocking=True)
515 +
516 + self.w13_lora_b_stacked[1][
517 + index, :, : slliced_w3_lora_b.shape[1], : slliced_w3_lora_b.shape[2]
518 + ].copy_(slliced_w3_lora_b, non_blocking=True)
519 +
520 + self.w2_lora_a_stacked[0][
521 + index, :, : sliced_w2_lora_a.shape[1], : sliced_w2_lora_a.shape[2]
522 + ].copy_(sliced_w2_lora_a, non_blocking=True)
523 +
524 + self.w2_lora_b_stacked[0][
525 + index, :, : sliced_w2_lora_b.shape[1], : sliced_w2_lora_b.shape[2]
526 + ].copy_(sliced_w2_lora_b, non_blocking=True)
527
528 def forward(self, *args, **kwargs):
529 return self.base_layer.forward(*args, **kwargs)
@@ -506,12 +549,12 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
549 source_layer: nn.Module,
550 lora_config: LoRAConfig,
551 packed_modules_list: list,
509 - model_config: PretrainedConfig | None,
552 + model_config: PretrainedConfig | None = None,
553 ) -> bool:
554 """Returns True if the layer can be replaced by this LoRA layer."""
512 - # return type(source_layer) is FusedMoE
555
514 - return type(source_layer) is FusedMoE and len(packed_modules_list) == 2
556 + # source_layer is FusedMoE or SharedFusedMoE
557 + return isinstance(source_layer, FusedMoE) and len(packed_modules_list) == 2
558
559
560 class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
@@ -555,6 +598,9 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
598 model_config: PretrainedConfig | None = None,
599 ) -> None:
600 """Initializes lora matrices."""
601 +
602 + assert isinstance(model_config, PretrainedConfig)
603 + self._base_model = model_config.architectures[0]
604 self.max_loras = lora_config.max_loras
605 self.fully_sharded = lora_config.fully_sharded_loras
606
@@ -565,20 +611,7 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
611 self._create_lora_a_weights(max_loras, lora_config)
612 self._create_lora_b_weights(max_loras, lora_config)
613
568 - def _slice_w13_a(self, w13_lora_a: torch.Tensor) -> torch.Tensor:
569 - if self.tp_size == 1 or not self.fully_sharded:
570 - return w13_lora_a
571 -
572 - # w13_lora_a shape (num_experts,rank,input_size)
573 - current_lora_rank = w13_lora_a.shape[1]
574 - assert current_lora_rank % self.tp_size == 0
575 -
576 - sliced_rank = current_lora_rank // self.tp_size
577 - start_idx = self.tp_rank * sliced_rank
578 - end_idx = (self.tp_rank + 1) * sliced_rank
579 - return w13_lora_a[:, start_idx:end_idx, :]
580 -
581 - def _slice_w13_b(self, w13_lora_b: torch.Tensor, is_interleave: bool = True):
614 + def _slice_w13_b(self, w13_lora_b: torch.Tensor):
615 if self.tp_size == 1:
616 return w13_lora_b
617
@@ -586,7 +619,8 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
619 shard_size = self.base_layer.intermediate_size_per_partition
620 start_idx = self.tp_rank * shard_size
621 end_idx = (self.tp_rank + 1) * shard_size
589 - if is_interleave:
622 + # HACK: Currently, only GPT-OSS is in interleaved order
623 + if self._base_model == "GptOssForCausalLM":
624 # For models like GPT-OSS, the weights of w1 (gate_proj) and w3 (up_proj)
625 # in the interleaved order, and corresponding LoRA need to be processed.
626 w1_lora_b = w13_lora_b[:, ::2, :]
@@ -606,28 +640,6 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
640
641 return torch.cat([sliced_w1_lora_b, sliced_w3_lora_b], dim=1)
642
609 - def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor:
610 - if self.tp_size == 1:
611 - return w2_lora_a
612 - # w2_lora_a shape (num_experts,rank,input_size)
613 - shard_size = self.base_layer.intermediate_size_per_partition
614 - start_idx = self.tp_rank * shard_size
615 - end_idx = (self.tp_rank + 1) * shard_size
616 -
617 - return w2_lora_a[:, :, start_idx:end_idx]
618 -
619 - def _slice_w2_b(self, w2_lora_b: torch.Tensor) -> torch.Tensor:
620 - if self.tp_size == 1 or not self.fully_sharded:
621 - return w2_lora_b
622 - # Based on S-LoRA, we slice W2 B along the hidden_size dim.
623 - # w2_lora_b shape (num_experts,output_size,rank)
624 - current_lora_size = w2_lora_b.shape[1]
625 -
626 - sliced_size = current_lora_size // self.tp_size
627 - start_idx = self.tp_rank * sliced_size
628 - end_idx = (self.tp_rank + 1) * sliced_size
629 - return w2_lora_b[:, start_idx:end_idx, :]
630 -
643 def set_lora(
644 self,
645 index: int,
@@ -658,7 +670,7 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
670 w2_lora_b = w2_lora_b.permute(1, 0, 2)
671
672 sliced_w13_lora_a = self._slice_w13_a(w13_lora_a)
661 - sliced_w13_lora_b = self._slice_w13_b(w13_lora_b, is_interleave=True)
673 + sliced_w13_lora_b = self._slice_w13_b(w13_lora_b)
674
675 sliced_w2_lora_a = self._slice_w2_a(w2_lora_a)
676 sliced_w2_lora_b = self._slice_w2_b(w2_lora_b)
@@ -711,8 +723,8 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
723 source_layer: nn.Module,
724 lora_config: LoRAConfig,
725 packed_modules_list: list,
714 - model_config: PretrainedConfig | None,
726 + model_config: PretrainedConfig | None = None,
727 ) -> bool:
728 """Returns True if the layer can be replaced by this LoRA layer."""
717 -
718 - return type(source_layer) is FusedMoE and len(packed_modules_list) == 1
729 + # source_layer is FusedMoE or SharedFusedMoE
730 + return isinstance(source_layer, FusedMoE) and len(packed_modules_list) == 1
vllm/lora/layers/logits_processor.py
+1 -1
@@ -197,7 +197,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
197 source_layer: nn.Module,
198 lora_config: LoRAConfig,
199 packed_modules_list: list,
200 - model_config: PretrainedConfig | None,
200 + model_config: PretrainedConfig | None = None,
201 ) -> bool:
202 # Special handling for the LogitsProcessor.
203 return False
vllm/lora/layers/replicated_linear.py
+1 -1
@@ -53,7 +53,7 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):
53 source_layer: nn.Module,
54 lora_config: LoRAConfig,
55 packed_modules_list: list,
56 - model_config: PretrainedConfig | None,
56 + model_config: PretrainedConfig | None = None,
57 ) -> bool:
58 return type(source_layer) is ReplicatedLinear
59
vllm/lora/layers/row_parallel_linear.py
+2 -2
@@ -87,7 +87,7 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
87 source_layer: nn.Module,
88 lora_config: LoRAConfig,
89 packed_modules_list: list,
90 - model_config: PretrainedConfig | None,
90 + model_config: PretrainedConfig | None = None,
91 ) -> bool:
92 return type(source_layer) is RowParallelLinear
93
@@ -164,7 +164,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):
164 source_layer: nn.Module,
165 lora_config: LoRAConfig,
166 packed_modules_list: list,
167 - model_config: PretrainedConfig | None,
167 + model_config: PretrainedConfig | None = None,
168 ) -> bool:
169 # specifying kwargs so they can be easily accessed in decorator
170 return super().can_replace_layer(
vllm/lora/layers/vocal_parallel_embedding.py
+1 -1
@@ -131,7 +131,7 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
131 source_layer: nn.Module,
132 lora_config: LoRAConfig,
133 packed_modules_list: list,
134 - model_config: PretrainedConfig | None,
134 + model_config: PretrainedConfig | None = None,
135 ) -> bool:
136 return type(source_layer) is VocabParallelEmbedding
137
vllm/lora/lora_weights.py
+53
@@ -152,6 +152,59 @@ class PackedLoRALayerWeights(LoRALayerWeights):
152 )
153 return obj
154
155 + @classmethod
156 + def pack_moe(
157 + cls, loras: GenericSequence[Optional["LoRALayerWeights"]], module_name: str
158 + ) -> "PackedLoRALayerWeights":
159 + """Pack a list of LoRAs into a single LoRA.
160 +
161 + If LoRA is None, it signifies that the submodule does not have a LoRA.
162 + """
163 +
164 + first_lora = next(lora for lora in loras if lora is not None)
165 + assert first_lora is not None
166 + rank = first_lora.rank
167 + lora_alpha = first_lora.lora_alpha
168 + assert len(loras) % 3 == 0
169 + w1_lora_a_lst = []
170 + w2_lora_a_lst = []
171 + w3_lora_a_lst = []
172 + w1_lora_b_lst = []
173 + w2_lora_b_lst = []
174 + w3_lora_b_lst = []
175 + # TODO: Consider the case where some experts don't have LoRA added.
176 + for eid in range(len(loras) // 3):
177 + w1_lora = loras[eid * 3]
178 + w2_lora = loras[eid * 3 + 1]
179 + w3_lora = loras[eid * 3 + 2]
180 + assert w1_lora is not None
181 + assert w2_lora is not None
182 + assert w3_lora is not None
183 +
184 + w1_lora_a_lst.append(w1_lora.lora_a)
185 + w2_lora_a_lst.append(w2_lora.lora_a)
186 + w3_lora_a_lst.append(w3_lora.lora_a)
187 +
188 + w1_lora_b_lst.append(w1_lora.lora_b)
189 + w2_lora_b_lst.append(w2_lora.lora_b)
190 + w3_lora_b_lst.append(w3_lora.lora_b)
191 +
192 + w1_lora_a = torch.stack(w1_lora_a_lst, dim=0) # (num_experts,rank,input_size)
193 + w2_lora_a = torch.stack(w2_lora_a_lst, dim=0)
194 + w3_lora_a = torch.stack(w3_lora_a_lst, dim=0)
195 + w1_lora_b = torch.stack(w1_lora_b_lst, dim=0) # (num_experts,output_size,rank)
196 + w2_lora_b = torch.stack(w2_lora_b_lst, dim=0)
197 + w3_lora_b = torch.stack(w3_lora_b_lst, dim=0)
198 +
199 + obj = cls(
200 + module_name,
201 + rank,
202 + [lora_alpha, lora_alpha, lora_alpha],
203 + [w1_lora_a, w2_lora_a, w3_lora_a],
204 + [w1_lora_b, w2_lora_b, w3_lora_b],
205 + )
206 + return obj
207 +
208 def optimize(self) -> "PackedLoRALayerWeights":
209 """Optimize the LoRA by merging the scaling into lora_b."""
210 for i in range(len(self.lora_b)):
vllm/lora/models.py
+27 -23
@@ -13,7 +13,7 @@ from torch import nn
13
14 from vllm.config.lora import LoRAConfig
15 from vllm.logger import init_logger
16 -from vllm.lora.layers import BaseLayerWithLoRA, FusedMoEWithLoRA, LoRAMapping
16 +from vllm.lora.layers import BaseLayerWithLoRA, FusedMoE3DWithLoRA, LoRAMapping
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
@@ -151,16 +151,13 @@ class LoRAModel:
151 if pin_memory:
152 loras[module_name].lora_b = loras[module_name].lora_b.pin_memory()
153
154 - for lora in loras.values():
155 - lora.optimize()
156 -
154 return cls(lora_model_id, peft_helper.r, loras)
155
156 @classmethod
157 def from_local_checkpoint(
158 cls,
159 lora_dir: str,
163 - expected_lora_modules: list[str],
160 + expected_lora_modules: set[str],
161 peft_helper: PEFTHelper,
162 *,
163 lora_model_id: int | None = None,
@@ -190,10 +187,7 @@ class LoRAModel:
187 lora_tensor_path = os.path.join(lora_dir, "adapter_model.safetensors")
188 lora_bin_file_path = os.path.join(lora_dir, "adapter_model.bin")
189 lora_pt_file_path = os.path.join(lora_dir, "adapter_model.pt")
193 - # new_embeddings_tensor_path = os.path.join(
194 - # lora_dir, "new_embeddings.safetensors"
195 - # )
196 - # new_embeddings_bin_file_path = os.path.join(lora_dir, "new_embeddings.bin")
190 +
191 tensors: dict[str, torch.Tensor] = {}
192 unexpected_modules: list[list[str] | str] = []
193
@@ -201,18 +195,19 @@ class LoRAModel:
195 for lora_module in modules.keys(): # noqa
196 if is_base_embeddding_weights(lora_module):
197 continue
204 - module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper)
205 - # Handle FSDP file format where experts.base_layer is the
198 + # Handle PEFT file format where experts.base_layer is the
199 # gate_up_proj and experts is the down_proj
200 if "base_layer" in lora_module:
201 continue
202 + module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper)
203 # Case for expert lora weights
204 if ".experts" in module_name:
211 - if not any(
212 - module_name.endswith(ele) for ele in expected_lora_modules
213 - ):
205 + expert_idx = module_name.find(".experts")
206 + expert_suffix = module_name[expert_idx + 1 :]
207 + if expert_suffix not in expected_lora_modules:
208 unexpected_modules.append(module_name)
215 - elif module_name.split(".")[-1] not in expected_lora_modules:
209 +
210 + elif module_name.rsplit(".", 1)[-1] not in expected_lora_modules:
211 unexpected_modules.append(module_name)
212
213 if unexpected_modules:
@@ -358,9 +353,7 @@ class LoRAModelManager:
353 self.modules: dict[str, BaseLayerWithLoRA] = {}
354 # Dict instead of a set for compatibility with LRUCache.
355 self._last_mapping: LoRAMapping | None = None
361 - self._is_3d_moe_model = is_moe_model(self.model) and hasattr(
362 - self.model, "is_3d_moe_weight"
363 - )
356 + self._is_3d_moe_model = is_moe_model(self.model) and self.model.is_3d_moe_weight
357 self._create_lora_modules()
358
359 self.model.lora_manager = self
@@ -411,7 +404,7 @@ class LoRAModelManager:
404 continue
405 # Note (gnovack) - If MOE lora weights are not split into
406 # num_experts chunks, we split them here
414 - if isinstance(module, FusedMoEWithLoRA) and torch.is_tensor(
407 + if isinstance(module, FusedMoE3DWithLoRA) and torch.is_tensor(
408 module_lora.lora_a
409 ):
410 # Handle PEFT file format where experts.base_layer is the
@@ -679,7 +672,10 @@ class LoRAModelManager:
672 "cpu",
673 )
674 subloras.append(lora)
682 - lora = PackedLoRALayerWeights.pack(subloras)
675 + if module.__class__.__name__ == "FusedMoEWithLoRA":
676 + lora = PackedLoRALayerWeights.pack_moe(subloras, module_name)
677 + else:
678 + lora = PackedLoRALayerWeights.pack(subloras)
679 model.loras[module_name] = lora
680 return model
681
@@ -739,13 +735,21 @@ class LoRAModelManager:
735 replaced_module_name = module_name.replace("model.", "")
736 if lora_model.check_lora_name(module_name):
737 module_name = replaced_module_name
742 - lora_model.loras[module_name] = PackedLoRALayerWeights.pack(
743 - replacement_loras
744 - )
738 + if module_name.endswith(".experts"):
739 + lora_model.loras[module_name] = PackedLoRALayerWeights.pack_moe(
740 + replacement_loras, module_name
741 + )
742 + else:
743 + lora_model.loras[module_name] = PackedLoRALayerWeights.pack(
744 + replacement_loras
745 + )
746 # Remove the modules that have been replaced.
747 for module in replaced_module:
748 lora_model.loras.pop(module, None)
749
750 + for lora in lora_model.loras.values():
751 + lora.optimize()
752 +
753 def _get_lora_layer_weights(
754 self, lora_model: LoRAModel, module_name: str
755 ) -> LoRALayerWeights | None:
vllm/lora/utils.py
+8 -9
@@ -170,16 +170,15 @@ def parse_fine_tuned_lora_name(
170
171 def is_base_embeddding_weights(name: str) -> bool:
172 # hardcoded subfixes for input & output embedding weights
173 - input_embedding_subfix = ".embed_tokens.base_layer.weight"
174 - output_embedding_subfix = ".lm_head.base_layer.weight"
175 -
176 - return name.endswith(input_embedding_subfix) or name.endswith(
177 - output_embedding_subfix
173 + embedding_suffixes = (
174 + ".embed_tokens.base_layer.weight",
175 + ".lm_head.base_layer.weight",
176 )
177 + return name.endswith(embedding_suffixes)
178
179
180 def is_regex_target_modules(
182 - load_modules: str | list[str], expected_lora_modules: list[str]
181 + load_modules: str | list[str], expected_lora_modules: set[str]
182 ) -> bool:
183 """
184 PEFT supports passing `target_modules` in the form of regular expressions,
@@ -195,8 +194,8 @@ def is_regex_target_modules(
194 except re.error:
195 return False
196
198 - def is_subset(sub_list, full_list):
199 - return set(sub_list).issubset(set(full_list))
197 + def is_subset(sub_list, full_set):
198 + return set(sub_list).issubset(full_set)
199
200 # Similar to PEFT's processing logic, regex-related operations are only
201 # executed when the load_modules is a `str`.
@@ -290,7 +289,7 @@ def process_packed_modules_mapping(model: nn.Module) -> dict[str, list[str]]:
289 # the expert indices are expanded based on the configured number
290 # of routed experts.
291 packed_modules_mapping = get_packed_modules_mapping(model)
293 - if not hasattr(model, "is_3d_moe_weight"):
292 + if not model.is_3d_moe_weight:
293 # 3D MoE LoRA does not need `packed_modules_mapping`
294 packed_modules_mapping["experts"] = [
295 weight_name.rstrip(".")
vllm/lora/worker_manager.py
+5 -5
@@ -88,15 +88,15 @@ class WorkerLoRAManager:
88 try:
89 supported_lora_modules = self._adapter_manager.supported_lora_modules
90 packed_modules_mapping = self._adapter_manager.packed_modules_mapping
91 - expected_lora_modules: list[str] = []
91 + expected_lora_lst: list[str] = []
92 for module in supported_lora_modules:
93 if module in packed_modules_mapping:
94 - expected_lora_modules.extend(packed_modules_mapping[module])
94 + expected_lora_lst.extend(packed_modules_mapping[module])
95 else:
96 - expected_lora_modules.append(module)
96 + expected_lora_lst.append(module)
97 if module == "experts":
98 - expected_lora_modules.append(module)
99 - expected_lora_modules = list(set(expected_lora_modules))
98 + expected_lora_lst.append(module)
99 + expected_lora_modules = set(expected_lora_lst)
100 lora_path = get_adapter_absolute_path(lora_request.lora_path)
101
102 peft_helper = PEFTHelper.from_local_dir(
vllm/model_executor/models/interfaces.py
+1
@@ -336,6 +336,7 @@ class SupportsLoRA(Protocol):
336 There is no need to redefine this flag if this class is in the
337 MRO of your model class.
338 """
339 + is_3d_moe_weight: ClassVar[bool] = False
340 # The `embedding_module` and `embedding_padding_modules`
341 # are empty by default.
342 embedding_modules: ClassVar[dict[str, str]] = {}
vllm/model_executor/models/qwen3_vl_moe.py
+1
@@ -401,6 +401,7 @@ class Qwen3VLMoeMixtureOfExperts(MixtureOfExperts):
401 class Qwen3VLMoeForConditionalGeneration(
402 Qwen3VLForConditionalGeneration, Qwen3VLMoeMixtureOfExperts
403 ):
404 + is_3d_moe_weight: bool = True
405 packed_modules_mapping = {
406 "qkv_proj": [
407 "q_proj",