[Feat] Add EPLB support for GLM-5.3-Flash (#55119)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai> Co-authored-by: Jared Wen <jaredwen@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Chauncey committed
Sep 5, 2026 at 13:28 UTC
8369affa5428ca29a30f378d660fbffbad12e240
2 files changed
+69
-9
vllm/models/glm5next/nvidia/model.py
+68
-9
@@ -725,12 +725,24 @@ class Glm5NextModel(nn.Module):
725
if self.config.is_moe:
726
# Params for weights, fp8 weight scales, fp8 activation scales
727
# (param_name, weight_name, expert_id, shard_id)
728
+ # EPLB: the mapping enumerates physical experts, so it must cover
729
+ # the redundant replicas or their slots are never loaded.
730
+ num_redundant_experts = next(
731
+ (
732
+ layer.mlp.n_redundant_experts
733
+ for layer in self.layers
734
+ if isinstance(layer, Glm5NextDecoderLayer)
735
+ and isinstance(layer.mlp, Glm5NextMoE)
736
+ ),
737
+ 0,
738
+ )
739
expert_params_mapping = fused_moe_make_expert_params_mapping(
740
self,
741
ckpt_gate_proj_name="gate_proj",
742
ckpt_down_proj_name="down_proj",
743
ckpt_up_proj_name="up_proj",
744
num_experts=self.config.n_routed_experts,
745
+ num_redundant_experts=num_redundant_experts,
746
)
747
else:
748
expert_params_mapping = []
@@ -818,28 +830,38 @@ class Glm5NextModel(nn.Module):
830
weight_loader(param, loaded_weight, shard_id)
831
break
832
else:
821
- for idx, (
833
+ is_expert_weight = False
834
+ for (
835
param_name,
836
weight_name,
837
expert_id,
838
expert_shard_id,
826
- ) in enumerate(expert_params_mapping):
839
+ ) in expert_params_mapping:
840
if weight_name not in name:
841
continue
829
- name = name.replace(weight_name, param_name)
830
- if is_pp_missing_parameter(name, self):
842
+ # A checkpoint expert may map to several physical replicas
843
+ # under EPLB; keep `name` intact and try the next entry
844
+ # when this physical expert is not local to the rank.
845
+ is_expert_weight = True
846
+ name_mapped = name.replace(weight_name, param_name)
847
+ if is_pp_missing_parameter(name_mapped, self):
848
continue
832
- param = params_dict[name]
849
+ param = params_dict[name_mapped]
850
weight_loader = param.weight_loader
834
- weight_loader(
851
+ success = weight_loader(
852
param,
853
loaded_weight,
837
- name,
854
+ name_mapped,
855
expert_id=expert_id,
856
shard_id=expert_shard_id,
857
+ return_success=True,
858
)
841
- break
859
+ if success:
860
+ name = name_mapped
861
+ break
862
else:
863
+ if is_expert_weight:
864
+ continue
865
# Skip loading extra bias for GPTQ models.
866
if (
867
name.endswith(".bias")
@@ -961,7 +983,7 @@ class Glm5NextForCausalLM(
983
dummy_inputs=Glm4vDummyInputsBuilder,
984
)
985
class Glm5NextForConditionalGeneration(
964
- Glm4vForConditionalGeneration, HasInnerState, IsHybrid
986
+ Glm4vForConditionalGeneration, HasInnerState, IsHybrid, MixtureOfExperts
987
):
988
# The text model (KDA + dense-MLA + MoE) is a hybrid mamba model. The
989
# multimodal wrapper must declare the same interfaces so vLLM treats it as
@@ -1036,10 +1058,47 @@ class Glm5NextForConditionalGeneration(
1058
architectures=["Glm5NextForCausalLM"],
1059
)
1060
1061
+ self.set_moe_parameters()
1062
+
1063
# Glm5NextForCausalLM does not implement make_empty_intermediate_tensors,
1064
# so pipeline parallelism is gated off (consistent with the text-only
1065
# model) and we intentionally do not alias it here.
1066
1067
+ def set_moe_parameters(self) -> None:
1068
+ self.moe_mlp_layers = [
1069
+ layer.mlp
1070
+ for layer in self.language_model.model.layers
1071
+ if isinstance(layer, Glm5NextDecoderLayer)
1072
+ and isinstance(layer.mlp, Glm5NextMoE)
1073
+ ]
1074
+ self.moe_layers = [moe.experts for moe in self.moe_mlp_layers]
1075
+ self.num_moe_layers = len(self.moe_layers)
1076
+ if not self.num_moe_layers:
1077
+ return
1078
+ example_moe = self.moe_mlp_layers[0]
1079
+ self.num_expert_groups = self.config.text_config.n_group
1080
+ self.num_logical_experts = example_moe.n_logical_experts
1081
+ self.num_physical_experts = example_moe.n_physical_experts
1082
+ self.num_local_physical_experts = example_moe.n_local_physical_experts
1083
+ self.num_routed_experts = example_moe.n_routed_experts
1084
+ self.num_shared_experts = example_moe.n_shared_experts
1085
+ self.num_redundant_experts = example_moe.n_redundant_experts
1086
+
1087
+ def update_physical_experts_metadata(
1088
+ self,
1089
+ num_physical_experts: int,
1090
+ num_local_physical_experts: int,
1091
+ ) -> None:
1092
+ if not self.num_moe_layers:
1093
+ return
1094
+ assert self.num_local_physical_experts == num_local_physical_experts
1095
+ self.num_physical_experts = num_physical_experts
1096
+ self.num_redundant_experts = num_physical_experts - self.num_logical_experts
1097
+ for moe in self.moe_mlp_layers:
1098
+ moe.n_physical_experts = num_physical_experts
1099
+ moe.n_redundant_experts = self.num_redundant_experts
1100
+ moe.experts.update_expert_map()
1101
+
1102
def get_encoder_cudagraph_config(self):
1103
# This vision tower does not produce the absolute position embedding
1104
# buffer used by GLM4V.
vllm/models/glm5next/nvidia/mtp.py
+1
@@ -303,6 +303,7 @@ class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts):
303
ckpt_down_proj_name="down_proj",
304
ckpt_up_proj_name="up_proj",
305
num_experts=self.config.n_routed_experts,
306
+ num_redundant_experts=self.num_redundant_experts,
307
)
308
309
params_dict = dict(self.named_parameters())