[fixed] fix pre-commit

Author: charlotteroes

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

@@ -30,7 +30,7 @@ def __init__(
         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.expert_down_proj_etp = None
         self.w2_list = [None] * self.n_routed_experts
         self.quant_method = None
         self.lock = threading.Lock()
@@ -39,7 +39,7 @@ 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):
 
@@ -99,65 +99,64 @@ 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 
+        # 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_
+            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
+
+                # 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
+                # 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)
+                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)
+                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
+                    # 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_)
 
-                    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]
-
+                            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_)
-        
+                    # 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]
-
+                            self.expert_down_proj_etp[i_experts_ep, :] = self.experts_up_projs[i_experts_ep]
 
     def load_hf_weights(self, weights):
         if os.environ.get("ETP_MODE_ENABLED") == "true":

lightllm/common/deepseek2_mem_manager.py

@@ -11,12 +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
+        # 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 = 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

@@ -403,42 +403,42 @@ 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_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
+        # 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
+        # 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 )  
+        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_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
-
+        # 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.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(), 
+            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(), 
+            final_hidden_states.contiguous(),
             rank_self,
             gating_scaling_factor,
             num_experts,
@@ -447,12 +447,11 @@ def _moe_ffn_etp(
             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
+            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)
-