v3-fix

Author: shihaobai

lightllm/common/basemodel/basemodel.py

@@ -99,7 +99,8 @@ def _init_config(self):
         repair_config(self.config, same_names=["num_hidden_layers", "n_layer"])
         if self.finetune_config:
             self.config["vocab_size"] = self.finetune_config.vocab_size
-
+        # self.config["num_hidden_layers"] = 4
+        # self.config["n_layer"] = 4
         return
 
     @final

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -154,7 +154,7 @@ def _decompress_kv(self, kv, infer_state: Deepseek2InferStateInfo, layer_weight:
             )
 
         # CC
-        compressed_kv = compressed_kv.view(-1, layer_weight.kv_lora_rank)
+        compressed_kv = compressed_kv.view(-1, layer_weight.kv_lora_rank).contiguous()
         k_nope = self.alloc_tensor(
             [compressed_kv.shape[0], self.tp_q_head_num_, self.qk_nope_head_dim],
             dtype=compressed_kv.dtype,
@@ -163,10 +163,8 @@ def _decompress_kv(self, kv, infer_state: Deepseek2InferStateInfo, layer_weight:
             k_nope.shape,
             dtype=compressed_kv.dtype,
         )
-        wk = layer_weight.k_b_proj_.weight.view(-1, layer_weight.kv_lora_rank).T
-        wv = layer_weight.v_b_proj_.weight.transpose(0, 1).reshape(layer_weight.kv_lora_rank, -1)
-        torch.mm(compressed_kv, wk, out=k_nope.reshape(compressed_kv.shape[0], -1))
-        torch.mm(compressed_kv, wv, out=v.reshape(compressed_kv.shape[0], -1))
+        layer_weight.cc_k_b_proj_.mm(compressed_kv, out=k_nope.reshape(compressed_kv.shape[0], -1))
+        layer_weight.cc_v_b_proj_.mm(compressed_kv, out=v.reshape(compressed_kv.shape[0], -1))
         return k_nope, k_rope, v
 
     def _context_attention_kernel_with_CC(

lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py

@@ -160,6 +160,13 @@ def load_hf_weights(self, weights):
             weights[f"model.layers.{self.layer_num_}.self_attn.k_b_proj.weight"] = self._load_kb(kv_b_proj_)
             weights[f"model.layers.{self.layer_num_}.self_attn.v_b_proj.weight"] = self._load_vb(kv_b_proj_)
 
+            weights[f"model.layers.{self.layer_num_}.self_attn.cc_k_b_proj.weight"] = self._load_kb(kv_b_proj_).reshape(
+                -1, self.kv_lora_rank
+            )
+            weights[f"model.layers.{self.layer_num_}.self_attn.cc_v_b_proj.weight"] = (
+                self._load_vb(kv_b_proj_).transpose(0, 1).reshape(self.kv_lora_rank, -1).transpose(0, 1).contiguous()
+            )
+
         if (
             self.quant_cfg.quantized_weight
             and f"model.layers.{self.layer_num_}.self_attn.kv_b_proj." + self.weight_scale_suffix in weights
@@ -178,6 +185,17 @@ def load_hf_weights(self, weights):
                 f"model.layers.{self.layer_num_}.self_attn.v_b_proj." + self.weight_scale_suffix
             ] = self._load_vb_scale(kv_b_proj_scale_, block_size)
 
+            weights[
+                f"model.layers.{self.layer_num_}.self_attn.cc_k_b_proj." + self.weight_scale_suffix
+            ] = self._load_kb_scale(kv_b_proj_scale_, block_size).reshape(-1, self.kv_lora_rank // block_size)
+            weights[f"model.layers.{self.layer_num_}.self_attn.cc_v_b_proj." + self.weight_scale_suffix] = (
+                self._load_vb_scale(kv_b_proj_scale_, block_size)
+                .transpose(0, 1)
+                .reshape(self.kv_lora_rank // block_size, -1)
+                .transpose(0, 1)
+                .contiguous()
+            )
+
         return super().load_hf_weights(weights)
 
     def _set_quantization(self):
@@ -237,6 +255,21 @@ def _init_qkvo(self):
             act_scale_suffix=self.act_scale_suffix,
         )
 
+        self.cc_k_b_proj_ = ROWMMWeight(
+            f"model.layers.{self.layer_num_}.self_attn.cc_k_b_proj.weight",
+            self.data_type_,
+            split_n_embed=self.tp_q_head_num_ * self.qk_nope_head_dim,
+            weight_scale_suffix=self.weight_scale_suffix,
+            act_scale_suffix=self.act_scale_suffix,
+        )
+        self.cc_v_b_proj_ = ROWMMWeight(
+            f"model.layers.{self.layer_num_}.self_attn.cc_v_b_proj.weight",
+            self.data_type_,
+            split_n_embed=self.tp_q_head_num_ * self.qk_nope_head_dim,
+            weight_scale_suffix=self.weight_scale_suffix,
+            act_scale_suffix=self.act_scale_suffix,
+        )
+
         self.o_weight_ = COLMMWeight(
             f"model.layers.{self.layer_num_}.self_attn.o_proj.weight",
             self.data_type_,