Update modeling_deepseek.py

modeling_deepseek.py

@@ -552,7 +552,9 @@ class DeepseekV2MoE(nn.Module):
             self.ep_rank = 0
             self.experts = nn.ModuleList(
                 [
-                    DeepseekV2MLP(config, intermediate_size=config.moe_intermediate_size)
+                    DeepseekV2MLP(
+                        config, intermediate_size=config.moe_intermediate_size
+                    )
                     for i in range(config.n_routed_experts)
                 ]
             )
@@ -577,7 +579,7 @@ class DeepseekV2MoE(nn.Module):
             for i, expert in enumerate(self.experts):
                 y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i])
             y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
-            y = y.view(*orig_shape)
+            y = y.to(hidden_states.dtype).view(*orig_shape)
             y = AddAuxiliaryLoss.apply(y, aux_loss)
         else:
             y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
@@ -1023,7 +1025,11 @@ class DeepseekV2FlashAttention2(DeepseekV2Attention):
             elif torch.is_autocast_enabled():
                 target_dtype = torch.get_autocast_gpu_dtype()
             else:
-                target_dtype = self.q_proj.weight.dtype if self.q_lora_rank is None else self.q_a_proj.weight.dtype
+                target_dtype = (
+                    self.q_proj.weight.dtype
+                    if self.q_lora_rank is None
+                    else self.q_a_proj.weight.dtype
+                )
 
             logger.warning_once(
                 f"The input hidden states seems to be silently casted in float32, this might be related to"