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, 23:45 UTC
[Model] Fix Ernie4.5MoE e_score_correction_bias parameter (#21586)
Signed-off-by: zhouchong <zhouchong03@baidu.com> Co-authored-by: zhouchong <zhouchong03@baidu.com>
xyxinyang committed
Jul 25, 2025 at 21:02 UTC
c72f049cb4c98b92cb85fc791e74d9220a263cf2
1 file changed
+17
-8
vllm/model_executor/models/ernie45_moe.py
+17
-8
@@ -123,14 +123,19 @@ class Ernie4_5_MoeMoE(nn.Module):
123
quant_config=None,
124
prefix=f"{prefix}.gate")
125
126
- self.experts = FusedMoE(num_experts=config.moe_num_experts,
127
- top_k=config.moe_k,
128
- hidden_size=config.hidden_size,
129
- intermediate_size=config.moe_intermediate_size,
130
- reduce_results=False,
131
- renormalize=True,
132
- quant_config=quant_config,
133
- prefix=f"{prefix}.experts")
126
+ self.gate.e_score_correction_bias = nn.Parameter(
127
+ torch.empty(config.moe_num_experts))
128
+
129
+ self.experts = FusedMoE(
130
+ num_experts=config.moe_num_experts,
131
+ top_k=config.moe_k,
132
+ hidden_size=config.hidden_size,
133
+ intermediate_size=config.moe_intermediate_size,
134
+ reduce_results=False,
135
+ renormalize=True,
136
+ quant_config=quant_config,
137
+ prefix=f"{prefix}.experts",
138
+ e_score_correction_bias=self.gate.e_score_correction_bias)
139
140
if self.moe_num_shared_experts is not None:
141
intermediate_size = (config.moe_intermediate_size *
@@ -459,6 +464,10 @@ class Ernie4_5_MoeModel(nn.Module):
464
if "mtp" in name:
465
continue
466
467
+ if "e_score_correction_bias" in name:
468
+ name = name.replace("moe_statics", "gate")
469
+ loaded_weight = loaded_weight.squeeze(0)
470
+
471
for (param_name, weight_name, shard_id) in stacked_params_mapping:
472
# Skip non-stacked layers and experts (experts handled below).
473
if weight_name not in name: