Upload modeling_deepseek.py

modeling_deepseek.py

@@ -27,6 +27,7 @@ import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.library import custom_op
 
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
@@ -469,11 +470,17 @@ class MoEGate(nn.Module):
 
         return topk_idx, topk_weight
 
-class DeepseekV3MoE(nn.Module):
-    """
-    A mixed expert module containing shared experts.
-    """
+@torch.library.custom_op("deepseek::moe_infer_op", mutates_args=())
+def moe_infer_fake(x: torch.Tensor, gate_proj_weight: torch.Tensor, up_proj_weight: torch.Tensor, down_proj_weight: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor) -> torch.Tensor:
+    final_out = torch.empty_like(x)
+    return final_out
 
+# FakeTensor 커널 등록
+@moe_infer_fake.register_fake
+def _(x, gate_proj_weight, up_proj_weight, down_proj_weight, topk_ids, topk_weight):
+    return torch.empty_like(x)
+
+class DeepseekV3MoE(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.config = config
@@ -487,9 +494,7 @@ class DeepseekV3MoE(nn.Module):
             self.experts = nn.ModuleList(
                 [
                     (
-                        DeepseekV3MLP(
-                            config, intermediate_size=config.moe_intermediate_size
-                        )
+                        DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size)
                         if i >= self.ep_rank * self.experts_per_rank
                         and i < (self.ep_rank + 1) * self.experts_per_rank
                         else None
@@ -503,18 +508,14 @@ class DeepseekV3MoE(nn.Module):
             self.ep_rank = 0
             self.experts = nn.ModuleList(
                 [
-                    DeepseekV3MLP(
-                        config, intermediate_size=config.moe_intermediate_size
-                    )
+                    DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size)
                     for i in range(config.n_routed_experts)
                 ]
             )
         self.gate = MoEGate(config)
         if config.n_shared_experts is not None:
             intermediate_size = config.moe_intermediate_size * config.n_shared_experts
-            self.shared_experts = DeepseekV3MLP(
-                config=config, intermediate_size=intermediate_size
-            )
+            self.shared_experts = DeepseekV3MLP(config=config, intermediate_size=intermediate_size)
 
     def forward(self, hidden_states):
         identity = hidden_states
@@ -530,79 +531,29 @@ class DeepseekV3MoE(nn.Module):
 
     @torch.no_grad()
     def moe_infer(self, x, topk_ids, topk_weight):
-        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
-        cnts.scatter_(1, topk_ids, 1)
-        tokens_per_expert = cnts.sum(dim=0)
-        idxs = topk_ids.view(-1).argsort()
-        sorted_tokens = x[idxs // topk_ids.shape[1]]
-        sorted_tokens_shape = sorted_tokens.shape
-        if self.ep_size > 1:
-            tokens_per_ep_rank = tokens_per_expert.view(self.ep_size, -1).sum(dim=1)
-            tokens_per_expert_group = tokens_per_expert.new_empty(
-                tokens_per_expert.shape[0]
-            )
-            dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert)
-            output_splits = (
-                tokens_per_expert_group.view(self.ep_size, -1)
-                .sum(1)
-                .cpu()
-                .numpy()
-                .tolist()
-            )
-            gathered_tokens = sorted_tokens.new_empty(
-                tokens_per_expert_group.sum(dim=0).cpu().item(), sorted_tokens.shape[1]
-            )
-            input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist()
-            dist.all_to_all(
-                list(gathered_tokens.split(output_splits)),
-                list(sorted_tokens.split(input_split_sizes)),
-            )
-            tokens_per_expert_post_gather = tokens_per_expert_group.view(
-                self.ep_size, self.experts_per_rank
-            ).sum(dim=0)
-            gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32)
-            s = 0
-            for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):
-                gatherd_idxs[s : s + k] = i % self.experts_per_rank
-                s += k
-            gatherd_idxs = gatherd_idxs.argsort()
-            sorted_tokens = gathered_tokens[gatherd_idxs]
-            tokens_per_expert = tokens_per_expert_post_gather
-        tokens_per_expert = tokens_per_expert.cpu().numpy()
-
-        outputs = []
-        start_idx = 0
-        for i, num_tokens in enumerate(tokens_per_expert):
-            end_idx = start_idx + num_tokens
-            if num_tokens == 0:
-                continue
-            expert = self.experts[i + self.ep_rank * self.experts_per_rank]
-            tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
-            expert_out = expert(tokens_for_this_expert)
-            outputs.append(expert_out)
-            start_idx = end_idx
-
-        outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
-        if self.ep_size > 1:
-            new_x = torch.empty_like(outs)
-            new_x[gatherd_idxs] = outs
-            gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
-            dist.all_to_all(
-                list(gathered_tokens.split(input_split_sizes)),
-                list(new_x.split(output_splits)),
-            )
-            outs = gathered_tokens
-
-        new_x = torch.empty_like(outs)
-        new_x[idxs] = outs
-        final_out = (
-            new_x.view(*topk_ids.shape, -1)
-            .type(topk_weight.dtype)
-            .mul_(topk_weight.unsqueeze(dim=-1))
-            .sum(dim=1)
-            .type(new_x.dtype)
+        # self.experts MLP모듈별 weight 추출
+        gate_proj_weight = []
+        up_proj_weight = []
+        down_proj_weight = []
+        for i in range(len(self.experts)):
+            expert = self.experts[i]
+            if expert is not None:
+                gate_proj_weight.append(expert.gate_proj.weight.unsqueeze(0))
+                up_proj_weight.append(expert.up_proj.weight.unsqueeze(0))
+                down_proj_weight.append(expert.down_proj.weight.unsqueeze(0))
+            
+        gate_proj_weight = torch.cat(gate_proj_weight, dim=0)  # [num_experts, hidden_size, intermediate_size]
+        up_proj_weight = torch.cat(up_proj_weight, dim=0)      # [num_experts, hidden_size, intermediate_size] 
+        down_proj_weight = torch.cat(down_proj_weight, dim=0)  # [num_experts, intermediate_size, hidden_size]
+            
+        return moe_infer_fake(
+            x=x,
+            gate_proj_weight=gate_proj_weight,
+            up_proj_weight=up_proj_weight,
+            down_proj_weight=down_proj_weight,
+            topk_ids=topk_ids,
+            topk_weight=topk_weight
         )
-        return final_out
 
 
 # Copied from transformers.models.llama.modeling_llama.repeat_kv