solve format

Author: shihaobai

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

@@ -7,7 +7,9 @@
 
 
 class FusedMoeWeight(BaseWeight):
-    def __init__(self, gate_proj_name, down_proj_name, up_proj_name, weight_prefix, n_routed_experts, split_inter_size, data_type):
+    def __init__(
+        self, gate_proj_name, down_proj_name, up_proj_name, weight_prefix, n_routed_experts, split_inter_size, data_type
+    ):
         super().__init__()
         self.w1_weight_name = gate_proj_name
         self.w2_weight_name = down_proj_name
@@ -22,53 +24,47 @@ def __init__(self, gate_proj_name, down_proj_name, up_proj_name, weight_prefix,
         self.w2_list = [None] * self.n_routed_experts
         self.quant_method = None
         self.lock = threading.Lock()
-    
+
     def set_quant_method(self, quant_method):
         self.quant_method = quant_method
         if self.quant_method is not None:
             self.quant_method.is_moe = True
 
-    def experts(
-            self,
-            input_tensor,
-            router_logits,
-            top_k,
-            renormalize,
-            use_grouped_topk,
-            topk_group,
-            num_expert_group
-        ):
+    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,
             use_grouped_topk=use_grouped_topk,
             top_k=top_k,
             renormalize=renormalize,
             topk_group=topk_group,
-            num_expert_group=num_expert_group
+            num_expert_group=num_expert_group,
         )
         if self.quant_method is not None:
-            fused_experts(input_tensor,
-                         w1=self.w1[0],
-                         w2=self.w2[0],
-                         topk_weights=topk_weights,
-                         topk_ids=topk_ids,
-                         inplace=False,
-                         use_fp8_w8a8=True,
-                         use_int8_w8a16=False,
-                         w1_scale=self.w1[1],
-                         w2_scale=self.w2[1],
-                         a1_scale=None,
-                         a2_scale=None)
+            fused_experts(
+                input_tensor,
+                w1=self.w1[0],
+                w2=self.w2[0],
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                inplace=False,
+                use_fp8_w8a8=True,
+                use_int8_w8a16=False,
+                w1_scale=self.w1[1],
+                w2_scale=self.w2[1],
+                a1_scale=None,
+                a2_scale=None,
+            )
             return
-        fused_experts(hidden_states=input_tensor,
+        fused_experts(
+            hidden_states=input_tensor,
             w1=self.w1,
             w2=self.w2,
             topk_weights=topk_weights,
             topk_ids=topk_ids,
-            inplace=True
+            inplace=True,
         )
-        
+
     def fuse(self):
         with self.lock:
             if (
@@ -120,15 +116,14 @@ def load_hf_weights(self, 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:
             return cpu_tensor.contiguous().to(self.data_type_).cuda()
         else:
             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

lightllm/common/basemodel/layer_weights/transformer_layer_weight.py

@@ -2,7 +2,7 @@
 
 # from lightllm.common.layers.mm import MM
 from .base_layer_weight import BaseLayerWeight
-from .meta_weights import MMWeight,FusedMoeWeight
+from .meta_weights import MMWeight, FusedMoeWeight
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)

lightllm/common/quantization/vllm_quant.py

@@ -64,10 +64,9 @@ def quantize(self, weight: torch.Tensor):
             return self.quantize_moe(weight)
         qweight, weight_scale = ops.scaled_fp8_quant(weight.cuda(), scale=None, use_per_token_if_dynamic=True)
         return qweight.transpose(0, 1), weight_scale
-    
+
     def quantize_moe(self, weight):
         num_experts = weight.shape[0]
-        out_dim = weight.shape[1]
         qweights = []
         weight_scales = []
         qweights = torch.empty_like(weight, dtype=torch.float8_e4m3fn).cuda()

lightllm/models/baichuan7b/layer_weights/transformer_layer_weight.py

@@ -9,7 +9,7 @@ class BaiChuan7bTransformerLayerWeight(LlamaTransformerLayerWeight):
     def __init__(self, layer_num, tp_rank, world_size, data_type, network_config, mode=[], quant_cfg=None):
         super().__init__(layer_num, tp_rank, world_size, data_type, network_config, mode, quant_cfg)
         return
-      
+
     def _init_config(self):
         self.network_config_["num_key_value_heads"] = self.network_config_["num_attention_heads"]
         self.n_embed = self.network_config_["hidden_size"]

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -9,6 +9,7 @@
     context_attention_fwd,
     context_attention_fwd_no_prompt_cache,
 )
+
 from lightllm.models.deepseek2.triton_kernel.flash_decoding import token_decode_attention_flash_decoding
 from lightllm.models.deepseek2.layer_infer.fused_moe import fused_experts, grouped_topk
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
@@ -20,7 +21,9 @@
 
 
 class Deepseek2TransformerLayerInfer(LlamaTransformerLayerInfer):
-    def __init__(self, layer_num, tp_rank, world_size, network_config, mode=[], disable_qk_absorb=False, disable_vo_absorb=False):
+    def __init__(
+        self, layer_num, tp_rank, world_size, network_config, mode=[], disable_qk_absorb=False, disable_vo_absorb=False
+    ):
         self.tp_k_head_num_ = 1
         self.tp_v_head_num_ = 1
         self.qk_nope_head_dim = network_config["qk_nope_head_dim"]
@@ -207,7 +210,7 @@ def _moe_ffn(
             renormalize=self.norm_topk_prob,
             use_grouped_topk=self.n_group,
             topk_group=self.topk_group,
-            num_expert_group=self.n_group
+            num_expert_group=self.n_group,
         )
 
         hidden_states.mul_(self.routed_scaling_factor)

lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py

@@ -2,7 +2,14 @@
 import math
 import numpy as np
 from lightllm.common.basemodel import TransformerLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, COLMMWeight, NormWeight, CustomMMWeight, FusedMoeWeight, CustomBMMWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import (
+    ROWMMWeight,
+    COLMMWeight,
+    NormWeight,
+    CustomMMWeight,
+    FusedMoeWeight,
+    CustomBMMWeight,
+)
 from functools import partial
 
 
@@ -19,16 +26,32 @@ def fuse_q_kb(self, A, B):
     k_nope_proj_ = k_b_proj_.unsqueeze(0)
     k_nope_proj_ = k_nope_proj_.to(torch.float64)
 
-    return self._cuda(torch.matmul(q_nope_proj_, k_nope_proj_).view(-1, self.tp_q_head_num_ * self.kv_lora_rank).transpose(0, 1))
+    return self._cuda(
+        torch.matmul(q_nope_proj_, k_nope_proj_).view(-1, self.tp_q_head_num_ * self.kv_lora_rank).transpose(0, 1)
+    )
+
 
 def fuse_vb_o(self, A, B):
     v_b_proj_ = A.weight
-    o_weight_ = B.weight.transpose(0, 1).view(self.tp_q_head_num_, self.qk_nope_head_dim, -1).contiguous().to(self.data_type_).cpu()
-    return self._cuda(torch.matmul(v_b_proj_.to(torch.float64), o_weight_.to(torch.float64)).view(-1, self.network_config_["hidden_size"]))
+    o_weight_ = (
+        B.weight.transpose(0, 1)
+        .view(self.tp_q_head_num_, self.qk_nope_head_dim, -1)
+        .contiguous()
+        .to(self.data_type_)
+        .cpu()
+    )
+    return self._cuda(
+        torch.matmul(v_b_proj_.to(torch.float64), o_weight_.to(torch.float64)).view(
+            -1, self.network_config_["hidden_size"]
+        )
+    )
+
 
 def load_q_rope(self, A, q_weight_):
     q_split_n_embed_with_rope = A.split_n_embed
-    q_weight_ = q_weight_[q_split_n_embed_with_rope * self.tp_rank_ : q_split_n_embed_with_rope * (self.tp_rank_ + 1), :]
+    q_weight_ = q_weight_[
+        q_split_n_embed_with_rope * self.tp_rank_ : q_split_n_embed_with_rope * (self.tp_rank_ + 1), :
+    ]
     q_weight_ = q_weight_.transpose(0, 1).contiguous()
     q_nope_proj_, q_rope_proj_ = torch.split(
         q_weight_.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim),
@@ -37,6 +60,7 @@ def load_q_rope(self, A, q_weight_):
     )
     return self._cuda(q_rope_proj_.reshape(-1, self.qk_rope_head_dim * self.tp_q_head_num_).transpose(0, 1))
 
+
 def load_kb(self, A, kv_b_proj_):
     kv_b_proj_ = kv_b_proj_
     k_b_proj_ = kv_b_proj_.view(self.num_attention_heads, self.qk_nope_head_dim * 2, self.kv_lora_rank)[
@@ -47,22 +71,31 @@ def load_kb(self, A, kv_b_proj_):
         return k_b_proj_.contiguous().to(self.data_type_).cpu()
     return self._cuda(k_b_proj_)
 
+
 def load_vb(self, A, kv_b_proj_):
     kv_b_proj_ = kv_b_proj_
-    v_b_proj_ = kv_b_proj_.T.view(
-        self.kv_lora_rank,
-        self.num_attention_heads,
-        self.qk_nope_head_dim * 2,
-    )[:, :, self.qk_nope_head_dim :].transpose(0, 1)
-    v_b_proj_ = v_b_proj_[
-        self.tp_q_head_num_ * self.tp_rank_ : self.tp_q_head_num_ * (self.tp_rank_ + 1), :, :
-    ]
+    v_b_proj_ = kv_b_proj_.T.view(self.kv_lora_rank, self.num_attention_heads, self.qk_nope_head_dim * 2,)[
+        :, :, self.qk_nope_head_dim :
+    ].transpose(0, 1)
+    v_b_proj_ = v_b_proj_[self.tp_q_head_num_ * self.tp_rank_ : self.tp_q_head_num_ * (self.tp_rank_ + 1), :, :]
     if A.wait_fuse:
         return v_b_proj_.contiguous().to(self.data_type_).cpu()
     return self._cuda(v_b_proj_)
 
+
 class Deepseek2TransformerLayerWeight(TransformerLayerWeight):
-    def __init__(self, layer_num, tp_rank, world_size, data_type, network_config, mode=[], quant_cfg=None, disable_qk_absorb=False, disable_vo_absorb=False):
+    def __init__(
+        self,
+        layer_num,
+        tp_rank,
+        world_size,
+        data_type,
+        network_config,
+        mode=[],
+        quant_cfg=None,
+        disable_qk_absorb=False,
+        disable_vo_absorb=False,
+    ):
         super().__init__(layer_num, tp_rank, world_size, data_type, network_config, mode, quant_cfg)
         self.is_moe = (
             self.network_config_["n_routed_experts"] is not None
@@ -86,9 +119,11 @@ def __init__(self, layer_num, tp_rank, world_size, data_type, network_config, mo
                 self.fuse_pairs = {"q_b_proj_&k_b_proj_": "fuse_qk_weight_"}
         if not self.disable_vo_absorb:
             self.fuse_pairs["v_b_proj_&o_weight_"] = "fuse_vo_weight_"
-        self.fuse_pairs.update({
-            "gate_proj&up_proj": "gate_up_proj",
-        })
+        self.fuse_pairs.update(
+            {
+                "gate_proj&up_proj": "gate_up_proj",
+            }
+        )
 
         self.init_qkvo()
         if self.is_moe:
@@ -115,7 +150,10 @@ def init_qkvo(self):
             rope_weight_name = f"model.layers.{self.layer_num_}.self_attn.q_proj.weight"
         else:
             self.q_a_proj_ = ROWMMWeight(
-                f"model.layers.{self.layer_num_}.self_attn.q_a_proj.weight", self.data_type_, self.q_lora_rank, disable_tp=True
+                f"model.layers.{self.layer_num_}.self_attn.q_a_proj.weight",
+                self.data_type_,
+                self.q_lora_rank,
+                disable_tp=True,
             )
             self.q_b_proj_ = CustomMMWeight(
                 f"model.layers.{self.layer_num_}.self_attn.q_b_proj.weight",
@@ -126,10 +164,7 @@ def init_qkvo(self):
             )
             rope_weight_name = f"model.layers.{self.layer_num_}.self_attn.q_b_proj.weight"
         self.q_rope_proj_ = CustomMMWeight(
-            rope_weight_name,
-            self.data_type_,
-            q_split_n_embed_with_rope,
-            custom_load=partial(load_q_rope, self)
+            rope_weight_name, self.data_type_, q_split_n_embed_with_rope, custom_load=partial(load_q_rope, self)
         )
         self.kv_a_proj_with_mqa_ = ROWMMWeight(
             f"model.layers.{self.layer_num_}.self_attn.kv_a_proj_with_mqa.weight",
@@ -142,50 +177,47 @@ def init_qkvo(self):
             self.data_type_,
             None,
             wait_fuse=not self.disable_qk_absorb,
-            custom_load=partial(load_kb, self)
+            custom_load=partial(load_kb, self),
         )
         self.v_b_proj_ = CustomBMMWeight(
             f"model.layers.{self.layer_num_}.self_attn.kv_b_proj.weight",
             self.data_type_,
             None,
             wait_fuse=not self.disable_vo_absorb,
             custom_load=partial(load_vb, self),
-            custom_fuse=partial(fuse_vb_o, self)
+            custom_fuse=partial(fuse_vb_o, self),
         )
         self.o_weight_ = COLMMWeight(
-            f"model.layers.{self.layer_num_}.self_attn.o_proj.weight", self.data_type_, q_split_n_embed, wait_fuse=not self.disable_vo_absorb,
+            f"model.layers.{self.layer_num_}.self_attn.o_proj.weight",
+            self.data_type_,
+            q_split_n_embed,
+            wait_fuse=not self.disable_vo_absorb,
         )
 
     def _load_mlp(self, mlp_prefix, split_inter_size):
         self.gate_proj = ROWMMWeight(
             f"{mlp_prefix}.gate_proj.weight", self.data_type_, split_inter_size, wait_fuse=True
         )
-        self.up_proj = ROWMMWeight(
-            f"{mlp_prefix}.up_proj.weight", self.data_type_, split_inter_size, wait_fuse=True
-        )
-        self.down_proj = COLMMWeight(
-            f"{mlp_prefix}.down_proj.weight", self.data_type_, split_inter_size
-        )
+        self.up_proj = ROWMMWeight(f"{mlp_prefix}.up_proj.weight", self.data_type_, split_inter_size, wait_fuse=True)
+        self.down_proj = COLMMWeight(f"{mlp_prefix}.down_proj.weight", self.data_type_, split_inter_size)
 
     def init_moe(self):
         moe_intermediate_size = self.network_config_["moe_intermediate_size"]
         self.moe_gate = ROWMMWeight(
             f"model.layers.{self.layer_num_}.mlp.gate.weight", self.data_type_, moe_intermediate_size, disable_tp=True
         )
-        shared_intermediate_size = (
-            moe_intermediate_size * self.network_config_["n_shared_experts"]
-        )
+        shared_intermediate_size = moe_intermediate_size * self.network_config_["n_shared_experts"]
         shared_split_inter_size = shared_intermediate_size // self.world_size_
         self._load_mlp(f"model.layers.{self.layer_num_}.mlp.shared_experts", shared_split_inter_size)
-        
+
         self.experts = FusedMoeWeight(
             gate_proj_name="gate_proj",
             down_proj_name="down_proj",
             up_proj_name="up_proj",
             weight_prefix=f"model.layers.{self.layer_num_}.mlp.experts",
             n_routed_experts=self.n_routed_experts,
             split_inter_size=moe_intermediate_size // self.world_size_,
-            data_type=self.data_type_
+            data_type=self.data_type_,
         )
 
     def init_ffn(self):

lightllm/models/deepseek2/model.py

@@ -66,7 +66,7 @@ def _init_weights(self):
                 mode=self.mode,
                 quant_cfg=self.quant_cfg,
                 disable_qk_absorb=self.disable_qk_absorb,
-                disable_vo_absorb=self.disable_vo_absorb
+                disable_vo_absorb=self.disable_vo_absorb,
             )
             for i in range(self.config["n_layer"])
         ]
@@ -80,7 +80,7 @@ def _init_weights(self):
         self.pre_post_weight.verify_load()
         [weight.verify_load() for weight in self.trans_layers_weight]
         return
-    
+
     def _init_infer_layer(self):
         self.pre_infer = self.pre_layer_infer_class(
             tp_rank=self.tp_rank_, world_size=self.world_size_, network_config=self.config, mode=self.mode
@@ -90,7 +90,13 @@ def _init_infer_layer(self):
         )
         self.layers_infer = [
             self.transformer_layer_infer_class(
-                i, tp_rank=self.tp_rank_, world_size=self.world_size_, network_config=self.config, mode=self.mode, disable_qk_absorb=self.disable_qk_absorb, disable_vo_absorb=self.disable_vo_absorb
+                i,
+                tp_rank=self.tp_rank_,
+                world_size=self.world_size_,
+                network_config=self.config,
+                mode=self.mode,
+                disable_qk_absorb=self.disable_qk_absorb,
+                disable_vo_absorb=self.disable_vo_absorb,
             )
             for i in range(self.config["n_layer"])
         ]