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

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

[Bugfix] fix deepseek fp16 scale bug (#14809)

Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>

Jinzhen Lin committed Apr 9, 2025 at 04:56 UTC db104221845e44e0b85336a54e44eb8850b8a5bc

1 file changed +24 -10

vllm/model_executor/models/deepseek_v2.py

+24 -10

@@ -160,14 +160,16 @@ class DeepseekV2MoE(nn.Module):

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

163 + # Fix FP16 overflow

164 + # See DeepseekV2DecoderLayer for more details.

165 final_hidden_states = self.experts(hidden_states=hidden_states,

166 router_logits=router_logits)

167 if shared_output is not None:

168 if hidden_states.dtype != torch.float16:

169 final_hidden_states = final_hidden_states + shared_output

170 else:

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

171 + # Fix FP16 overflow

172 + # See DeepseekV2DecoderLayer for more details.

173 final_hidden_states = final_hidden_states + shared_output \

174 * (1. / self.routed_scaling_factor)

175 if self.tp_size > 1:

@@ -499,6 +501,7 @@ class DeepseekV2DecoderLayer(nn.Module):

501 # DecoderLayers are created with `make_layers` which passes the prefix

502 # with the layer's index.

503 layer_idx = int(prefix.split(sep='.')[-1])

504 + self.layer_idx = layer_idx

505 if model_config.use_mla:

506 attn_cls = DeepseekV2MLAAttention

507 else:

@@ -561,19 +564,30 @@ class DeepseekV2DecoderLayer(nn.Module):

564 hidden_states=hidden_states,

565 )

566

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

567 + if hidden_states.dtype == torch.float16:

568 + # Fix FP16 overflow

569 + # We scale both hidden_states and residual before

570 + # rmsnorm, and rmsnorm result would not affect by scale.

571 hidden_states *= 1. / self.routed_scaling_factor

572 + if self.layer_idx == 0:

573 + # The residual is shared by all layers, we only scale it on

574 + # first layer.

575 + residual *= 1. / self.routed_scaling_factor

576 +

577 + # Fully Connected

578 hidden_states, residual = self.post_attention_layernorm(

579 hidden_states, residual)

580 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

581 +

582 + if isinstance(self.mlp,

583 + DeepseekV2MLP) and hidden_states.dtype == torch.float16:

584 + # Fix FP16 overflow

585 + # Scaling the DeepseekV2MLP output, it is the input of

586 + # input_layernorm of next decoder layer.

587 + # The scaling of DeepseekV2MOE output would be done in the forward

588 + # of DeepseekV2MOE

589 hidden_states *= 1. / self.routed_scaling_factor

576 - residual *= 1. / self.routed_scaling_factor

590 +

591 return hidden_states, residual

592

593