@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1073ba68b0
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, 16:00 UTC

[LoRA] Optimize 3D MoE logic (#29222)

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

Jee Jee Li committed Nov 24, 2025 at 10:27 UTC 1073ba68b0a741de6de0a7230795acba151084f2
11 files changed +395 -103
tests/lora/test_gptoss_tp.py
+6 -1
@@ -1,6 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import pytest
5 +
6 import vllm
7 from vllm.lora.request import LoRARequest
8
@@ -84,14 +86,17 @@ def test_gpt_oss_lora(gptoss20b_lora_files):
86
87
88 @multi_gpu_test(num_gpus=2)
87 -def test_gpt_oss_lora_tp2(gptoss20b_lora_files):
89 +@pytest.mark.parametrize("fully_sharded_loras", [False, True])
90 +def test_gpt_oss_lora_tp2(gptoss20b_lora_files, fully_sharded_loras):
91 llm = vllm.LLM(
92 MODEL_PATH,
93 max_model_len=1024,
94 enable_lora=True,
95 max_loras=2,
96 max_lora_rank=8,
97 + max_num_seqs=16,
98 tensor_parallel_size=2,
99 + fully_sharded_loras=fully_sharded_loras,
100 compilation_config=vllm.config.CompilationConfig( # Avoid OOM
101 cudagraph_specialize_lora=False,
102 ),
vllm/lora/layers/__init__.py
+2 -1
@@ -11,7 +11,7 @@ from vllm.lora.layers.column_parallel_linear import (
11 QKVParallelLinearWithLoRA,
12 QKVParallelLinearWithShardedLoRA,
13 )
14 -from vllm.lora.layers.fused_moe import FusedMoEWithLoRA
14 +from vllm.lora.layers.fused_moe import FusedMoE3DWithLoRA, FusedMoEWithLoRA
15 from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
16 from vllm.lora.layers.replicated_linear import ReplicatedLinearWithLoRA
17 from vllm.lora.layers.row_parallel_linear import (
@@ -38,4 +38,5 @@ __all__ = [
38 "ReplicatedLinearWithLoRA",
39 "LoRAMapping",
40 "FusedMoEWithLoRA",
41 + "FusedMoE3DWithLoRA",
42 ]
vllm/lora/layers/base.py
+2 -2
@@ -42,8 +42,8 @@ class BaseLayerWithLoRA(nn.Module):
42 def set_lora(
43 self,
44 index: int,
45 - lora_a: torch.Tensor,
46 - lora_b: torch.Tensor,
45 + lora_a: torch.Tensor | list[torch.Tensor],
46 + lora_b: torch.Tensor | list[torch.Tensor],
47 ):
48 """Overwrites lora tensors at index."""
49 ...
vllm/lora/layers/base_linear.py
+4 -2
@@ -94,13 +94,15 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
94 def set_lora(
95 self,
96 index: int,
97 - lora_a: torch.Tensor,
98 - lora_b: torch.Tensor,
97 + lora_a: torch.Tensor | list[torch.Tensor],
98 + lora_b: torch.Tensor | list[torch.Tensor],
99 ):
100 # Except for QKVParallelLinearWithLoRA and
101 # MergedColumnParallelLinearWithLoRA, all other linear LoRA layers
102 # store weights in a tuple of size 1. These two layers will
103 # override this function.
104 + assert isinstance(lora_a, torch.Tensor)
105 + assert isinstance(lora_b, torch.Tensor)
106 assert (
107 len(self.lora_a_stacked) == len(self.lora_b_stacked) == self.n_slices == 1
108 )
vllm/lora/layers/column_parallel_linear.py
+2 -2
@@ -246,8 +246,8 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
246 def set_lora(
247 self,
248 index: int,
249 - lora_a: torch.Tensor,
250 - lora_b: torch.Tensor,
249 + lora_a: torch.Tensor | list[torch.Tensor],
250 + lora_b: torch.Tensor | list[torch.Tensor],
251 ):
252 self.reset_lora(index)
253
vllm/lora/layers/fused_moe.py
+286 -63
@@ -42,7 +42,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
42 self.tp_size = get_tensor_model_parallel_world_size()
43 self.tp_rank = get_tensor_model_parallel_rank()
44 self.device = base_layer.w2_weight.device
45 - self.w13_slices = 2
45 + self._w13_slices = 2
46 self._inject_lora_into_fused_moe()
47
48 def _normalize_keys(self, config: dict[str, int | None]) -> dict[str, int | None]:
@@ -160,7 +160,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
160 op_prefix="w13",
161 num_loras=self.max_loras,
162 rank=max_lora_rank,
163 - num_slices=self.w13_slices,
163 + num_slices=self._w13_slices,
164 M=M,
165 layer=layer,
166 top_k=top_k,
@@ -230,7 +230,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
230 CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
231 num_tokens = hidden_states.size(0)
232 M = min(num_tokens, CHUNK_SIZE)
233 - max_lora_rank = self.w2_lora_a_stacked.shape[-2]
233 + max_lora_rank = self.w2_lora_a_stacked[0].shape[-2]
234 shrink_config, expand_config = self._get_lora_moe_configs(
235 op_prefix="w2",
236 num_loras=self.max_loras,
@@ -258,8 +258,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
258 self.punica_wrapper.add_lora_fused_moe(
259 intermediate_cache3,
260 intermediate_cache2,
261 - (self.w2_lora_a_stacked,),
262 - (self.w2_lora_b_stacked,),
261 + self.w2_lora_a_stacked,
262 + self.w2_lora_b_stacked,
263 topk_weights,
264 sorted_token_ids_lora,
265 expert_ids_lora,
@@ -292,22 +292,12 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
292 self.base_layer.quant_method, m_fused_moe_fn
293 )
294
295 - def create_lora_weights(
295 + def _create_lora_a_weights(
296 self,
297 max_loras: int,
298 lora_config: LoRAConfig,
299 - model_config: PretrainedConfig | None = None,
300 - ) -> None:
301 - """Initializes lora matrices."""
302 - assert self.w13_slices == 2
303 - self.max_loras = lora_config.max_loras
304 - self.fully_sharded = lora_config.fully_sharded_loras
305 -
306 - self.adapter_enabled = torch.tensor(
307 - [0] * (max_loras + 1), dtype=torch.int, device=self.device
308 - )
309 -
310 - self.w13_lora_a_stacked = tuple(
299 + ):
300 + self.w13_lora_a_stacked: tuple[torch.Tensor, ...] = tuple(
301 torch.zeros(
302 (
303 max_loras,
@@ -320,10 +310,23 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
310 dtype=lora_config.lora_dtype,
311 device=self.device,
312 )
323 - for _ in range(self.w13_slices)
313 + for _ in range(self._w13_slices)
314 + )
315 + self.w2_lora_a_stacked: tuple[torch.Tensor, ...] = (
316 + torch.zeros(
317 + (
318 + max_loras,
319 + self.base_layer.local_num_experts,
320 + lora_config.max_lora_rank,
321 + self.base_layer.intermediate_size_per_partition,
322 + ),
323 + dtype=lora_config.lora_dtype,
324 + device=self.device,
325 + ),
326 )
327
326 - self.w13_lora_b_stacked = tuple(
328 + def _create_lora_b_weights(self, max_loras: int, lora_config: LoRAConfig):
329 + self.w13_lora_b_stacked: tuple[torch.Tensor, ...] = tuple(
330 torch.zeros(
331 (
332 max_loras,
@@ -334,34 +337,42 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
337 dtype=lora_config.lora_dtype,
338 device=self.device,
339 )
337 - for _ in range(self.w13_slices)
340 + for _ in range(self._w13_slices)
341 )
339 -
340 - self.w2_lora_a_stacked = torch.zeros(
341 - (
342 - max_loras,
343 - self.base_layer.local_num_experts,
344 - lora_config.max_lora_rank,
345 - self.base_layer.intermediate_size_per_partition,
342 + self.w2_lora_b_stacked: tuple[torch.Tensor, ...] = (
343 + torch.zeros(
344 + (
345 + max_loras,
346 + self.base_layer.local_num_experts,
347 + self.base_layer.hidden_size
348 + if not self.fully_sharded
349 + else divide(self.base_layer.hidden_size, self.tp_size),
350 + lora_config.max_lora_rank,
351 + ),
352 + dtype=lora_config.lora_dtype,
353 + device=self.device,
354 ),
347 - dtype=lora_config.lora_dtype,
348 - device=self.device,
355 )
350 - self.w2_lora_b_stacked = torch.zeros(
351 - (
352 - max_loras,
353 - self.base_layer.local_num_experts,
354 - self.base_layer.hidden_size
355 - if not self.fully_sharded
356 - else divide(self.base_layer.hidden_size, self.tp_size),
357 - lora_config.max_lora_rank,
358 - ),
359 - dtype=lora_config.lora_dtype,
360 - device=self.device,
356 +
357 + def create_lora_weights(
358 + self,
359 + max_loras: int,
360 + lora_config: LoRAConfig,
361 + model_config: PretrainedConfig | None = None,
362 + ) -> None:
363 + """Initializes lora matrices."""
364 + self.max_loras = lora_config.max_loras
365 + self.fully_sharded = lora_config.fully_sharded_loras
366 +
367 + self.adapter_enabled = torch.tensor(
368 + [0] * (max_loras + 1), dtype=torch.int, device=self.device
369 )
370
371 + self._create_lora_a_weights(max_loras, lora_config)
372 + self._create_lora_b_weights(max_loras, lora_config)
373 # They will be used by 'LoRALayerWeights.create_dummy_lora_weights'
374 # to create a dummy LoRA weights.
375 + # TODO Optimize this section
376 self.lora_a_stacked = []
377 self.lora_b_stacked = []
378 for lora_id in range(max_loras):
@@ -370,36 +381,43 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
381 self.lora_a_stacked.append(
382 self.w13_lora_a_stacked[0][lora_id][experts_id]
383 )
373 - self.lora_a_stacked.append(self.w2_lora_a_stacked[lora_id][experts_id])
384 self.lora_a_stacked.append(
375 - self.w13_lora_a_stacked[1][lora_id][experts_id]
385 + self.w2_lora_a_stacked[0][lora_id][experts_id]
386 )
387
388 self.lora_b_stacked.append(
389 self.w13_lora_b_stacked[0][lora_id][experts_id]
390 )
381 - self.lora_b_stacked.append(self.w2_lora_b_stacked[lora_id][experts_id])
391 + self.lora_b_stacked.append(
392 + self.w2_lora_b_stacked[0][lora_id][experts_id]
393 + )
394 +
395 + self.lora_a_stacked.append(
396 + self.w13_lora_a_stacked[1][lora_id][experts_id]
397 + )
398 self.lora_b_stacked.append(
399 self.w13_lora_b_stacked[1][lora_id][experts_id]
400 )
401
402 def reset_lora(self, index: int):
403 """Resets the lora weights at index back to 0."""
388 - for pos in range(self.w13_slices):
404 + for pos in range(self._w13_slices):
405 self.w13_lora_a_stacked[pos][index] = 0
406 self.w13_lora_b_stacked[pos][index] = 0
407
392 - self.w2_lora_a_stacked[index] = 0
393 - self.w2_lora_b_stacked[index] = 0
408 + self.w2_lora_a_stacked[0][index] = 0
409 + self.w2_lora_b_stacked[0][index] = 0
410 self.adapter_enabled[index] = 0
411
412 def set_lora(
413 self,
414 index: int,
399 - lora_a: torch.Tensor,
400 - lora_b: torch.Tensor,
415 + lora_a: torch.Tensor | list[torch.Tensor],
416 + lora_b: torch.Tensor | list[torch.Tensor],
417 ):
418 """Overwrites lora tensors at index."""
419 + assert isinstance(lora_a, list)
420 + assert isinstance(lora_b, list)
421 self.reset_lora(index)
422 self.adapter_enabled[index] = 1
423 for eid in range(len(lora_a) // 3):
@@ -432,7 +450,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
450 w1_lora_a = w1_lora_a[w13_start_idx:w13_end_idx, :]
451 w3_lora_a = w3_lora_a[w13_start_idx:w13_end_idx, :]
452
435 - w2_shard_size = self.w2_lora_b_stacked[index, eid].shape[0]
453 + w2_shard_size = self.w2_lora_b_stacked[0][index, eid].shape[0]
454 w2_start_idx = self.tp_rank * w2_shard_size
455 w2_end_idx = (self.tp_rank + 1) * w2_shard_size
456 w2_lora_b = w2_lora_b[w2_start_idx:w2_end_idx, :]
@@ -454,14 +472,32 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
472 index, eid, : w3_lora_b.shape[0], : w3_lora_b.shape[1]
473 ].copy_(w3_lora_b, non_blocking=True)
474
457 - self.w2_lora_a_stacked[
475 + self.w2_lora_a_stacked[0][
476 index, eid, : w2_lora_a.shape[0], : w2_lora_a.shape[1]
477 ].copy_(w2_lora_a, non_blocking=True)
478
461 - self.w2_lora_b_stacked[
479 + self.w2_lora_b_stacked[0][
480 index, eid, : w2_lora_b.shape[0], : w2_lora_b.shape[1]
481 ].copy_(w2_lora_b, non_blocking=True)
482
483 + def forward(self, *args, **kwargs):
484 + return self.base_layer.forward(*args, **kwargs)
485 +
486 + def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs):
487 + return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs)
488 +
489 + @property
490 + def _shared_experts(self):
491 + return self.base_layer._shared_experts
492 +
493 + @property
494 + def quant_method(self):
495 + return self.base_layer.quant_method
496 +
497 + @property
498 + def is_internal_router(self) -> bool:
499 + return self.base_layer.is_internal_router
500 +
501 @classmethod
502 def can_replace_layer(
503 cls,
@@ -472,22 +508,209 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
508 ) -> bool:
509 """Returns True if the layer can be replaced by this LoRA layer."""
510 # return type(source_layer) is FusedMoE
475 - return isinstance(source_layer, FusedMoE)
511
477 - def forward(self, *args, **kwargs):
478 - return self.base_layer.forward(*args, **kwargs)
512 + return type(source_layer) is FusedMoE and len(packed_modules_list) == 2
513
480 - def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs):
481 - return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs)
514 +
515 +class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
516 + def __init__(self, base_layer):
517 + super().__init__(base_layer)
518 + self._w13_slices = 1
519 +
520 + def _create_lora_b_weights(self, max_loras, lora_config):
521 + self.w13_lora_b_stacked: tuple[torch.Tensor] = tuple(
522 + torch.zeros(
523 + (
524 + max_loras,
525 + self.base_layer.local_num_experts,
526 + self.base_layer.intermediate_size_per_partition * 2,
527 + lora_config.max_lora_rank,
528 + ),
529 + dtype=lora_config.lora_dtype,
530 + device=self.device,
531 + )
532 + for _ in range(self._w13_slices)
533 + )
534 + self.w2_lora_b_stacked: tuple[torch.Tensor] = (
535 + torch.zeros(
536 + (
537 + max_loras,
538 + self.base_layer.local_num_experts,
539 + self.base_layer.hidden_size
540 + if not self.fully_sharded
541 + else divide(self.base_layer.hidden_size, self.tp_size),
542 + lora_config.max_lora_rank,
543 + ),
544 + dtype=lora_config.lora_dtype,
545 + device=self.device,
546 + ),
547 + )
548 +
549 + def create_lora_weights(
550 + self,
551 + max_loras: int,
552 + lora_config: LoRAConfig,
553 + model_config: PretrainedConfig | None = None,
554 + ) -> None:
555 + """Initializes lora matrices."""
556 + self.max_loras = lora_config.max_loras
557 + self.fully_sharded = lora_config.fully_sharded_loras
558 +
559 + self.adapter_enabled = torch.tensor(
560 + [0] * (max_loras + 1), dtype=torch.int, device=self.device
561 + )
562 +
563 + self._create_lora_a_weights(max_loras, lora_config)
564 + self._create_lora_b_weights(max_loras, lora_config)
565 +
566 + def _slice_w13_a(self, w13_lora_a: torch.Tensor) -> torch.Tensor:
567 + if self.tp_size == 1 or not self.fully_sharded:
568 + return w13_lora_a
569 +
570 + # w13_lora_a shape (num_experts,rank,input_size)
571 + current_lora_rank = w13_lora_a.shape[1]
572 + assert current_lora_rank % self.tp_size == 0
573 +
574 + sliced_rank = current_lora_rank // self.tp_size
575 + start_idx = self.tp_rank * sliced_rank
576 + end_idx = (self.tp_rank + 1) * sliced_rank
577 + return w13_lora_a[:, start_idx:end_idx, :]
578 +
579 + def _slice_w13_b(self, w13_lora_b: torch.Tensor, is_interleave: bool = True):
580 + if self.tp_size == 1:
581 + return w13_lora_b
582 +
583 + # w13_lora_b shape (num_experts,output_size,rank)
584 + shard_size = self.base_layer.intermediate_size_per_partition
585 + start_idx = self.tp_rank * shard_size
586 + end_idx = (self.tp_rank + 1) * shard_size
587 + if is_interleave:
588 + # For models like GPT-OSS, the weights of w1 (gate_proj) and w3 (up_proj)
589 + # in the interleaved order, and corresponding LoRA need to be processed.
590 + w1_lora_b = w13_lora_b[:, ::2, :]
591 + w3_lora_b = w13_lora_b[:, 1::2, :]
592 + sliced_w1_lora_b = w1_lora_b[:, start_idx:end_idx, :]
593 + sliced_w3_lora_b = w3_lora_b[:, start_idx:end_idx, :]
594 +
595 + return torch.stack([sliced_w1_lora_b, sliced_w3_lora_b], dim=2).flatten(
596 + 1, 2
597 + )
598 + else:
599 + slice_size = w13_lora_b.shape[1] // 2
600 + w1_lora_b = w13_lora_b[:, :slice_size, :]
601 + w3_lora_b = w13_lora_b[:, slice_size:, :]
602 + sliced_w1_lora_b = w1_lora_b[:, start_idx:end_idx, :]
603 + sliced_w3_lora_b = w3_lora_b[:, start_idx:end_idx, :]
604 +
605 + return torch.cat([sliced_w1_lora_b, sliced_w3_lora_b], dim=1)
606 +
607 + def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor:
608 + if self.tp_size == 1:
609 + return w2_lora_a
610 + # w2_lora_a shape (num_experts,rank,input_size)
611 + shard_size = self.base_layer.intermediate_size_per_partition
612 + start_idx = self.tp_rank * shard_size
613 + end_idx = (self.tp_rank + 1) * shard_size
614 +
615 + return w2_lora_a[:, :, start_idx:end_idx]
616 +
617 + def _slice_w2_b(self, w2_lora_b: torch.Tensor) -> torch.Tensor:
618 + if self.tp_size == 1 or not self.fully_sharded:
619 + return w2_lora_b
620 + # Based on S-LoRA, we slice W2 B along the hidden_size dim.
621 + # w2_lora_b shape (num_experts,output_size,rank)
622 + current_lora_size = w2_lora_b.shape[1]
623 +
624 + sliced_size = current_lora_size // self.tp_size
625 + start_idx = self.tp_rank * sliced_size
626 + end_idx = (self.tp_rank + 1) * sliced_size
627 + return w2_lora_b[:, start_idx:end_idx, :]
628 +
629 + def set_lora(
630 + self,
631 + index: int,
632 + lora_a: torch.Tensor | list[torch.Tensor],
633 + lora_b: torch.Tensor | list[torch.Tensor],
634 + ):
635 + """Overwrites lora tensors at index."""
636 + # Make mypy happy
637 + assert isinstance(lora_a, list)
638 + assert isinstance(lora_b, list)
639 + assert len(lora_a) == len(lora_b) == 2
640 +
641 + self.reset_lora(index)
642 + self.adapter_enabled[index] = 1
643 +
644 + num_experts = self.w13_lora_a_stacked[0].shape[1]
645 + w13_lora_a, w2_lora_a = lora_a
646 + w13_lora_b, w2_lora_b = lora_b
647 +
648 + # (num_experts,rank,input_size)
649 + w13_lora_a = w13_lora_a.reshape(num_experts, -1, w13_lora_a.shape[-1])
650 + w2_lora_a = w2_lora_a.reshape(num_experts, -1, w2_lora_a.shape[-1])
651 + # (output_size,num_experts,rank)
652 + w13_lora_b = w13_lora_b.reshape(w13_lora_b.shape[0], num_experts, -1)
653 + w2_lora_b = w2_lora_b.reshape(w2_lora_b.shape[0], num_experts, -1)
654 + # (num_experts,output_size,rank)
655 + w13_lora_b = w13_lora_b.permute(1, 0, 2)
656 + w2_lora_b = w2_lora_b.permute(1, 0, 2)
657 +
658 + sliced_w13_lora_a = self._slice_w13_a(w13_lora_a)
659 + sliced_w13_lora_b = self._slice_w13_b(w13_lora_b, is_interleave=True)
660 +
661 + sliced_w2_lora_a = self._slice_w2_a(w2_lora_a)
662 + sliced_w2_lora_b = self._slice_w2_b(w2_lora_b)
663 +
664 + self.w13_lora_a_stacked[0][
665 + index, :, : sliced_w13_lora_a.shape[1], : sliced_w13_lora_a.shape[2]
666 + ].copy_(sliced_w13_lora_a, non_blocking=True)
667 + self.w2_lora_a_stacked[0][
668 + index, :, : sliced_w2_lora_a.shape[1], : sliced_w2_lora_a.shape[2]
669 + ].copy_(sliced_w2_lora_a, non_blocking=True)
670 +
671 + self.w13_lora_b_stacked[0][
672 + index, :, : sliced_w13_lora_b.shape[1], : sliced_w13_lora_b.shape[2]
673 + ].copy_(sliced_w13_lora_b, non_blocking=True)
674 + self.w2_lora_b_stacked[0][
675 + index, :, : sliced_w2_lora_b.shape[1], : sliced_w2_lora_b.shape[2]
676 + ].copy_(sliced_w2_lora_b, non_blocking=True)
677
678 @property
484 - def _shared_experts(self):
485 - return self.base_layer._shared_experts
679 + def w13_input_size(self):
680 + """
681 + Full size
682 + """
683 + return self.w13_lora_a_stacked[0].shape[-1]
684
685 @property
488 - def quant_method(self):
489 - return self.base_layer.quant_method
686 + def w13_output_size(self):
687 + """
688 + Full size
689 + """
690 + return self.w13_lora_b_stacked[0].shape[-2] * self.tp_size
691
692 @property
492 - def is_internal_router(self) -> bool:
493 - return self.base_layer.is_internal_router
693 + def w2_input_size(self):
694 + """
695 + Full size
696 + """
697 + return self.w2_lora_a_stacked[0].shape[-1] * self.tp_size
698 +
699 + @property
700 + def w2_output_size(self):
701 + """
702 + Full size
703 + """
704 + return self.w2_lora_a_stacked[0].shape[-2]
705 +
706 + @classmethod
707 + def can_replace_layer(
708 + cls,
709 + source_layer: nn.Module,
710 + lora_config: LoRAConfig,
711 + packed_modules_list: list,
712 + model_config: PretrainedConfig | None,
713 + ) -> bool:
714 + """Returns True if the layer can be replaced by this LoRA layer."""
715 +
716 + return type(source_layer) is FusedMoE and len(packed_modules_list) == 1
vllm/lora/layers/logits_processor.py
+4 -2
@@ -128,9 +128,11 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
128 def set_lora(
129 self,
130 index: int,
131 - lora_a: torch.Tensor,
132 - lora_b: torch.Tensor,
131 + lora_a: torch.Tensor | list[torch.Tensor],
132 + lora_b: torch.Tensor | list[torch.Tensor],
133 ):
134 + assert isinstance(lora_a, torch.Tensor)
135 + assert isinstance(lora_b, torch.Tensor)
136 self.reset_lora(index)
137 self.lora_a_stacked[index, 0, : lora_a.shape[0], : lora_a.shape[1]].copy_(
138 lora_a, non_blocking=True
vllm/lora/layers/vocal_parallel_embedding.py
+5 -2
@@ -77,12 +77,15 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
77 def set_lora(
78 self,
79 index: int,
80 - lora_a: torch.Tensor,
81 - lora_b: torch.Tensor,
80 + lora_a: torch.Tensor | list[torch.Tensor],
81 + lora_b: torch.Tensor | list[torch.Tensor],
82 ):
83 + assert isinstance(lora_a, torch.Tensor)
84 + assert isinstance(lora_b, torch.Tensor)
85 self.reset_lora(index)
86 # NOTE self.lora_a_stacked is row-major, and lora_a is col-major,
87 # so we need transpose here
88 +
89 self.lora_a_stacked[index, : lora_a.shape[1], : lora_a.shape[0]].copy_(
90 lora_a.T, non_blocking=True
91 )
vllm/lora/models.py
+75 -24
@@ -22,11 +22,13 @@ from vllm.lora.utils import (
22 from_layer_logits_processor,
23 get_supported_lora_modules,
24 is_base_embeddding_weights,
25 + is_moe_model,
26 is_regex_target_modules,
27 parse_fine_tuned_lora_name,
28 process_packed_modules_mapping,
29 replace_submodule,
30 )
31 +from vllm.model_executor.layers.fused_moe import FusedMoE
32 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
33 from vllm.model_executor.models import SupportsLoRA, supports_multimodal
34 from vllm.model_executor.models.interfaces import is_pooling_model
@@ -356,7 +358,11 @@ class LoRAModelManager:
358 self.modules: dict[str, BaseLayerWithLoRA] = {}
359 # Dict instead of a set for compatibility with LRUCache.
360 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 + )
364 self._create_lora_modules()
365 +
366 self.model.lora_manager = self
367
368 def __len__(self) -> int:
@@ -400,22 +406,36 @@ class LoRAModelManager:
406 self.lora_index_to_id[index] = lora_model.id
407 for module_name, module in self.modules.items():
408 module_lora = self._get_lora_layer_weights(lora_model, module_name)
403 - if module_lora:
404 - # Note (gnovack) - If MOE lora weights are not split into
405 - # num_experts chunks, we split them here
406 - if isinstance(module, FusedMoEWithLoRA) and torch.is_tensor(
407 - module_lora.lora_a
408 - ):
409 - # Handle FSDP file format where experts.base_layer is the
410 - # gate_up_proj and experts is the down_proj
411 - gate_up_proj_lora = self._get_lora_layer_weights(
412 - lora_model, module_name + ".base_layer"
413 - )
414 -
415 - assert gate_up_proj_lora is not None
416 - assert module_lora is not None
417 -
418 - down_proj_lora = module_lora
409 + if not module_lora:
410 + module.reset_lora(index)
411 + continue
412 + # Note (gnovack) - If MOE lora weights are not split into
413 + # num_experts chunks, we split them here
414 + if isinstance(module, FusedMoEWithLoRA) and torch.is_tensor(
415 + module_lora.lora_a
416 + ):
417 + # Handle PEFT file format where experts.base_layer is the
418 + # gate_up_proj and experts is the down_proj
419 + gate_up_proj_lora = self._get_lora_layer_weights(
420 + lora_model, module_name + ".base_layer"
421 + )
422 + down_proj_lora = module_lora
423 + # FIXME Edge case where LoRA is not added to gate_up_proj
424 + # or down_proj
425 + assert gate_up_proj_lora is not None
426 + assert down_proj_lora is not None
427 + if self._is_3d_moe_model:
428 + module_lora.lora_a = [
429 + gate_up_proj_lora.lora_a,
430 + down_proj_lora.lora_a,
431 + ]
432 + module_lora.lora_b = [
433 + gate_up_proj_lora.lora_b,
434 + down_proj_lora.lora_b,
435 + ]
436 + else:
437 + # Some 3D MoE models haven't added the `is_3d_moe_weight`
438 + # attribute yet, so fallback here
439 num_experts = module_lora.lora_a.shape[0] // module_lora.rank
440
441 gate_proj_a = gate_up_proj_lora.lora_a.chunk(num_experts, dim=0)
@@ -444,14 +464,12 @@ class LoRAModelManager:
464
465 module_lora.lora_a = lora_a
466 module_lora.lora_b = lora_b
467 + module.set_lora(
468 + index,
469 + module_lora.lora_a,
470 + module_lora.lora_b,
471 + )
472
448 - module.set_lora(
449 - index,
450 - module_lora.lora_a,
451 - module_lora.lora_b,
452 - )
453 - else:
454 - module.reset_lora(index)
473 return True
474
475 def _deactivate_adapter(self, lora_id: int):
@@ -512,6 +530,13 @@ class LoRAModelManager:
530 continue
531 parts = module_name.split(".")[-1]
532 packed_moduled_lst = self.packed_modules_mapping.get(parts, [])
533 + if isinstance(module, FusedMoE):
534 + # packed_moduled_lst is used here to just determine whether to
535 + # instantiate FusedMoE3DWithLoRA or FusedMoEWithLoRA, and the
536 + # difference between these two LoRA layers is whether the
537 + # LoRA weights of w1 and w3 have already been fused on disk.
538 +
539 + packed_moduled_lst = ["w13"] if self._is_3d_moe_model else ["w1", "w3"]
540 new_module = replace_submodule(
541 self.model,
542 module_name,
@@ -560,6 +585,7 @@ class LoRAModelManager:
585 self._register_packed_modules(module_name)
586 # All lora layers share the same punica_wrapper based on reference.
587 new_module.set_mapping(self.punica_wrapper)
588 + pass
589
590 def register_module(self, module_name: str, module: "BaseLayerWithLoRA"):
591 assert isinstance(module, BaseLayerWithLoRA), (
@@ -605,6 +631,30 @@ class LoRAModelManager:
631 module.lora_a_stacked[0].dtype,
632 "cpu",
633 )
634 + model.loras[module_name] = lora
635 + elif module.__class__.__name__ == "FusedMoE3DWithLoRA":
636 + # Case for 3D moe model
637 + # w2
638 + lora = LoRALayerWeights.create_dummy_lora_weights(
639 + module_name,
640 + module.w2_input_size,
641 + module.w2_output_size,
642 + rank * module.w2_lora_a_stacked[0].shape[1], # rank*num_experts
643 + module.w2_lora_a_stacked[0].dtype,
644 + "cpu",
645 + )
646 + model.loras[module_name] = lora
647 + # w13
648 + lora = LoRALayerWeights.create_dummy_lora_weights(
649 + module_name,
650 + module.w13_input_size,
651 + module.w13_output_size,
652 + rank
653 + * module.w13_lora_a_stacked[0].shape[1], # rank*num_experts
654 + module.w13_lora_a_stacked[0].dtype,
655 + "cpu",
656 + )
657 + model.loras[module_name + ".base_layer"] = lora
658 else:
659 lora = LoRALayerWeights.create_dummy_lora_weights(
660 module_name,
@@ -614,6 +664,7 @@ class LoRAModelManager:
664 module.lora_a_stacked[0].dtype,
665 "cpu",
666 )
667 + model.loras[module_name] = lora
668 else:
669 parts = module_name.split(".")
670 replacements = self.packed_modules_mapping[parts[-1]]
@@ -629,7 +680,7 @@ class LoRAModelManager:
680 )
681 subloras.append(lora)
682 lora = PackedLoRALayerWeights.pack(subloras)
632 - model.loras[module_name] = lora
683 + model.loras[module_name] = lora
684 return model
685
686 def _match_target_modules(self, module_name: str):
vllm/lora/utils.py
+8 -4
@@ -23,6 +23,7 @@ from vllm.lora.layers import (
23 BaseLayerWithLoRA,
24 ColumnParallelLinearWithLoRA,
25 ColumnParallelLinearWithShardedLoRA,
26 + FusedMoE3DWithLoRA,
27 FusedMoEWithLoRA,
28 LogitsProcessorWithLoRA,
29 MergedColumnParallelLinearWithLoRA,
@@ -62,6 +63,7 @@ _all_lora_classes: set[type[BaseLayerWithLoRA]] = {
63 MergedQKVParallelLinearWithShardedLoRA,
64 RowParallelLinearWithShardedLoRA,
65 FusedMoEWithLoRA,
66 + FusedMoE3DWithLoRA,
67 }
68
69
@@ -288,10 +290,12 @@ def process_packed_modules_mapping(model: nn.Module) -> dict[str, list[str]]:
290 # the expert indices are expanded based on the configured number
291 # of routed experts.
292 packed_modules_mapping = get_packed_modules_mapping(model)
291 -
292 - packed_modules_mapping["experts"] = [
293 - weight_name.rstrip(".") for _, weight_name, _, _ in moe_packed_mapping
294 - ]
293 + if not hasattr(model, "is_3d_moe_weight"):
294 + # 3D MoE LoRA does not need `packed_modules_mapping`
295 + packed_modules_mapping["experts"] = [
296 + weight_name.rstrip(".")
297 + for _, weight_name, _, _ in moe_packed_mapping
298 + ]
299
300 return packed_modules_mapping
301 else:
vllm/model_executor/models/gpt_oss.py
+1
@@ -656,6 +656,7 @@ class GptOssModel(nn.Module):
656
657
658 class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
659 + is_3d_moe_weight: bool = True
660 packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
661
662 hf_to_vllm_mapper = WeightsMapper(