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 14, 05:30 UTC
Support expert parallel load balancing in Transformers backend (#26287)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Oct 6, 2025 at 12:20 UTC
0340f455539f8bd8c3da07d99e1796a4e4d8a5a9
2 files changed
+76
-27
vllm/model_executor/models/transformers.py
+13
-11
@@ -40,7 +40,7 @@ from vllm.config import (
40
)
41
from vllm.config.multimodal import BaseDummyOptions
42
from vllm.config.utils import getattr_iter
43
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
43
+from vllm.distributed import get_pp_group, get_tp_group
44
from vllm.distributed.utils import get_pp_indices
45
from vllm.logger import init_logger
46
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -506,9 +506,7 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
506
self.quant_config: Optional[QuantizationConfig] = vllm_config.quant_config
507
508
self.pp_group = get_pp_group()
509
- self.pp_size = self.pp_group.world_size
510
- self.pp_rank = self.pp_group.rank_in_group
511
- self.tp_size = get_tensor_model_parallel_world_size()
509
+ self.tp_group = get_tp_group()
510
511
# Weights to skip in `self.load_weights`
512
self.skip_prefixes: list[str] = []
@@ -576,7 +574,7 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
574
"""
575
Apply the model's pipeline parallelization plan.
576
"""
579
- if self.pp_size <= 1:
577
+ if self.pp_group.world_size <= 1:
578
return
579
580
if not self.model.supports_pp_plan:
@@ -613,7 +611,9 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
611
612
# Module list
613
start_layer, end_layer = get_pp_indices(
616
- self.text_config.num_hidden_layers, self.pp_rank, self.pp_size
614
+ self.text_config.num_hidden_layers,
615
+ self.pp_group.rank_in_group,
616
+ self.pp_group.world_size,
617
)
618
layers_name = pp_plan[module_list_idx]
619
layers = getattr(self.model, layers_name)
@@ -638,7 +638,7 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
638
"""
639
tp_plan = self.model.tp_plan
640
641
- if not tp_plan and self.tp_size > 1:
641
+ if not tp_plan and self.tp_group.world_size > 1:
642
tip = get_feature_request_tip(
643
self.model_config.model, self.model_config.trust_remote_code
644
)
@@ -687,7 +687,9 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
687
head_size = self.model_config.get_head_size()
688
num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config)
689
start, end = get_pp_indices(
690
- self.text_config.num_hidden_layers, self.pp_rank, self.pp_size
690
+ self.text_config.num_hidden_layers,
691
+ self.pp_group.rank_in_group,
692
+ self.pp_group.world_size,
693
)
694
695
attention_instances = {}
@@ -749,7 +751,7 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
751
intermediate_tensors: Optional[IntermediateTensors] = None,
752
inputs_embeds: Optional[torch.Tensor] = None,
753
) -> Union[torch.Tensor, IntermediateTensors]:
752
- if not get_pp_group().is_first_rank:
754
+ if not self.pp_group.is_first_rank:
755
assert intermediate_tensors is not None
756
input_ids = None
757
inputs_embeds = intermediate_tensors["hidden_states"]
@@ -773,7 +775,7 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
775
return_dict=False,
776
)[0][0, ...] # we remove batch dimension for now
777
776
- if not get_pp_group().is_last_rank:
778
+ if not self.pp_group.is_last_rank:
779
return IntermediateTensors({"hidden_states": hidden_states})
780
781
return hidden_states
@@ -811,7 +813,7 @@ class TransformersForCausalLM(TransformersBase):
813
if self.text_config.tie_word_embeddings:
814
self.skip_prefixes.append("lm_head.")
815
814
- if get_pp_group().is_last_rank:
816
+ if self.pp_group.is_last_rank:
817
self.unpadded_vocab_size = self.text_config.vocab_size
818
self.lm_head = ParallelLMHead(
819
self.text_config.vocab_size,
vllm/model_executor/models/transformers_moe.py
+63
-16
@@ -30,6 +30,7 @@ from vllm.model_executor.layers.fused_moe import FusedMoE
30
from vllm.platforms import current_platform
31
from vllm.utils import direct_register_custom_op
32
33
+from .interfaces import MixtureOfExperts
34
from .transformers import (
35
TransformersBase,
36
TransformersForCausalLM,
@@ -116,17 +117,41 @@ direct_register_custom_op(
117
)
118
119
119
-class TransformersMoEBase(TransformersBase):
120
+class TransformersMoEBase(TransformersBase, MixtureOfExperts):
121
def __init__(self, *, vllm_config, prefix=""):
122
self.check_version("4.57.0.dev0", "MoE models support")
123
+ self.ep_group = get_ep_group()
124
super().__init__(vllm_config=vllm_config, prefix=prefix)
125
124
- if self.parallel_config.enable_eplb:
125
- raise NotImplementedError(
126
- "Transformers backend does not support expert parallel load "
127
- "balancing yet."
126
+ def set_eplb_state(
127
+ self,
128
+ expert_load_view: torch.Tensor,
129
+ logical_to_physical_map: torch.Tensor,
130
+ logical_replica_count: torch.Tensor,
131
+ ):
132
+ for moe_layer_idx, mlp_layer in enumerate(self.mlp_layers):
133
+ mlp_layer.experts.set_eplb_state(
134
+ moe_layer_idx=moe_layer_idx,
135
+ expert_load_view=expert_load_view,
136
+ logical_to_physical_map=logical_to_physical_map,
137
+ logical_replica_count=logical_replica_count,
138
)
139
140
+ def update_physical_experts_metadata(
141
+ self,
142
+ num_physical_experts: int,
143
+ num_local_physical_experts: int,
144
+ ):
145
+ assert self.num_local_physical_experts == num_local_physical_experts
146
+ self.num_physical_experts = num_physical_experts
147
+ self.num_local_physical_experts = num_local_physical_experts
148
+ self.num_redundant_experts = num_physical_experts - self.num_logical_experts
149
+ for mlp in self.mlp_layers:
150
+ mlp.n_local_physical_experts = num_local_physical_experts
151
+ mlp.n_physical_experts = num_physical_experts
152
+ mlp.n_redundant_experts = self.num_redundant_experts
153
+ mlp.experts.update_expert_map()
154
+
155
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
156
"""
157
Params for weights, fp8 weight scales, fp8 activation scales
@@ -138,6 +163,8 @@ class TransformersMoEBase(TransformersBase):
163
("w1", "w2", "w3"), # Granite, Mixtral, Phi MoE style
164
("linear", "linear_1", "linear_v"), # Grok1 style
165
]
166
+ num_experts = self.model_config.get_num_experts()
167
+ num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts
168
expert_mapping = []
169
for gate_proj, down_proj, up_proj in ckpt_names:
170
expert_mapping.extend(
@@ -145,8 +172,8 @@ class TransformersMoEBase(TransformersBase):
172
ckpt_gate_proj_name=gate_proj,
173
ckpt_down_proj_name=down_proj,
174
ckpt_up_proj_name=up_proj,
148
- num_experts=self.model_config.get_num_experts(),
149
- num_redundant_experts=0, # TODO: enable EPLB
175
+ num_experts=num_experts,
176
+ num_redundant_experts=num_redundant_experts,
177
)
178
)
179
return expert_mapping
@@ -167,12 +194,15 @@ class TransformersMoEBase(TransformersBase):
194
195
# If there are shared experts, the results are
196
# reduced after mlp.forward() not inside FusedMoE
170
- num_experts_shared = getattr_iter(
197
+ num_shared_experts = getattr_iter(
198
text_config,
172
- ["num_experts_shared", "n_shared_experts", "moe_num_shared_experts"],
199
+ [
200
+ "n_shared_experts", # DeepSeek, Docs, GLM
201
+ "moe_num_shared_experts", # Aria, Ernie
202
+ ],
203
0,
204
)
175
- reduce_results = num_experts_shared == 0
205
+ reduce_results = num_shared_experts == 0
206
207
def add_all_reduce(mlp: nn.Module):
208
"""Adds an all-reduce to the output of `mlp.forward()`."""
@@ -207,13 +237,23 @@ class TransformersMoEBase(TransformersBase):
237
# Expert mapping for `AutoWeightsLoader`
238
expert_mapping = self.get_expert_mapping()
239
210
- # Configs
211
- parallel_config = self.parallel_config
212
- eplb_config = parallel_config.eplb_config
213
-
240
# Expert parallel load balancing kwargs
215
- enable_eplb = parallel_config.enable_eplb
216
- num_redundant_experts = eplb_config.num_redundant_experts
241
+ enable_eplb = self.parallel_config.enable_eplb
242
+ num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts
243
+
244
+ # MixtureOfExperts mixin settings
245
+ ep_size = self.ep_group.world_size
246
+
247
+ self.mlp_layers = [] # Used for MixtureOfExperts methods
248
+ self.expert_weights = []
249
+ self.num_moe_layers = 0
250
+ self.num_expert_groups = 1 if num_expert_group is None else num_expert_group
251
+ self.num_logical_experts = num_experts
252
+ self.num_physical_experts = num_experts + num_redundant_experts
253
+ self.num_local_physical_experts = self.num_physical_experts // ep_size
254
+ self.num_routed_experts = num_experts
255
+ self.num_shared_experts = num_shared_experts
256
+ self.num_redundant_experts = num_redundant_experts
257
258
# Recursively fuse MoE layers
259
def _recursive_replace(module: nn.Module, prefix: str):
@@ -235,6 +275,9 @@ class TransformersMoEBase(TransformersBase):
275
for mlp_param_name, _ in mlp.named_parameters():
276
if "shared_expert" in mlp_param_name:
277
reduce_results = False
278
+ # If the config does not specify num_shared_experts, but
279
+ # the model has shared experts, we assume there is one.
280
+ self.num_shared_experts = 1
281
break
282
# Replace experts module with FusedMoE
283
fused_experts = TransformersFusedMoE(
@@ -258,6 +301,10 @@ class TransformersMoEBase(TransformersBase):
301
)
302
mlp.experts = fused_experts
303
log_replacement(qual_name, experts, fused_experts)
304
+ # Update MixtureOfExperts mixin state
305
+ self.mlp_layers.append(mlp)
306
+ self.expert_weights.append(fused_experts.get_expert_weights())
307
+ self.num_moe_layers += 1
308
# If results are not all-reduced in FusedMoE, ensure they
309
# are all-reduced at the end of mlp.forward() if tensor
310
# parallel or expert parallel is enabled