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 10, 17:45 UTC
pin lora_b moe weights on cpu (#31317)
Signed-off-by: gnovack <gnovack@amazon.com>
gnovack committed
Jan 5, 2026 at 12:15 UTC
c77a993cc297f23ae6f0c3ff1efb2d6d022bdfd2
2 files changed
+89
-65
vllm/lora/layers/fused_moe.py
-11
@@ -671,20 +671,9 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
671
self.reset_lora(index)
672
self.adapter_enabled[index] = 1
673
674
- num_experts = self.w13_lora_a_stacked[0].shape[1]
674
w13_lora_a, w2_lora_a = lora_a
675
w13_lora_b, w2_lora_b = lora_b
676
678
- # (num_experts,rank,input_size)
679
- w13_lora_a = w13_lora_a.reshape(num_experts, -1, w13_lora_a.shape[-1])
680
- w2_lora_a = w2_lora_a.reshape(num_experts, -1, w2_lora_a.shape[-1])
681
- # (output_size,rank,num_experts)
682
- w13_lora_b = w13_lora_b.reshape(w13_lora_b.shape[0], -1, num_experts)
683
- w2_lora_b = w2_lora_b.reshape(w2_lora_b.shape[0], -1, num_experts)
684
- # (num_experts,output_size,rank)
685
- w13_lora_b = w13_lora_b.permute(2, 0, 1)
686
- w2_lora_b = w2_lora_b.permute(2, 0, 1)
687
-
677
sliced_w13_lora_a = self._slice_w13_a(w13_lora_a)
678
sliced_w13_lora_b = self._slice_w13_b(w13_lora_b)
679
vllm/lora/model_manager.py
+89
-54
@@ -256,61 +256,7 @@ class LoRAModelManager:
256
if not module_lora:
257
module.reset_lora(index)
258
continue
259
- # Note (gnovack) - If MOE lora weights are not split into
260
- # num_experts chunks, we split them here
261
- if isinstance(module, FusedMoE3DWithLoRA) and torch.is_tensor(
262
- module_lora.lora_a
263
- ):
264
- # Handle PEFT file format where experts.base_layer is the
265
- # gate_up_proj and experts is the down_proj
266
- gate_up_proj_lora = self._get_lora_layer_weights(
267
- lora_model, module_name + ".base_layer"
268
- )
269
- down_proj_lora = module_lora
270
- # FIXME Edge case where LoRA is not added to gate_up_proj
271
- # or down_proj
272
- assert gate_up_proj_lora is not None
273
- assert down_proj_lora is not None
274
- if self._is_3d_moe_model:
275
- module_lora.lora_a = [
276
- gate_up_proj_lora.lora_a,
277
- down_proj_lora.lora_a,
278
- ]
279
- module_lora.lora_b = [
280
- gate_up_proj_lora.lora_b,
281
- down_proj_lora.lora_b,
282
- ]
283
- else:
284
- # Some 3D MoE models haven't added the `is_3d_moe_weight`
285
- # attribute yet, so fallback here
286
- num_experts = module_lora.lora_a.shape[0] // module_lora.rank
287
-
288
- gate_proj_a = gate_up_proj_lora.lora_a.chunk(num_experts, dim=0)
289
- up_proj_a = gate_up_proj_lora.lora_a.chunk(num_experts, dim=0)
290
-
291
- gate_proj_b = gate_up_proj_lora.lora_b[::2, ...].chunk(
292
- num_experts, dim=-1
293
- )
294
- up_proj_b = gate_up_proj_lora.lora_b[1::2, ...].chunk(
295
- num_experts, dim=-1
296
- )
297
-
298
- down_proj_a = down_proj_lora.lora_a.chunk(num_experts, dim=0)
299
- down_proj_b = down_proj_lora.lora_b.chunk(num_experts, dim=-1)
259
301
- lora_a = []
302
- lora_b = []
303
- for i in range(num_experts):
304
- lora_a.append(gate_proj_a[i])
305
- lora_a.append(down_proj_a[i])
306
- lora_a.append(up_proj_a[i])
307
-
308
- lora_b.append(gate_proj_b[i])
309
- lora_b.append(down_proj_b[i])
310
- lora_b.append(up_proj_b[i])
311
-
312
- module_lora.lora_a = lora_a
313
- module_lora.lora_b = lora_b
260
module.set_lora(
261
index,
262
module_lora.lora_a,
@@ -627,6 +573,10 @@ class LoRAModelManager:
573
for lora in lora_model.loras.values():
574
lora.optimize()
575
576
+ for module_name, module in self.modules.items():
577
+ if isinstance(module, FusedMoE3DWithLoRA):
578
+ self._stack_moe_lora_weights(lora_model, module, module_name)
579
+
580
first_lora: LoRALayerWeights = next(iter(lora_model.loras.values()))
581
assert first_lora.lora_a is not None
582
if isinstance(first_lora.lora_a, list):
@@ -653,6 +603,91 @@ class LoRAModelManager:
603
lora.lora_a = lora.lora_a.pin_memory()
604
lora.lora_b = lora.lora_b.pin_memory()
605
606
+ def _stack_moe_lora_weights(
607
+ self, lora_model: LoRAModel, module: FusedMoE3DWithLoRA, module_name: str
608
+ ):
609
+ module_lora = self._get_lora_layer_weights(lora_model, module_name)
610
+
611
+ # Note (gnovack) - If MOE lora weights are not split into
612
+ # num_experts chunks, we split them here
613
+ if module_lora and torch.is_tensor(module_lora.lora_a):
614
+ # Handle PEFT file format where experts.base_layer is the
615
+ # gate_up_proj and experts is the down_proj
616
+ gate_up_proj_lora = self._get_lora_layer_weights(
617
+ lora_model, module_name + ".base_layer"
618
+ )
619
+ down_proj_lora = module_lora
620
+ # FIXME Edge case where LoRA is not added to gate_up_proj
621
+ # or down_proj
622
+ assert gate_up_proj_lora is not None
623
+ assert down_proj_lora is not None
624
+ if self._is_3d_moe_model:
625
+ num_experts = module.w13_lora_a_stacked[0].shape[1]
626
+
627
+ # (num_experts,rank,input_size)
628
+ gate_up_proj_lora.lora_a = gate_up_proj_lora.lora_a.reshape(
629
+ num_experts, -1, gate_up_proj_lora.lora_a.shape[-1]
630
+ )
631
+ down_proj_lora.lora_a = down_proj_lora.lora_a.reshape(
632
+ num_experts, -1, down_proj_lora.lora_a.shape[-1]
633
+ )
634
+
635
+ # (output_size,num_experts,rank)
636
+ gate_up_proj_lora.lora_b = gate_up_proj_lora.lora_b.reshape(
637
+ gate_up_proj_lora.lora_b.shape[0], -1, num_experts
638
+ )
639
+ down_proj_lora.lora_b = down_proj_lora.lora_b.reshape(
640
+ down_proj_lora.lora_b.shape[0], -1, num_experts
641
+ )
642
+
643
+ # (num_experts,output_size,rank)
644
+ gate_up_proj_lora.lora_b = gate_up_proj_lora.lora_b.permute(
645
+ 2, 0, 1
646
+ ).contiguous()
647
+ down_proj_lora.lora_b = down_proj_lora.lora_b.permute(
648
+ 2, 0, 1
649
+ ).contiguous()
650
+
651
+ module_lora.lora_a = [
652
+ gate_up_proj_lora.lora_a,
653
+ down_proj_lora.lora_a,
654
+ ]
655
+ module_lora.lora_b = [
656
+ gate_up_proj_lora.lora_b,
657
+ down_proj_lora.lora_b,
658
+ ]
659
+ else:
660
+ # Some 3D MoE models haven't added the `is_3d_moe_weight`
661
+ # attribute yet, so fallback here
662
+ num_experts = module_lora.lora_a.shape[0] // module_lora.rank
663
+
664
+ gate_proj_a = gate_up_proj_lora.lora_a.chunk(num_experts, dim=0)
665
+ up_proj_a = gate_up_proj_lora.lora_a.chunk(num_experts, dim=0)
666
+
667
+ gate_proj_b = gate_up_proj_lora.lora_b[::2, ...].chunk(
668
+ num_experts, dim=-1
669
+ )
670
+ up_proj_b = gate_up_proj_lora.lora_b[1::2, ...].chunk(
671
+ num_experts, dim=-1
672
+ )
673
+
674
+ down_proj_a = down_proj_lora.lora_a.chunk(num_experts, dim=0)
675
+ down_proj_b = down_proj_lora.lora_b.chunk(num_experts, dim=-1)
676
+
677
+ lora_a = []
678
+ lora_b = []
679
+ for i in range(num_experts):
680
+ lora_a.append(gate_proj_a[i])
681
+ lora_a.append(down_proj_a[i])
682
+ lora_a.append(up_proj_a[i])
683
+
684
+ lora_b.append(gate_proj_b[i])
685
+ lora_b.append(down_proj_b[i])
686
+ lora_b.append(up_proj_b[i])
687
+
688
+ module_lora.lora_a = lora_a
689
+ module_lora.lora_b = lora_b
690
+
691
def _get_lora_layer_weights(
692
self, lora_model: LoRAModel, module_name: str
693
) -> LoRALayerWeights | None: