From 8369affa5428ca29a30f378d660fbffbad12e240 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sat, 5 Sep 2026 13:28:33 +0800 Subject: [PATCH] [Feat] Add EPLB support for GLM-5.3-Flash (#55119) Signed-off-by: chaunceyjiang Signed-off-by: Jared Wen Co-authored-by: Jared Wen Co-authored-by: Claude Fable 5.1 --- vllm/models/glm5next/nvidia/model.py | 77 ++++++++++++++++++++++++---- vllm/models/glm5next/nvidia/mtp.py | 1 + 2 files changed, 69 insertions(+), 9 deletions(-) diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py index 9988dd84ea..5d632df137 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/nvidia/model.py @@ -725,12 +725,24 @@ class Glm5NextModel(nn.Module): if self.config.is_moe: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) + # EPLB: the mapping enumerates physical experts, so it must cover + # the redundant replicas or their slots are never loaded. + num_redundant_experts = next( + ( + layer.mlp.n_redundant_experts + for layer in self.layers + if isinstance(layer, Glm5NextDecoderLayer) + and isinstance(layer.mlp, Glm5NextMoE) + ), + 0, + ) expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", num_experts=self.config.n_routed_experts, + num_redundant_experts=num_redundant_experts, ) else: expert_params_mapping = [] @@ -818,28 +830,38 @@ class Glm5NextModel(nn.Module): weight_loader(param, loaded_weight, shard_id) break else: - for idx, ( + is_expert_weight = False + for ( param_name, weight_name, expert_id, expert_shard_id, - ) in enumerate(expert_params_mapping): + ) in expert_params_mapping: if weight_name not in name: continue - name = name.replace(weight_name, param_name) - if is_pp_missing_parameter(name, self): + # A checkpoint expert may map to several physical replicas + # under EPLB; keep `name` intact and try the next entry + # when this physical expert is not local to the rank. + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name_mapped, self): continue - param = params_dict[name] + param = params_dict[name_mapped] weight_loader = param.weight_loader - weight_loader( + success = weight_loader( param, loaded_weight, - name, + name_mapped, expert_id=expert_id, shard_id=expert_shard_id, + return_success=True, ) - break + if success: + name = name_mapped + break else: + if is_expert_weight: + continue # Skip loading extra bias for GPTQ models. if ( name.endswith(".bias") @@ -961,7 +983,7 @@ class Glm5NextForCausalLM( dummy_inputs=Glm4vDummyInputsBuilder, ) class Glm5NextForConditionalGeneration( - Glm4vForConditionalGeneration, HasInnerState, IsHybrid + Glm4vForConditionalGeneration, HasInnerState, IsHybrid, MixtureOfExperts ): # The text model (KDA + dense-MLA + MoE) is a hybrid mamba model. The # multimodal wrapper must declare the same interfaces so vLLM treats it as @@ -1036,10 +1058,47 @@ class Glm5NextForConditionalGeneration( architectures=["Glm5NextForCausalLM"], ) + self.set_moe_parameters() + # Glm5NextForCausalLM does not implement make_empty_intermediate_tensors, # so pipeline parallelism is gated off (consistent with the text-only # model) and we intentionally do not alias it here. + def set_moe_parameters(self) -> None: + self.moe_mlp_layers = [ + layer.mlp + for layer in self.language_model.model.layers + if isinstance(layer, Glm5NextDecoderLayer) + and isinstance(layer.mlp, Glm5NextMoE) + ] + self.moe_layers = [moe.experts for moe in self.moe_mlp_layers] + self.num_moe_layers = len(self.moe_layers) + if not self.num_moe_layers: + return + example_moe = self.moe_mlp_layers[0] + self.num_expert_groups = self.config.text_config.n_group + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_shared_experts = example_moe.n_shared_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + if not self.num_moe_layers: + return + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for moe in self.moe_mlp_layers: + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + def get_encoder_cudagraph_config(self): # This vision tower does not produce the absolute position embedding # buffer used by GLM4V. diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/nvidia/mtp.py index ae1f7f56b5..35b452f44f 100644 --- a/vllm/models/glm5next/nvidia/mtp.py +++ b/vllm/models/glm5next/nvidia/mtp.py @@ -303,6 +303,7 @@ class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts): ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", num_experts=self.config.n_routed_experts, + num_redundant_experts=self.num_redundant_experts, ) params_dict = dict(self.named_parameters())