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

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 15, 19:45 UTC

[Bugfix] Fix FP16 overflow for DeepSeek V2 (#13232)

Signed-off-by: Yida Wu <yida.wu@amd.com>

Concurrensee committed Mar 10, 2025 at 22:46 UTC c982ac572248dc394ccce076c42277e8364a67ee

1 file changed +24 -4

vllm/model_executor/models/deepseek_v2.py

+24 -4

@@ -155,11 +155,21 @@ class DeepseekV2MoE(nn.Module):

155 shared_output = self.shared_experts(hidden_states)

156 # router_logits: (num_tokens, n_experts)

157 router_logits, _ = self.gate(hidden_states)

158 - final_hidden_states = self.experts(

159 - hidden_states=hidden_states,

160 - router_logits=router_logits) * self.routed_scaling_factor

158 + if hidden_states.dtype != torch.float16:

159 + final_hidden_states = self.experts(

160 + hidden_states=hidden_states,

161 + router_logits=router_logits) * self.routed_scaling_factor

162 + else:

163 + # This is a special case to avoid FP16 overflow

164 + final_hidden_states = self.experts(hidden_states=hidden_states,

165 + router_logits=router_logits)

166 if shared_output is not None:

162 - final_hidden_states = final_hidden_states + shared_output

167 + if hidden_states.dtype != torch.float16:

168 + final_hidden_states = final_hidden_states + shared_output

169 + else:

170 + # This is a special case to avoid FP16 overflow

171 + final_hidden_states = final_hidden_states + shared_output \

172 + * (1. / self.routed_scaling_factor)

173 if self.tp_size > 1:

174 final_hidden_states = tensor_model_parallel_all_reduce(

175 final_hidden_states)

@@ -531,6 +541,7 @@ class DeepseekV2DecoderLayer(nn.Module):

541 eps=config.rms_norm_eps)

542 self.post_attention_layernorm = RMSNorm(config.hidden_size,

543 eps=config.rms_norm_eps)

544 + self.routed_scaling_factor = config.routed_scaling_factor

545

546 def forward(

547 self,

@@ -551,9 +562,18 @@ class DeepseekV2DecoderLayer(nn.Module):

562 )

563

564 # Fully Connected

565 + if isinstance(self.mlp, DeepseekV2MoE) and \

566 + hidden_states.dtype == torch.float16:

567 + # This is a special case to avoid FP16 overflow

568 + hidden_states *= 1. / self.routed_scaling_factor

569 hidden_states, residual = self.post_attention_layernorm(

570 hidden_states, residual)

571 hidden_states = self.mlp(hidden_states)

572 + if isinstance(self.mlp, DeepseekV2MLP) and \

573 + hidden_states.dtype == torch.float16:

574 + # This is a special case to avoid FP16 overflow

575 + hidden_states *= 1. / self.routed_scaling_factor

576 + residual *= 1. / self.routed_scaling_factor

577 return hidden_states, residual

578

579