Update modelling_deepseek.py

modelling_deepseek.py

@@ -265,7 +265,30 @@ class MoEGate(nn.Module):
             topk_weight = topk_weight / denominator
 
         # Expert-level computation auxiliary loss
-        aux_loss = None
+        # (was absent before)
+        if self.training and self.alpha > 0.0:
+            scores_for_aux = scores
+            aux_topk = self.top_k
+            # always compute aux loss based on the naive greedy topk method
+            topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
+            if self.seq_aux:
+                scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
+                ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device, dtype=torch.float32)
+                ce.scatter_add_(
+                    1,
+                    topk_idx_for_aux_loss,
+                    torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device, dtype=torch.float32)
+                )
+                ce.div_(seq_len * aux_topk / self.n_routed_experts)
+                aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
+            else:
+                mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
+                ce = mask_ce.float().mean(0)
+                Pi = scores_for_aux.mean(0)
+                fi = ce * self.n_routed_experts
+                aux_loss = (Pi * fi).sum() * self.alpha
+        else:
+            aux_loss = None
         return topk_idx, topk_weight.to(hidden_states.dtype), aux_loss
 
 
@@ -314,11 +337,11 @@ class DeepseekMoE(nn.Module):
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
         flat_topk_idx = topk_idx.view(-1)
         if self.training:
-            y = self.moe_train(hidden_states, flat_topk_idx, topk_weight.view(-1, 1))
+            y = self.moe_train(hidden_states, flat_topk_idx, topk_weight) # removed unnecessary .view(-1, 1)
             y = y.view(*orig_shape)
             y = AddAuxiliaryLoss.apply(y, aux_loss)
         else:
-            y = self.moe_infer(hidden_states, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
+            y = self.moe_infer(hidden_states, flat_topk_idx, topk_weight).view(*orig_shape) # removed unnecessary .view(-1, 1)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
         return y
@@ -329,7 +352,7 @@ class DeepseekMoE(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)
-        return y
+        return y.to(hidden_states.dtype) # .sum() in previous line returns fp32 tensor
 
     @torch.no_grad()
     def moe_infer(self, x, flat_expert_indices, flat_expert_weights):