[feature]moe etp done, without group greed

Author: charlotteroes

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight.py

@@ -3,6 +3,7 @@
 from lightllm.utils.dist_utils import get_world_size, get_rank
 import threading
 from lightllm.common.quantization import vLLMFP8w8a8QuantizationMethod
+import os
 
 try:
     HAS_VLLM = True
@@ -28,6 +29,8 @@ def __init__(
         self.tp_rank_ = get_rank()
         self.experts_up_projs = [None] * self.n_routed_experts
         self.experts_gate_projs = [None] * self.n_routed_experts
+        self.expert_gate_up_proj_etp = None
+        self.expert_down_proj_etp  = None
         self.w2_list = [None] * self.n_routed_experts
         self.quant_method = None
         self.lock = threading.Lock()
@@ -36,9 +39,10 @@ def set_quant_method(self, quant_method):
         if isinstance(quant_method, vLLMFP8w8a8QuantizationMethod):
             self.quant_method = quant_method
             if self.quant_method is not None:
-                self.quant_method.is_moe = True
+                self.quant_method.is_moe = True     
 
     def experts(self, input_tensor, router_logits, top_k, renormalize, use_grouped_topk, topk_group, num_expert_group):
+
         topk_weights, topk_ids = FusedMoE.select_experts(
             hidden_states=input_tensor,
             router_logits=router_logits,
@@ -95,27 +99,90 @@ def _fuse(self):
                 delattr(self, "experts_up_projs")
                 delattr(self, "experts_gate_projs")
 
+
+    def _load_hf_weights_etp(self, weights):
+        world_size_ = get_world_size()
+        assert self.n_routed_experts % world_size_ == 0
+        n_expert_ep = self.n_routed_experts // world_size_
+
+        #tp to ep here 
+        expert_gate_up_proj_last = None
+        expert_down_proj_last = None
+        
+        for i_experts_ep in range(n_expert_ep):
+            expert_up_proj = None
+            expert_gate_proj = None
+            expert_gate_up_proj = None
+            expert_down_proj = None
+            i_experts = i_experts_ep + n_expert_ep*self.tp_rank_
+
+            if f"{self.weight_prefix}.{i_experts}.up_proj.weight" in weights:
+                expert_up_proj = weights[f"{self.weight_prefix}.{i_experts}.up_proj.weight"]
+                
+                #self.experts_up_proj[i_experts] = expert_up_proj
+
+            if f"{self.weight_prefix}.{i_experts}.gate_proj.weight" in weights:
+                expert_gate_proj = weights[f"{self.weight_prefix}.{i_experts}.gate_proj.weight"]
+                #self.experts_gate_proj[i_experts] = expert_gate_proj
+
+            if expert_gate_proj is not None and expert_up_proj is not None:
+                expert_gate_up_proj = torch.cat([expert_gate_proj, expert_up_proj], dim=0)
+                self.experts_gate_projs[i_experts_ep] = expert_gate_up_proj #self._cuda(expert_gate_up_proj)
+                expert_gate_up_proj_last = expert_gate_up_proj
+                
+            if f"{self.weight_prefix}.{i_experts}.down_proj.weight" in weights:
+                expert_down_proj = weights[f"{self.weight_prefix}.{i_experts}.down_proj.weight"]
+                self.experts_up_projs[i_experts_ep] = expert_down_proj #self._cuda(expert_down_proj)
+                expert_down_proj_last = expert_down_proj
+
+            with self.lock:
+                if expert_gate_up_proj_last is not None:
+                    #package, if there is broken experts
+
+                    if self.expert_gate_up_proj_etp  is None:
+                        self.expert_gate_up_proj_etp = torch.zeros( (n_expert_ep,) + expert_gate_up_proj_last.shape,
+                            dtype=expert_gate_up_proj_last.dtype).cuda(self.tp_rank_)
+                    
+                    for i_experts_ep in range(n_expert_ep):
+                        if self.experts_gate_projs[i_experts_ep] is not None:
+                            self.expert_gate_up_proj_etp[i_experts_ep,:] = self.experts_gate_projs[i_experts_ep]
+
+
+                if expert_down_proj_last is not None:
+                    #package, if there is broken experts
+                    if self.expert_down_proj_etp  is None:
+                        self.expert_down_proj_etp = torch.zeros( (n_expert_ep,) + expert_down_proj_last.shape,
+                            dtype=expert_down_proj_last.dtype).cuda(self.tp_rank_)
+        
+                    for i_experts_ep in range(n_expert_ep):
+                        if self.experts_up_projs[i_experts_ep] is not None:
+                            self.expert_down_proj_etp[i_experts_ep,:] = self.experts_up_projs[i_experts_ep]
+
+
     def load_hf_weights(self, weights):
-        for i_experts in range(self.n_routed_experts):
-            w1_weight = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.weight"
-            w2_weight = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.weight"
-            w3_weight = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.weight"
-
-            if w1_weight in weights:
-                self.experts_gate_projs[i_experts] = weights[w1_weight][
-                    self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1), :
-                ]
-            if w3_weight in weights:
-                self.experts_up_projs[i_experts] = weights[w3_weight][
-                    self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1), :
-                ]
-
-            if w2_weight in weights:
-                self.w2_list[i_experts] = weights[w2_weight][
-                    :, self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1)
-                ]
-
-        self._fuse()
+        if os.environ.get("ETP_MODE_ENABLED") == "true":
+            self._load_hf_weights_etp(weights)
+        else:
+            for i_experts in range(self.n_routed_experts):
+                w1_weight = f"{self.weight_prefix}.{i_experts}.{self.w1_weight_name}.weight"
+                w2_weight = f"{self.weight_prefix}.{i_experts}.{self.w2_weight_name}.weight"
+                w3_weight = f"{self.weight_prefix}.{i_experts}.{self.w3_weight_name}.weight"
+
+                if w1_weight in weights:
+                    self.experts_gate_projs[i_experts] = weights[w1_weight][
+                        self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1), :
+                    ]
+                if w3_weight in weights:
+                    self.experts_up_projs[i_experts] = weights[w3_weight][
+                        self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1), :
+                    ]
+
+                if w2_weight in weights:
+                    self.w2_list[i_experts] = weights[w2_weight][
+                        :, self.split_inter_size * self.tp_rank_ : self.split_inter_size * (self.tp_rank_ + 1)
+                    ]
+
+            self._fuse()
 
     def _cuda(self, cpu_tensor):
         if self.tp_rank_ is None:
@@ -124,4 +191,7 @@ def _cuda(self, cpu_tensor):
             return cpu_tensor.contiguous().to(self.data_type_).cuda(self.tp_rank_)
 
     def verify_load(self):
-        return self.w1 is not None and self.w2 is not None
+        if os.environ.get("ETP_MODE_ENABLED") == "true":
+            return True
+        else:
+            return self.w1 is not None and self.w2 is not None

lightllm/common/deepseek2_mem_manager.py

@@ -1,4 +1,5 @@
 import torch
+import os
 
 from .mem_manager import MemoryManager
 from typing import List
@@ -10,7 +11,12 @@ def get_cell_size(self):
 
     def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
         self.kv_buffer = torch.empty((layer_num, size, head_num, head_dim), dtype=dtype, device="cuda")
-
+        #todo, etp or edp use the same work buffer here
+        #also it can be used for any kernels for work buffer witout save info only
+        if os.environ.get("ETP_MODE_ENABLED") == "true":
+            self.work_buffer = torch.empty(1024*1024*1024,dtype=torch.bfloat16,  device="cuda") 
+            self.work_buffer.share_memory_()
+            
     def alloc_kv_move_buffer(self, max_req_total_len):
         self.kv_move_buffer = torch.empty(
             (1, max_req_total_len + 8, self.head_num, self.head_dim), dtype=self.dtype, device="cuda"

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -18,7 +18,7 @@
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
 from functools import partial
 from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
-
+import os
 
 class Deepseek2TransformerLayerInfer(LlamaTransformerLayerInfer):
     def __init__(
@@ -30,6 +30,9 @@ def __init__(
         self.qk_rope_head_dim = network_config["qk_rope_head_dim"]
         self.q_lora_rank = network_config["q_lora_rank"]
         self.kv_lora_rank = network_config["kv_lora_rank"]
+
+        self.n_routed_experts = network_config["n_routed_experts"]
+
         self.is_moe = (
             network_config["n_routed_experts"] is not None
             and layer_num >= network_config["first_k_dense_replace"]
@@ -64,7 +67,10 @@ def _bind_attention(self):
         )
         self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
         if self.is_moe:
-            self._ffn = partial(Deepseek2TransformerLayerInfer._moe_ffn, self)
+            if os.environ.get("ETP_MODE_ENABLED") == "true":
+                self._ffn = partial(Deepseek2TransformerLayerInfer._moe_ffn_etp, self)
+            else:
+                self._ffn = partial(Deepseek2TransformerLayerInfer._moe_ffn, self)
         else:
             self._ffn = partial(LlamaTransformerLayerInfer._ffn, self)
 
@@ -196,6 +202,7 @@ def _copy_kv_to_mem_cache_normal(self, buffer, mem_index, mem_manager):
     def _moe_ffn(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
     ) -> torch.Tensor:
+
         hidden_states = input.view(-1, self.embed_dim_)
         num_tokens, hidden_dim = hidden_states.shape
 
@@ -219,3 +226,61 @@ def _moe_ffn(
             hidden_states.add_(shared_output)
 
         return hidden_states.view(num_tokens, hidden_dim)
+
+    def _moe_ffn_etp(
+        self, input, infer_state: LlamaInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
+    ) -> torch.Tensor:
+        world_size_ = self.world_size_
+        num_local_experts = self.n_shared_experts // world_size_
+        local_expert_offset = self.tp_rank_ * num_local_experts
+        num_experts_per_token = self.num_experts_per_tok
+        num_experts = self.n_routed_experts
+        num_expert_groups = self.n_group
+        num_groups_per_token = self.topk_group
+        gating_scaling_factor = self.routed_scaling_factor
+        gating_normalize_prob = self.norm_topk_prob
+        rank_self = self.tp_rank_
+
+        hidden_states = input.view(-1, self.embed_dim_)
+        num_tokens, hidden_dim = hidden_states.shape
+
+        final_hidden_states = torch.empty(num_tokens,hidden_dim,device=hidden_states.device,
+            dtype = hidden_states.dtype )  
+
+        #router_logits_len = hidden_states.shape[0]*layer_weight.moe_gate.shape[1]        
+        router_logits = layer_weight.moe_gate.mm(hidden_states)
+
+        #now some parameter is not supported yet 
+        #assert gating_normalize_prob is False
+        #assert num_expert_groups<=1
+
+
+      
+        import lightllm_moe_etp_kernel
+        lightllm_moe_etp_kernel.moe_fused_all(
+            router_logits.contiguous(),
+            hidden_states.contiguous(),
+            layer_weight.gate_up_proj.weight.contiguous(), #transpose
+            layer_weight.down_proj.weight.contiguous(),    #transpose
+            layer_weight.experts.expert_gate_up_proj_etp.contiguous(),
+            layer_weight.experts.expert_down_proj_etp.contiguous(), 
+            infer_state.mem_manager.work_buffer.contiguous(), 
+            infer_state.mem_manager.work_buffer.nelement(),
+            final_hidden_states.contiguous(), 
+            rank_self,
+            gating_scaling_factor,
+            num_experts,
+            num_experts_per_token,
+            num_tokens,
+            world_size_,
+            True,
+            hidden_dim,
+            layer_weight.gate_up_proj.weight.size(1)//2,
+            layer_weight.experts.expert_gate_up_proj_etp.size(1)//2,
+            self.n_shared_experts is not None
+        )
+
+        router_logits = None
+
+        return final_hidden_states.view(num_tokens, hidden_dim)
+