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 3, 01:45 UTC
[LoRA] Reduce memory of 2D weights when EP is set (#42737)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
May 22, 2026 at 21:41 UTC
15f7cd33dc8bd4d2270b70ba49d511827d2413ff
4 files changed
+236
-8
tests/lora/test_moe_lora_ep_load.py
new
+90
@@ -0,0 +1,90 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+End-to-end correctness test for 2D MoE LoRA expert-parallel
5
+load-time slicing
6
+"""
7
+
8
+import pytest
9
+import torch
10
+
11
+from vllm.lora.lora_model import LoRAModel, MoEEPLoadSpec
12
+from vllm.lora.peft_helper import PEFTHelper
13
+
14
+NUM_LAYERS = 48
15
+GLOBAL_NUM_EXPERTS = 128
16
+LOCAL_NUM_EXPERTS = 64 # ep_size = 2
17
+EXPERT_PROJECTIONS = ("down_proj", "gate_proj", "up_proj")
18
+NON_EXPERT_MODULES = ("q_proj", "k_proj", "v_proj", "o_proj", "gate")
19
+
20
+
21
+def _expected_lora_modules() -> set[str]:
22
+ """Replicate the set ``WorkerLoRAManager._load_adapter`` would build
23
+ from this model's ``packed_modules_mapping``."""
24
+ expected: set[str] = set(NON_EXPERT_MODULES)
25
+ for expert in range(GLOBAL_NUM_EXPERTS):
26
+ for proj in EXPERT_PROJECTIONS:
27
+ expected.add(f"experts.{expert}.{proj}")
28
+ return expected
29
+
30
+
31
+def _load(lora_dir, peft_helper, *, moe_ep_spec, lora_id):
32
+ return LoRAModel.from_local_checkpoint(
33
+ lora_dir,
34
+ _expected_lora_modules(),
35
+ peft_helper=peft_helper,
36
+ lora_model_id=lora_id,
37
+ device="cpu",
38
+ moe_ep_spec=moe_ep_spec,
39
+ )
40
+
41
+
42
+@pytest.mark.parametrize("ep_rank", [0, 1])
43
+def test_moe_lora_ep2_real_qwen3moe(qwen3moe_lora_files, ep_rank):
44
+ """ep_size=2 against the real Qwen3-MoE adapter: each rank's loaded
45
+ LoRA has the right size, the right expert membership, and the
46
+ right tensor values."""
47
+ peft_helper = PEFTHelper.from_local_dir(
48
+ qwen3moe_lora_files, max_position_embeddings=4096
49
+ )
50
+
51
+ # Baseline: no spec → loads every expert × projection × layer plus
52
+ # all non-expert LoRA modules.
53
+ ground_truth = _load(qwen3moe_lora_files, peft_helper, moe_ep_spec=None, lora_id=1)
54
+ expected_baseline = (
55
+ GLOBAL_NUM_EXPERTS * len(EXPERT_PROJECTIONS) * NUM_LAYERS
56
+ + len(NON_EXPERT_MODULES) * NUM_LAYERS
57
+ )
58
+ assert len(ground_truth.loras) == expected_baseline
59
+
60
+ # Sliced load: only this rank's experts; non-expert LoRA is untouched.
61
+ spec = MoEEPLoadSpec(
62
+ ep_rank=ep_rank,
63
+ local_num_experts=LOCAL_NUM_EXPERTS,
64
+ global_num_experts=GLOBAL_NUM_EXPERTS,
65
+ )
66
+ sliced = _load(
67
+ qwen3moe_lora_files,
68
+ peft_helper,
69
+ moe_ep_spec=spec,
70
+ lora_id=100 + ep_rank,
71
+ )
72
+
73
+ expected_sliced = (
74
+ LOCAL_NUM_EXPERTS * len(EXPERT_PROJECTIONS) * NUM_LAYERS
75
+ + len(NON_EXPERT_MODULES) * NUM_LAYERS
76
+ )
77
+ assert len(sliced.loras) == expected_sliced
78
+
79
+ expert_start = ep_rank * LOCAL_NUM_EXPERTS
80
+ expert_end = expert_start + LOCAL_NUM_EXPERTS
81
+
82
+ for name, lora in sliced.loras.items():
83
+ gt = ground_truth.loras[name]
84
+ torch.testing.assert_close(lora.lora_a, gt.lora_a)
85
+ torch.testing.assert_close(lora.lora_b, gt.lora_b)
86
+ if ".experts." in name:
87
+ expert_idx = int(name.split(".experts.")[-1].split(".")[0])
88
+ assert expert_start <= expert_idx < expert_end, (
89
+ f"non-local expert {expert_idx} leaked: {name}"
90
+ )
vllm/lora/lora_model.py
+55
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import os
5
+from dataclasses import dataclass
6
7
import safetensors
8
import torch
@@ -21,6 +22,41 @@ from vllm.utils.platform_utils import is_pin_memory_available
22
logger = init_logger(__name__)
23
24
25
+@dataclass(frozen=True)
26
+class MoEEPLoadSpec:
27
+ """Per-expert-parallel slicing metadata for one FusedMoE LoRA module.
28
+
29
+ Threaded into the LoRA loader so per-expert weights from EP ranks
30
+ other than this one can be skipped before they ever hit CPU memory.
31
+ """
32
+
33
+ ep_rank: int
34
+ local_num_experts: int
35
+ global_num_experts: int
36
+
37
+
38
+_EXPERTS_SEPARATOR = ".experts."
39
+
40
+
41
+def _is_remote_expert_key(raw_name: str, spec: "MoEEPLoadSpec") -> bool:
42
+ """
43
+ Decide whether a checkpoint key belongs to a non-local expert.
44
+ """
45
+ pos = raw_name.find(_EXPERTS_SEPARATOR)
46
+ if pos < 0:
47
+ return False
48
+ idx_start = pos + len(_EXPERTS_SEPARATOR)
49
+ idx_end = raw_name.find(".", idx_start)
50
+ if idx_end < 0:
51
+ return False
52
+ idx_str = raw_name[idx_start:idx_end]
53
+ if not idx_str.isdigit():
54
+ return False
55
+ expert_idx = int(idx_str)
56
+ local_start = spec.ep_rank * spec.local_num_experts
57
+ return not (local_start <= expert_idx < local_start + spec.local_num_experts)
58
+
59
+
60
class LoRAModel:
61
"""A LoRA fine-tuned model."""
62
@@ -141,6 +177,7 @@ class LoRAModel:
177
weights_mapper: WeightsMapper | None = None,
178
tensorizer_config_dict: dict | None = None,
179
skip_prefixes: list[str] | None = None,
180
+ moe_ep_spec: MoEEPLoadSpec | None = None,
181
) -> "LoRAModel":
182
"""Create a LoRAModel from a local checkpoint.
183
@@ -156,6 +193,11 @@ class LoRAModel:
193
skip_prefixes: List of module name prefixes to skip during loading.
194
Models can define this to skip modules not used in inference
195
(e.g., MTP layers). Format: ["mtp."]
196
+ moe_ep_spec: When 2D FusedMoE LoRA modules are present with
197
+ expert parallelism enabled, the (ep_rank, local, global)
198
+ slicing metadata shared across all MoE layers. Non-local
199
+ expert weights are skipped at read time instead of being
200
+ loaded and discarded later.
201
202
Returns:
203
Loaded LoRA Model.
@@ -227,6 +269,10 @@ class LoRAModel:
269
# Load tensors if there are only expected modules.
270
check_unexpected_modules(f)
271
for module in f.keys(): # noqa
272
+ if moe_ep_spec is not None and _is_remote_expert_key(
273
+ module, moe_ep_spec
274
+ ):
275
+ continue
276
tensors[module] = f.get_tensor(module)
277
elif os.path.isfile(lora_bin_file_path) or os.path.isfile(lora_pt_file_path):
278
lora_file_path = (
@@ -236,6 +282,15 @@ class LoRAModel:
282
)
283
tensors = torch.load(lora_file_path, map_location=device, weights_only=True)
284
check_unexpected_modules(tensors)
285
+ if moe_ep_spec is not None:
286
+ # `.bin`/`.pt` adapters can't be lazy-loaded, but pruning
287
+ # the dict here still frees the non-local expert tensors
288
+ # before the dtype cast / pin_memory work that follows.
289
+ tensors = {
290
+ k: v
291
+ for k, v in tensors.items()
292
+ if not _is_remote_expert_key(k, moe_ep_spec)
293
+ }
294
else:
295
raise ValueError(f"{lora_dir} doesn't contain tensors")
296
vllm/lora/model_manager.py
+90
-8
@@ -18,7 +18,7 @@ from vllm.lora.layers import (
18
LoRAMapping,
19
LoRAMappingType,
20
)
21
-from vllm.lora.lora_model import LoRAModel
21
+from vllm.lora.lora_model import LoRAModel, MoEEPLoadSpec
22
from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
23
from vllm.lora.punica_wrapper import PunicaWrapperBase, get_punica_wrapper
24
from vllm.lora.utils import (
@@ -109,8 +109,8 @@ class LoRAModelManager:
109
# Dict instead of a set for compatibility with LRUCache.
110
self._last_mapping: LoRAMapping | None = None
111
is_moe = is_moe_model(self.model)
112
- # Whether the underlying model class declares 3D fused MoE weights.
113
- self._model_is_3d_moe = is_moe and self.model.is_3d_moe_weight
112
+ self._is_moe = is_moe
113
+
114
# When the engine is started with enable_mixed_moe_lora_format=True
115
# we force the universal 2D wrapper (FusedMoEWithLoRA) regardless of
116
# the model's 3D flag, so 2D and 3D adapters can coexist.
@@ -118,15 +118,22 @@ class LoRAModelManager:
118
is_moe and lora_config.enable_mixed_moe_lora_format
119
)
120
self._is_3d_moe_model = (
121
- self._model_is_3d_moe and not self._enable_mixed_moe_lora_format
121
+ self._is_moe
122
+ and self.model.is_3d_moe_weight
123
+ and not self._enable_mixed_moe_lora_format
124
)
125
self.packed_modules_mapping = process_packed_modules_mapping(
126
self.model, force_2d_moe=self._enable_mixed_moe_lora_format
127
)
128
self._is_non_gated_moe = is_moe and self.model.is_non_gated_moe
129
+ self._use_ep = bool(
130
+ vllm_config and vllm_config.parallel_config.enable_expert_parallel
131
+ )
132
self._init_punica_wrapper(max_num_batched_tokens, vllm_config)
133
self._create_lora_modules()
134
135
+ self.moe_ep_load_spec: MoEEPLoadSpec | None = self._build_moe_ep_load_spec()
136
+
137
self.model.lora_manager = self
138
139
def _init_punica_wrapper(
@@ -715,15 +722,21 @@ class LoRAModelManager:
722
723
def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None:
724
for module_name, new_module_names in self.packed_modules.items():
725
+ # For 2D FusedMoE modules with EP, narrow the per-expert
726
+ # sub-module list to this rank's owned experts so pack_moe
727
+ # produces a tensor sized to local_num_experts directly.
728
+ packed_module_names = new_module_names
729
+ if module_name.endswith(".experts"):
730
+ new_module_names = self._restrict_to_local_experts(
731
+ module_name, new_module_names
732
+ )
733
replacement_loras: list[LoRALayerWeights | None] = []
719
- replaced_module: set[str] = set()
734
has_replacement = False
735
for r in new_module_names:
736
lora = self._get_lora_layer_weights(lora_model, r)
737
replacement_loras.append(lora)
738
if lora:
739
has_replacement = True
726
- replaced_module.add(r)
740
if not has_replacement:
741
continue
742
for i in range(len(replacement_loras)):
@@ -749,8 +762,11 @@ class LoRAModelManager:
762
lora_model.loras[module_name] = PackedLoRALayerWeights.pack(
763
replacement_loras
764
)
752
- # Remove the modules that have been replaced.
753
- for module in replaced_module:
765
+ # Drop every candidate sub-module, including non-local expert
766
+ # entries that were loaded but did not contribute to the
767
+ # packed result. Without this they would keep extra CPU
768
+ # memory alive after pack_moe.
769
+ for module in packed_module_names:
770
lora_model.loras.pop(module, None)
771
772
for lora in lora_model.loras.values():
@@ -998,6 +1014,13 @@ class LoRAModelManager:
1014
is active each rank only owns local_num_experts; without this slice
1015
the CPU LoRAModel keeps the full global weight and set_lora has to
1016
re-slice on every activation.
1017
+
1018
+ With the load-time / pack-time slicing in
1019
+ ``_restrict_to_local_experts``, the stacked tensors already match
1020
+ ``local_num_experts`` and the inner branch becomes a no-op. The
1021
+ guard remains so checkpoints that bypassed the pre-slicing (e.g.
1022
+ ``.bin``/``.pt`` adapters with weights mappers we don't recognize)
1023
+ still get sliced here.
1024
"""
1025
if not module.base_layer.use_ep:
1026
return
@@ -1022,6 +1045,65 @@ class LoRAModelManager:
1045
module_lora.lora_a = new_lora_a
1046
module_lora.lora_b = new_lora_b
1047
1048
+ def _restrict_to_local_experts(
1049
+ self, module_name: str, new_module_names: list[str]
1050
+ ) -> list[str]:
1051
+ """Narrow a flat expert-major sub-module list to this rank's experts.
1052
+
1053
+ ``new_module_names`` is produced by
1054
+ ``FusedMoE.make_expert_params_mapping`` and is ordered
1055
+ ``[e=0,w1, e=0,w2, e=0,w3, e=1,w1, ...]`` (non-gated MoE has 2
1056
+ entries per expert instead of 3). When the module is a 2D
1057
+ ``FusedMoEWithLoRA`` with EP enabled, we slice the list to the
1058
+ contiguous block of experts owned by this rank so the downstream
1059
+ ``pack_moe`` call only consumes local weights and produces a
1060
+ tensor sized to ``local_num_experts`` directly.
1061
+
1062
+ Returns the original list unchanged for non-MoE modules, the 3D
1063
+ MoE path (handled separately by ``_stack_moe_lora_weights``),
1064
+ modules without EP, or layouts we cannot cleanly partition.
1065
+ """
1066
+ module = self.modules.get(module_name)
1067
+ if not isinstance(module, FusedMoEWithLoRA):
1068
+ return new_module_names
1069
+ if isinstance(module, FusedMoE3DWithLoRA):
1070
+ return new_module_names
1071
+ if not getattr(module.base_layer, "use_ep", False):
1072
+ return new_module_names
1073
+ global_num_experts = module.base_layer.global_num_experts
1074
+ local_num_experts = module.base_layer.local_num_experts
1075
+ ep_rank = module.base_layer.ep_rank
1076
+ if global_num_experts <= 0 or len(new_module_names) % global_num_experts != 0:
1077
+ return new_module_names
1078
+ per_expert = len(new_module_names) // global_num_experts
1079
+ start = ep_rank * local_num_experts * per_expert
1080
+ end = start + local_num_experts * per_expert
1081
+ return new_module_names[start:end]
1082
+
1083
+ def _build_moe_ep_load_spec(self) -> MoEEPLoadSpec | None:
1084
+ """
1085
+ Per-rank slicing metadata for 2D FusedMoE LoRA modules.
1086
+ """
1087
+ if not self._use_ep or not self._is_moe:
1088
+ return None
1089
+ module = next(
1090
+ (
1091
+ m
1092
+ for m in self.modules.values()
1093
+ if isinstance(m, FusedMoEWithLoRA)
1094
+ and not isinstance(m, FusedMoE3DWithLoRA)
1095
+ ),
1096
+ None,
1097
+ )
1098
+ if module is None:
1099
+ return None
1100
+ base = module.base_layer
1101
+ return MoEEPLoadSpec(
1102
+ ep_rank=base.ep_rank,
1103
+ local_num_experts=base.local_num_experts,
1104
+ global_num_experts=base.global_num_experts,
1105
+ )
1106
+
1107
def _get_lora_layer_weights(
1108
self, lora_model: LoRAModel, module_name: str
1109
) -> LoRALayerWeights | None:
vllm/lora/worker_manager.py
+1
@@ -140,6 +140,7 @@ class WorkerLoRAManager:
140
tensorizer_config_dict=lora_request.tensorizer_config_dict,
141
weights_mapper=hf_to_vllm_mapper,
142
skip_prefixes=lora_skip_prefixes,
143
+ moe_ep_spec=self._adapter_manager.moe_ep_load_spec,
144
)
145
# Stamp the on-disk MoE layout onto the loaded model so the
146
# adapter manager can route 3D-format checkpoints through the