add v2 gready

modeling_deepseek.py

@@ -398,7 +398,8 @@ class MoEGate(nn.Module):
         self.n_routed_experts = config.n_routed_experts
         self.routed_scaling_factor = config.routed_scaling_factor
         self.scoring_func = config.scoring_func
-        self.topk_method = config.topk_method
+        # self.topk_method = config.topk_method
+        self.topk_method == "gready"
         self.n_group = config.n_group
         self.topk_group = config.topk_group
 
@@ -459,6 +460,14 @@ class MoEGate(nn.Module):
                 tmp_scores, k=self.top_k, dim=-1, sorted=False
             )
             topk_weight = scores.gather(1, topk_idx)
+        elif self.topk_method == "gready":
+            """
+            impl from deepseek v2
+            https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/main/modeling_deepseek.py#L435
+            """
+            topk_weight, topk_idx = torch.topk(
+                scores, k=self.top_k, dim=-1, sorted=False
+            )            
         else:
             raise NotImplementedError(
                 f"insupportable TopK function for MoE gating: {self.topk_method}"
@@ -528,11 +537,31 @@ class DeepseekV3MoE(nn.Module):
         if not self.training:
             y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         else:
-            y = self.moe_train(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+            # y = self.moe_train(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+            y = self.moe_train_v2(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
         return y
 
+    def moe_train_v2(self, hidden_states, topk_idx, topk_weight):
+        """
+        impl from deepseek v2
+        https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/main/modeling_deepseek.py#L566
+        """
+        flat_topk_idx = topk_idx.view(-1)
+        
+        hidden_states = hidden_states.repeat_interleave(
+                self.num_experts_per_tok, dim=0
+        )
+        y = torch.empty_like(hidden_states)
+        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.type(hidden_states.dtype)
+        # y = y.view(*orig_shape)
+        # y = AddAuxiliaryLoss.apply(y, aux_loss)
+        return y
+    
     def moe_train(self, x, topk_ids, topk_weight):
         """
         self.ep_size = 1の想定