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
[EPLB] Support EPLB for Mixtral Model (#22842)
Signed-off-by: rouchenzi <ruochenwen@gmail.com> Signed-off-by: rouchenzi <40842833+rouchenzi@users.noreply.github.com> Co-authored-by: Bowen Wang <abmfy@icloud.com>
rouchenzi committed
Sep 17, 2025 at 00:27 UTC
b77bf34e531abb32c054a38747fa817d08395ae7
1 file changed
+137
-23
vllm/model_executor/models/mixtral.py
+137
-23
@@ -23,7 +23,8 @@
23
# See the License for the specific language governing permissions and
24
# limitations under the License.
25
"""Inference-only Mixtral model."""
26
-from collections.abc import Iterable
26
+import typing
27
+from collections.abc import Callable, Iterable
28
from itertools import islice
29
from typing import Optional, Union
30
@@ -33,8 +34,9 @@ from transformers import MixtralConfig
34
35
from vllm.attention import Attention
36
from vllm.compilation.decorators import support_torch_compile
36
-from vllm.config import CacheConfig, VllmConfig
37
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
38
+from vllm.distributed import (get_ep_group, get_pp_group,
39
+ get_tensor_model_parallel_world_size)
40
from vllm.model_executor.layers.fused_moe import FusedMoE
41
from vllm.model_executor.layers.layernorm import RMSNorm
42
from vllm.model_executor.layers.linear import (QKVParallelLinear,
@@ -50,8 +52,8 @@ from vllm.model_executor.model_loader.weight_utils import (
52
from vllm.model_executor.sampling_metadata import SamplingMetadata
53
from vllm.sequence import IntermediateTensors
54
53
-from .interfaces import SupportsLoRA, SupportsPP
54
-from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
55
+from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
56
+from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
57
make_empty_intermediate_tensors_factory, make_layers,
58
maybe_prefix)
59
@@ -74,10 +76,32 @@ class MixtralMoE(nn.Module):
76
quant_config: Optional[QuantizationConfig] = None,
77
tp_size: Optional[int] = None,
78
dp_size: Optional[int] = None,
77
- prefix: str = ""):
79
+ prefix: str = "",
80
+ enable_eplb: bool = False):
81
super().__init__()
82
self.hidden_size = hidden_size
83
84
+ self.ep_group = get_ep_group().device_group
85
+ self.ep_rank = self.ep_group.rank()
86
+ self.ep_size = self.ep_group.size()
87
+
88
+ # Expert Parallelism Load balancing settings.
89
+ vllm_config = get_current_vllm_config()
90
+ parallel_config = vllm_config.parallel_config
91
+ self.enable_eplb = enable_eplb
92
+
93
+ self.n_routed_experts = num_experts
94
+ self.n_logical_experts = num_experts
95
+ self.n_redundant_experts = (
96
+ parallel_config.eplb_config.num_redundant_experts)
97
+ self.n_physical_experts = (self.n_logical_experts +
98
+ self.n_redundant_experts)
99
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
100
+ self.physical_expert_start = (self.ep_rank *
101
+ self.n_local_physical_experts)
102
+ self.physical_expert_end = (self.physical_expert_start +
103
+ self.n_local_physical_experts)
104
+
105
# Gate always runs at half / full precision for now.
106
107
self.gate = ReplicatedLinear(hidden_size,
@@ -97,7 +121,9 @@ class MixtralMoE(nn.Module):
121
quant_config=quant_config,
122
tp_size=tp_size,
123
dp_size=dp_size,
100
- prefix=f"{prefix}.experts")
124
+ prefix=f"{prefix}.experts",
125
+ enable_eplb=self.enable_eplb,
126
+ num_redundant_experts=self.n_redundant_experts)
127
128
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
129
# NOTE: hidden_states can have either 1D or 2D shape.
@@ -200,6 +226,7 @@ class MixtralDecoderLayer(nn.Module):
226
cache_config: Optional[CacheConfig] = None,
227
quant_config: Optional[QuantizationConfig] = None,
228
prefix: str = "",
229
+ enable_eplb: bool = False,
230
) -> None:
231
super().__init__()
232
self.hidden_size = config.hidden_size
@@ -221,7 +248,8 @@ class MixtralDecoderLayer(nn.Module):
248
hidden_size=config.hidden_size,
249
intermediate_size=config.intermediate_size,
250
quant_config=quant_config,
224
- prefix=f"{prefix}.block_sparse_moe")
251
+ prefix=f"{prefix}.block_sparse_moe",
252
+ enable_eplb=enable_eplb)
253
self.input_layernorm = RMSNorm(config.hidden_size,
254
eps=config.rms_norm_eps)
255
self.post_attention_layernorm = RMSNorm(config.hidden_size,
@@ -262,6 +290,7 @@ class MixtralModel(nn.Module):
290
cache_config = vllm_config.cache_config
291
quant_config = vllm_config.quant_config
292
lora_config = vllm_config.lora_config
293
+ parallel_config = vllm_config.parallel_config
294
295
self.config = config
296
self.quant_config = quant_config
@@ -276,10 +305,18 @@ class MixtralModel(nn.Module):
305
org_num_embeddings=config.vocab_size,
306
)
307
308
+ self.enable_eplb = parallel_config.enable_eplb
309
+ self.num_redundant_experts = (
310
+ parallel_config.eplb_config.num_redundant_experts)
311
+
312
self.start_layer, self.end_layer, self.layers = make_layers(
313
config.num_hidden_layers,
314
lambda prefix: MixtralDecoderLayer(
282
- config, cache_config, quant_config=quant_config, prefix=prefix
315
+ config,
316
+ cache_config,
317
+ quant_config=quant_config,
318
+ prefix=prefix,
319
+ enable_eplb=self.enable_eplb,
320
),
321
prefix=f"{prefix}.layers")
322
@@ -325,7 +362,8 @@ class MixtralModel(nn.Module):
362
ckpt_gate_proj_name="w1",
363
ckpt_down_proj_name="w2",
364
ckpt_up_proj_name="w3",
328
- num_experts=self.config.num_local_experts)
365
+ num_experts=self.config.num_local_experts,
366
+ num_redundant_experts=self.num_redundant_experts)
367
368
def load_weights(self, weights: Iterable[tuple[str,
369
torch.Tensor]]) -> set[str]:
@@ -373,26 +411,40 @@ class MixtralModel(nn.Module):
411
weight_loader(param, loaded_weight, shard_id)
412
break
413
else:
414
+ is_expert_weight = False
415
for mapping in expert_params_mapping:
416
param_name, weight_name, expert_id, shard_id = mapping
417
+
418
if weight_name not in name:
419
continue
380
- name = name.replace(weight_name, param_name)
420
+
421
+ is_expert_weight = True
422
+ name_mapped = name.replace(weight_name, param_name)
423
+
424
# Skip layers on other devices.
382
- if is_pp_missing_parameter(name, self):
425
+ if is_pp_missing_parameter(name_mapped, self):
426
continue
384
- if ((name.endswith(".bias") or name.endswith("_bias"))
385
- and name not in params_dict):
427
+
428
+ if ((name_mapped.endswith(".bias")
429
+ or name_mapped.endswith("_bias"))
430
+ and name_mapped not in params_dict):
431
continue
387
- param = params_dict[name]
388
- weight_loader = param.weight_loader
389
- weight_loader(param,
390
- loaded_weight,
391
- name,
392
- shard_id=shard_id,
393
- expert_id=expert_id)
394
- break
432
+
433
+ param = params_dict[name_mapped]
434
+ weight_loader = typing.cast(Callable[..., bool],
435
+ param.weight_loader)
436
+ success = weight_loader(param,
437
+ loaded_weight,
438
+ name_mapped,
439
+ shard_id=shard_id,
440
+ expert_id=expert_id,
441
+ return_success=True)
442
+ if success:
443
+ name = name_mapped
444
+ break
445
else:
446
+ if is_expert_weight:
447
+ continue
448
# Skip loading extra bias for GPTQ models.
449
if ((name.endswith(".bias") or name.endswith("_bias"))
450
and name not in params_dict):
@@ -413,7 +465,8 @@ class MixtralModel(nn.Module):
465
return loaded_params
466
467
416
-class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
468
+class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP,
469
+ MixtureOfExperts):
470
fall_back_to_pt_during_load = False
471
472
packed_modules_mapping = {
@@ -462,6 +515,67 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
515
self.make_empty_intermediate_tensors = (
516
self.model.make_empty_intermediate_tensors)
517
518
+ self.expert_weights = []
519
+ self.moe_layers: list[FusedMoE] = []
520
+ example_moe = None
521
+
522
+ for layer in self.model.layers:
523
+ if isinstance(layer, PPMissingLayer):
524
+ continue
525
+ assert isinstance(layer, MixtralDecoderLayer)
526
+ if hasattr(layer, "block_sparse_moe") and isinstance(
527
+ layer.block_sparse_moe, MixtralMoE):
528
+ example_moe = layer.block_sparse_moe
529
+ self.moe_layers.append(layer.block_sparse_moe.experts)
530
+
531
+ self.num_moe_layers = len(self.moe_layers)
532
+
533
+ if example_moe is None:
534
+ raise RuntimeError("No MixtralMoE layer found in model.layers.")
535
+
536
+ self.num_logical_experts = example_moe.n_logical_experts
537
+ self.num_physical_experts = example_moe.n_physical_experts
538
+ self.num_local_physical_experts = example_moe.n_local_physical_experts
539
+ self.num_routed_experts = example_moe.n_routed_experts
540
+ self.num_redundant_experts = example_moe.n_redundant_experts
541
+ self.num_expert_groups = 1
542
+ self.num_shared_experts = 0
543
+
544
+ def set_eplb_state(
545
+ self,
546
+ expert_load_view: torch.Tensor,
547
+ logical_to_physical_map: torch.Tensor,
548
+ logical_replica_count: torch.Tensor,
549
+ ) -> None:
550
+ for layer_idx, layer in enumerate(self.moe_layers):
551
+ # Register the expert weights.
552
+ self.expert_weights.append(layer.get_expert_weights())
553
+ layer.set_eplb_state(
554
+ moe_layer_idx=layer_idx,
555
+ expert_load_view=expert_load_view,
556
+ logical_to_physical_map=logical_to_physical_map,
557
+ logical_replica_count=logical_replica_count,
558
+ )
559
+
560
+ def update_physical_experts_metadata(
561
+ self,
562
+ num_physical_experts: int,
563
+ num_local_physical_experts: int,
564
+ ) -> None:
565
+ assert self.num_local_physical_experts == num_local_physical_experts
566
+ self.num_physical_experts = num_physical_experts
567
+ self.num_local_physical_experts = num_local_physical_experts
568
+ self.num_redundant_experts = (num_physical_experts -
569
+ self.num_logical_experts)
570
+ for layer in self.model.layers:
571
+ if hasattr(layer, "block_sparse_moe") and isinstance(
572
+ layer.block_sparse_moe, MixtralMoE):
573
+ moe = layer.block_sparse_moe
574
+ moe.n_local_physical_experts = num_local_physical_experts
575
+ moe.n_physical_experts = num_physical_experts
576
+ moe.n_redundant_experts = self.num_redundant_experts
577
+ moe.experts.update_expert_map()
578
+
579
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
580
return self.model.get_input_embeddings(input_ids)
581