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] Add EPLB support for hunyuan_v1 (#23078)
YiwenC committed
Sep 17, 2025 at 21:51 UTC
9d8a2d86d24b8afd849d18ddb4ef51cec1c0471d
2 files changed
+123
-16
vllm/model_executor/layers/fused_moe/layer.py
+2
-2
@@ -1508,8 +1508,8 @@ class FusedMoE(CustomOp):
1508
1509
return [
1510
weight.view(self.local_num_experts, -1) for name, weight in weights
1511
- if name not in NON_EXPERT_WEIGHTS
1512
- and not name.startswith("_shared_experts.")
1511
+ if name not in NON_EXPERT_WEIGHTS and weight.shape != torch.Size(
1512
+ []) and not name.startswith("_shared_experts.")
1513
]
1514
1515
def set_eplb_state(
vllm/model_executor/models/hunyuan_v1.py
+121
-14
@@ -23,7 +23,8 @@
23
# See the License for the specific language governing permissions and
24
# limitations under the License.
25
"""Inference-only HunYuan model compatible with HuggingFace weights."""
26
-from collections.abc import Iterable
26
+import typing
27
+from collections.abc import Callable, Iterable
28
from typing import Any, Optional, Union
29
30
import regex as re
@@ -33,8 +34,8 @@ from transformers import PretrainedConfig
34
35
from vllm.attention import Attention, AttentionType
36
from vllm.compilation.decorators import support_torch_compile
36
-from vllm.config import CacheConfig, VllmConfig
37
-from vllm.distributed import (get_pp_group,
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
tensor_model_parallel_all_reduce)
41
from vllm.model_executor.layers.activation import SiluAndMul
@@ -56,7 +57,7 @@ from vllm.model_executor.model_loader.weight_utils import (
57
from vllm.model_executor.sampling_metadata import SamplingMetadata
58
from vllm.sequence import IntermediateTensors
59
59
-from .interfaces import SupportsLoRA, SupportsPP
60
+from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
61
from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
62
make_layers, maybe_prefix)
63
@@ -355,10 +356,16 @@ class HunYuanSparseMoeBlock(nn.Module):
356
quant_config: Optional[QuantizationConfig] = None,
357
layer_id: int = -1,
358
prefix: str = "",
359
+ enable_eplb: bool = False,
360
):
361
super().__init__()
362
self.tp_size = get_tensor_model_parallel_world_size()
363
364
+ self.ep_group = get_ep_group().device_group
365
+ self.ep_rank = self.ep_group.rank()
366
+ self.ep_size = self.ep_group.size()
367
+ self.n_routed_experts = config.num_experts
368
+
369
if self.tp_size > config.num_experts:
370
raise ValueError(
371
f"Tensor parallel size {self.tp_size} is greater than "
@@ -379,8 +386,23 @@ class HunYuanSparseMoeBlock(nn.Module):
386
config.moe_intermediate_size, int) else
387
config.moe_intermediate_size[layer_id])
388
389
+ # Load balancing settings.
390
+ vllm_config = get_current_vllm_config()
391
+ eplb_config = vllm_config.parallel_config.eplb_config
392
+ self.enable_eplb = enable_eplb
393
+
394
+ self.n_logical_experts = self.n_routed_experts
395
+ self.n_redundant_experts = eplb_config.num_redundant_experts
396
+ self.n_physical_experts = (self.n_logical_experts +
397
+ self.n_redundant_experts)
398
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
399
+ self.physical_expert_start = (self.ep_rank *
400
+ self.n_local_physical_experts)
401
+ self.physical_expert_end = (self.physical_expert_start +
402
+ self.n_local_physical_experts)
403
+
404
self.experts = FusedMoE(
383
- num_experts=config.num_experts,
405
+ num_experts=self.n_routed_experts,
406
top_k=top_k,
407
hidden_size=config.hidden_size,
408
intermediate_size=intermediate_size,
@@ -388,6 +410,8 @@ class HunYuanSparseMoeBlock(nn.Module):
410
renormalize=top_k > 1,
411
quant_config=quant_config,
412
prefix=f"{prefix}.experts",
413
+ enable_eplb=self.enable_eplb,
414
+ num_redundant_experts=self.n_redundant_experts,
415
)
416
417
self.gate = ReplicatedLinear(config.hidden_size,
@@ -446,6 +470,7 @@ class HunYuanDecoderLayer(nn.Module):
470
quant_config: Optional[QuantizationConfig] = None,
471
prefix: str = "",
472
layer_id: int = -1,
473
+ enable_eplb: bool = False,
474
) -> None:
475
super().__init__()
476
assert layer_id >= 0
@@ -509,6 +534,7 @@ class HunYuanDecoderLayer(nn.Module):
534
quant_config=quant_config,
535
layer_id=layer_id,
536
prefix=f"{prefix}.mlp",
537
+ enable_eplb=enable_eplb,
538
)
539
else:
540
self.mlp = HunYuanMLP(
@@ -562,6 +588,9 @@ class HunYuanModel(nn.Module):
588
cache_config = vllm_config.cache_config
589
quant_config = vllm_config.quant_config
590
lora_config = vllm_config.lora_config
591
+ eplb_config = vllm_config.parallel_config.eplb_config
592
+ enable_eplb = vllm_config.parallel_config.enable_eplb
593
+ self.num_redundant_experts = eplb_config.num_redundant_experts
594
595
self.config = config
596
self.quant_config = quant_config
@@ -588,6 +617,7 @@ class HunYuanModel(nn.Module):
617
cache_config=cache_config,
618
quant_config=quant_config,
619
prefix=prefix,
620
+ enable_eplb=enable_eplb,
621
),
622
prefix=f"{prefix}.layers",
623
)
@@ -674,6 +704,7 @@ class HunYuanModel(nn.Module):
704
ckpt_down_proj_name="down_proj",
705
ckpt_up_proj_name="up_proj",
706
num_experts=self.config.num_experts,
707
+ num_redundant_experts=self.num_redundant_experts,
708
)
709
else:
710
return []
@@ -803,25 +834,43 @@ class HunYuanModel(nn.Module):
834
# Skip loading extra bias for GPTQ models.
835
if name.endswith(".bias") and name not in params_dict:
836
continue
837
+ is_expert_weight = False
838
for mapping in expert_params_mapping:
839
param_name, weight_name, expert_id, shard_id = mapping
840
if weight_name not in name:
841
continue
810
- name = name.replace(weight_name, param_name)
811
- # Skip layers on other devices.
812
- if is_pp_missing_parameter(name, self):
842
+ # this is an expert weight and should not be
843
+ # attempted to load as other weights later
844
+ is_expert_weight = True
845
+
846
+ # Do not modify `name` since the loop may continue here
847
+ # Instead, create a new variable
848
+ name_mapped = name.replace(weight_name, param_name)
849
+ if is_pp_missing_parameter(name_mapped, self):
850
continue
814
- param = params_dict[name]
815
- weight_loader = param.weight_loader
816
- weight_loader(
851
+ param = params_dict[name_mapped]
852
+ # We should ask the weight loader to return success or not
853
+ # here since otherwise we may skip experts with other
854
+ # available replicas.
855
+ weight_loader = typing.cast(Callable[..., bool],
856
+ param.weight_loader)
857
+ success = weight_loader(
858
param,
859
loaded_weight,
819
- name,
860
+ name_mapped,
861
shard_id=shard_id,
862
expert_id=expert_id,
863
+ return_success=True,
864
)
823
- break
865
+ if success:
866
+ name = name_mapped
867
+ break
868
else:
869
+ if is_expert_weight:
870
+ # We've checked that this is an expert weight
871
+ # However it's not mapped locally to this rank
872
+ # So we simply skip it
873
+ continue
874
# Remapping the name of FP8 kv-scale.
875
name = maybe_remap_kv_scale_name(name, params_dict)
876
if name is None:
@@ -841,7 +890,7 @@ class HunYuanModel(nn.Module):
890
return loaded_params
891
892
844
-class HunYuanV1Base(nn.Module, SupportsLoRA, SupportsPP):
893
+class HunYuanV1Base(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
894
packed_modules_mapping = {
895
"qkv_proj": [
896
"q_proj",
@@ -883,6 +932,64 @@ class HunYuanV1Base(nn.Module, SupportsLoRA, SupportsPP):
932
else:
933
self.lm_head = PPMissingLayer()
934
935
+ # Set MoE hyperparameters
936
+ self.expert_weights = []
937
+ self.num_expert_groups = 1
938
+ self.moe_layers: list[FusedMoE] = []
939
+ example_layer = None
940
+ for layer in self.model.layers:
941
+ if isinstance(layer, PPMissingLayer):
942
+ continue
943
+
944
+ assert isinstance(layer, HunYuanDecoderLayer)
945
+ if isinstance(layer.mlp, HunYuanSparseMoeBlock):
946
+ example_layer = layer.mlp
947
+ self.moe_layers.append(layer.mlp.experts)
948
+
949
+ if example_layer is None:
950
+ raise RuntimeError("No HunYuanMoE layer found in model.layers.")
951
+
952
+ self.num_moe_layers = len(self.moe_layers)
953
+ self.num_logical_experts = example_layer.n_logical_experts
954
+ self.num_physical_experts = example_layer.n_physical_experts
955
+ self.num_local_physical_experts = example_layer.n_local_physical_experts
956
+ self.num_routed_experts = example_layer.n_routed_experts
957
+ self.num_redundant_experts = example_layer.n_redundant_experts
958
+
959
+ def set_eplb_state(
960
+ self,
961
+ expert_load_view: torch.Tensor,
962
+ logical_to_physical_map: torch.Tensor,
963
+ logical_replica_count: torch.Tensor,
964
+ ) -> None:
965
+ for layer_idx, layer in enumerate(self.moe_layers):
966
+ self.expert_weights.append(layer.get_expert_weights())
967
+ # Register the expert weights.
968
+ layer.set_eplb_state(
969
+ moe_layer_idx=layer_idx,
970
+ expert_load_view=expert_load_view,
971
+ logical_to_physical_map=logical_to_physical_map,
972
+ logical_replica_count=logical_replica_count,
973
+ )
974
+
975
+ def update_physical_experts_metadata(
976
+ self,
977
+ num_physical_experts: int,
978
+ num_local_physical_experts: int,
979
+ ) -> None:
980
+ assert self.num_local_physical_experts == num_local_physical_experts
981
+ self.num_physical_experts = num_physical_experts
982
+ self.num_local_physical_experts = num_local_physical_experts
983
+ self.num_redundant_experts = (num_physical_experts -
984
+ self.num_logical_experts)
985
+ for layer in self.model.layers:
986
+ if isinstance(layer.mlp, HunYuanSparseMoeBlock):
987
+ moe = layer.mlp
988
+ moe.n_local_physical_experts = num_local_physical_experts
989
+ moe.n_physical_experts = num_physical_experts
990
+ moe.n_redundant_experts = self.num_redundant_experts
991
+ moe.experts.update_expert_map()
992
+
993
def forward(
994
self,
995
input_ids: torch.Tensor,