add model.combine_moe() to speedup

modeling_deepseekocr.py

@@ -1,6 +1,8 @@
 import os
 import math
 import re
+import gc
+import types
 from tqdm import tqdm
 from abc import ABC
 from typing import List, Optional, Tuple, Union
@@ -15,6 +17,7 @@ from torch.nn import CrossEntropyLoss
 from torchvision import transforms
 
 from transformers.cache_utils import Cache
+from transformers.activations import ACT2FN
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers import DeepseekV2Model, DeepseekV2ForCausalLM
 from transformers import DeepseekV2Config
@@ -1058,3 +1061,38 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
                 plt.close()
 
             result.save(f"{output_path}/result_with_boxes.jpg")
+
+    def combine_moe(self):
+        for layer in self.model.layers:
+            if isinstance(layer.mlp, DeepseekV2MoE):
+                moe_layer = layer.mlp
+                # combine experts
+                moe_layer.w1 = nn.Parameter(torch.stack([moe_layer.experts[i].gate_proj.weight.T for i in range(moe_layer.config.n_routed_experts)]), requires_grad=False)
+                moe_layer.w2 = nn.Parameter(torch.stack([moe_layer.experts[i].down_proj.weight.T for i in range(moe_layer.config.n_routed_experts)]), requires_grad=False)
+                moe_layer.w3 = nn.Parameter(torch.stack([moe_layer.experts[i].up_proj.weight.T for i in range(moe_layer.config.n_routed_experts)]), requires_grad=False)
+                del moe_layer.experts
+                gc.collect()
+                moe_layer.experts = None
+                moe_layer.act = ACT2FN[moe_layer.config.hidden_act]
+                moe_layer.forward = types.MethodType(new_forward_for_moe, moe_layer)
+
+
+def new_forward_for_moe(self, hidden_states):
+    batch_size, sequence_length, hidden_dim = hidden_states.shape
+    selected_experts, routing_weights = self.gate(hidden_states)
+    router_scores = torch.zeros(size=(batch_size * sequence_length, self.config.n_routed_experts), device=hidden_states.device, dtype=hidden_states.dtype)
+    # we cast back to the input dtype
+    routing_weights = routing_weights.to(hidden_states.dtype)
+    router_scores = torch.scatter_add(router_scores, -1, selected_experts, routing_weights)
+    hidden_states = hidden_states.view(-1, hidden_dim)
+    if self.config.n_shared_experts is not None:
+        shared_expert_output = self.shared_experts(hidden_states)
+
+    hidden_w1 = torch.matmul(hidden_states, self.w1)
+    hidden_w3 = torch.matmul(hidden_states, self.w3)
+    hidden_states = self.act(hidden_w1) * hidden_w3
+    hidden_states = torch.bmm(hidden_states, self.w2) * torch.transpose(router_scores, 0, 1).unsqueeze(-1)
+    final_hidden_states = hidden_states.sum(dim=0, dtype=hidden_states.dtype)
+    if self.config.n_shared_experts is not None:
+        hidden_states = final_hidden_states + shared_expert_output
+    return hidden_states.view(batch_size, sequence_length, hidden_dim)