@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ef6e0e7132

Fork 0

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

[Bugfix][Model]fix ernie45 moe gate&bias dtype to float32 (#25936)

Signed-off-by: wangyafeng <wangyafeng@baidu.com>

CSWYF3634076 committed Sep 30, 2025 at 19:11 UTC ef6e0e7132ec7bd3dadd10df7486a8cd10cf9c0e

2 files changed +13 -7

vllm/model_executor/models/ernie45_moe.py

+3 -2

@@ -120,11 +120,12 @@ class Ernie4_5_MoeMoE(nn.Module):

120 self.gate = ReplicatedLinear(config.hidden_size,

121 config.moe_num_experts,

122 bias=False,

123 + params_dtype=torch.float32,

124 quant_config=None,

125 prefix=f"{prefix}.gate")

126

127 self.gate.e_score_correction_bias = nn.Parameter(

127 - torch.empty(config.moe_num_experts))

128 + torch.empty(config.moe_num_experts, dtype=torch.float32))

129

130 self.experts = FusedMoE(

131 num_experts=config.moe_num_experts,

@@ -157,7 +158,7 @@ class Ernie4_5_MoeMoE(nn.Module):

158 if self.has_shared_experts:

159 shared_output = self.shared_experts(hidden_states)

160

160 - router_logits, _ = self.gate(hidden_states)

161 + router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))

162

163 final_hidden_states = self.experts(hidden_states=hidden_states,

164 router_logits=router_logits)

vllm/model_executor/models/ernie45_vl_moe.py

+10 -5

@@ -199,7 +199,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):

199

200 assert config.moe_num_experts[0] == config.moe_num_experts[1]

201 self.e_score_correction_bias = nn.Parameter(

202 - torch.empty(2, config.moe_num_experts[0]))

202 + torch.empty(2, config.moe_num_experts[0], dtype=torch.float32))

203

204 assert text_moe_layer_start_index <= text_moe_layer_end_index

205

@@ -209,6 +209,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):

209 config.hidden_size,

210 config.moe_num_experts[0],

211 bias=False,

212 + params_dtype=torch.float32,

213 quant_config=quant_config,

214 prefix=f"{prefix}.text_experts_gate")

215

@@ -238,6 +239,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):

239 config.hidden_size,

240 config.moe_num_experts[1],

241 bias=False,

242 + params_dtype=torch.float32,

243 quant_config=quant_config,

244 prefix=f"{prefix}.vision_experts_gate")

245

@@ -288,7 +290,8 @@ class Ernie4_5_VLMoeMoE(nn.Module):

290

291 if visual_token_mask is not None and visual_token_mask.all():

292 # only vision modal input

291 - router_logits, _ = self.vision_experts_gate(hidden_states)

293 + router_logits, _ = self.vision_experts_gate(

294 + hidden_states.to(dtype=torch.float32))

295 final_hidden_states = self.vision_experts(

296 hidden_states=hidden_states, router_logits=router_logits)

297 elif visual_token_mask is not None and visual_token_mask.any():

@@ -303,19 +306,21 @@ class Ernie4_5_VLMoeMoE(nn.Module):

306 vision_hidden_states = hidden_states[visual_token_mask].reshape(

307 -1, self.hidden_size)

308

306 - text_router_logits, _ = self.text_experts_gate(text_hidden_states)

309 + text_router_logits, _ = self.text_experts_gate(

310 + text_hidden_states.to(dtype=torch.float32))

311 final_hidden_states[text_token_mask] = self.text_experts(

312 hidden_states=text_hidden_states,

313 router_logits=text_router_logits).flatten()

314

315 vision_router_logits, _ = self.vision_experts_gate(

312 - vision_hidden_states)

316 + vision_hidden_states.to(dtype=torch.float32))

317 final_hidden_states[visual_token_mask] = self.vision_experts(

318 hidden_states=vision_hidden_states,

319 router_logits=vision_router_logits).flatten()

320 else:

321 # only text modal input

318 - text_router_logits, _ = self.text_experts_gate(hidden_states)

322 + text_router_logits, _ = self.text_experts_gate(

323 + hidden_states.to(dtype=torch.float32))

324

325 final_hidden_states = self.text_experts(

326 hidden_states=hidden_states, router_logits=text_router_logits)