update gemma

Author: shihaobai

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

@@ -3,8 +3,10 @@
     MMWeight,
     MultiMMWeight,
     ROWMMWeight,
+    ROWMMWeightNoTP,
     COLMMWeight,
     MultiROWMMWeight,
+    MultiROWMMWeightNoTP,
     CustomMMWeight,
     CustomBMMWeight,
 )

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

@@ -7,7 +7,8 @@ class MMWeightTpl(BaseWeightTpl):
     def __init__(self, data_type, split_n_embed):
         super().__init__()
         self.data_type_ = data_type
-        self.split_n_embed = split_n_embed
+        self.start = split_n_embed * self.tp_rank_
+        self.end = split_n_embed * (self.tp_rank_ + 1)
         self.quant_method = None
         self.weight = None
         self.bias = None
@@ -58,32 +59,35 @@ def __init__(self, weight_name, data_type, split_n_embed, bias_name=None):
         super().__init__(weight_name, data_type, split_n_embed, bias_name)
 
     def load_hf_weights(self, weights):
-        start = self.split_n_embed * self.tp_rank_
-        end = self.split_n_embed * (self.tp_rank_ + 1)
         weight = None
         if self.weight_name in weights:
             weight = weights[self.weight_name].to(self.data_type_)
-            self.weight = weight[start:end]
+            self.weight = weight[self.start : self.end]
         if self.bias_name in weights:
-            bias = weights[self.bias_name].to(self.data_type_)[start:end]
+            bias = weights[self.bias_name].to(self.data_type_)[self.start : self.end]
             self.bias = bias.cuda(self.tp_rank_)
         if weight is None:
             return
         self._post_load_weights()
         return
 
 
+class ROWMMWeightNoTP(MMWeight):
+    def __init__(self, weight_name, data_type, split_n_embed, bias_name=None):
+        super().__init__(weight_name, data_type, split_n_embed, bias_name)
+        self.start = 0
+        self.end = split_n_embed
+
+
 class COLMMWeight(MMWeight):
     def __init__(self, weight_name, data_type, split_n_embed, bias_name=None):
         super().__init__(weight_name, data_type, split_n_embed, bias_name)
 
     def load_hf_weights(self, weights):
-        start = self.split_n_embed * self.tp_rank_
-        end = self.split_n_embed * (self.tp_rank_ + 1)
         weight = None
         if self.weight_name in weights:
             weight = weights[self.weight_name].to(self.data_type_)
-            self.weight = weight[:, start:end]
+            self.weight = weight[:, self.start : self.end]
         if self.bias_name in weights:
             bias = weights[self.bias_name].to(self.data_type_)
             self.bias = (bias / self.world_size_).cuda(self.tp_rank_)
@@ -126,20 +130,25 @@ def _fuse(self):
         return self
 
     def load_hf_weights(self, weights):
-        start = self.split_n_embed * self.tp_rank_
-        end = self.split_n_embed * (self.tp_rank_ + 1)
         weight = None
         for i in range(len(self.weight_names)):
             if self.weight_names[i] in weights:
                 weight = weights[self.weight_names[i]].to(self.data_type_)
-                self.weights[i] = weight[start:end]
+                self.weights[i] = weight[self.start : self.end]
             if self.has_bias and self.bias_names[i] in weights:
                 bias = weights[self.bias_names[i]].to(self.data_type_)
-                self.biases[i] = bias[start:end]
+                self.biases[i] = bias[self.start : self.end]
         self._fuse()
         return
 
 
+class MultiROWMMWeightNoTP(MultiROWMMWeight):
+    def __init__(self, weight_names, data_type, split_n_embed, bias_names=None):
+        super().__init__(weight_names, data_type, split_n_embed, bias_names)
+        self.start = 0
+        self.end = split_n_embed
+
+
 class CustomMMWeight(ROWMMWeight):
     def __init__(
         self,

lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py

@@ -96,21 +96,25 @@ def __init__(
         disable_qk_absorb=False,
         disable_vo_absorb=False,
     ):
+        self.disable_qk_absorb = disable_qk_absorb
+        self.disable_vo_absorb = disable_vo_absorb
         super().__init__(layer_num, tp_rank, world_size, data_type, network_config, mode, quant_cfg)
+        return
+
+    def _parse_config(self):
+        super()._parse_config()
         self.is_moe = (
             self.network_config_["n_routed_experts"] is not None
             and self.layer_num_ >= self.network_config_["first_k_dense_replace"]
             and self.layer_num_ % self.network_config_["moe_layer_freq"] == 0
         )
-        self.tp_q_head_num_ = network_config["num_attention_heads"] // self.world_size_
+        self.tp_q_head_num_ = self.network_config_["num_attention_heads"] // self.world_size_
         self.n_routed_experts = self.network_config_["n_routed_experts"]
         self.q_lora_rank = self.network_config_["q_lora_rank"]
         self.qk_nope_head_dim = self.network_config_["qk_nope_head_dim"]
         self.qk_rope_head_dim = self.network_config_["qk_rope_head_dim"]
         self.num_attention_heads = self.network_config_["num_attention_heads"]
         self.kv_lora_rank = self.network_config_["kv_lora_rank"]
-        self.disable_qk_absorb = disable_qk_absorb
-        self.disable_vo_absorb = disable_vo_absorb
         self.fuse_pairs = {}
         if not self.disable_qk_absorb:
             if self.q_lora_rank is None:
@@ -125,16 +129,15 @@ def __init__(
             }
         )
 
-        self.init_qkvo()
+    def _init_weight(self):
+        self._init_qkvo()
         if self.is_moe:
-            self.init_moe()
+            self._init_moe()
         else:
-            self.init_ffn()
-        self.init_norm()
-        self.set_quantization()
-        return
+            self._init_ffn()
+        self._init_norm()
 
-    def init_qkvo(self):
+    def _init_qkvo(self):
         q_split_n_embed = self.qk_nope_head_dim * self.tp_q_head_num_
         q_split_n_embed_with_rope = (
             (self.qk_nope_head_dim + self.qk_rope_head_dim) * self.num_attention_heads // self.world_size_
@@ -201,7 +204,7 @@ def _load_mlp(self, mlp_prefix, 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):
+    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
@@ -220,12 +223,12 @@ def init_moe(self):
             data_type=self.data_type_,
         )
 
-    def init_ffn(self):
+    def _init_ffn(self):
         inter_size = self.network_config_["intermediate_size"]
         split_inter_size = inter_size // self.world_size_
         self._load_mlp(f"model.layers.{self.layer_num_}.mlp", split_inter_size)
 
-    def init_norm(self):
+    def _init_norm(self):
         self.att_norm_weight_ = NormWeight(f"model.layers.{self.layer_num_}.input_layernorm.weight", self.data_type_)
         self.ffn_norm_weight_ = NormWeight(
             f"model.layers.{self.layer_num_}.post_attention_layernorm.weight", self.data_type_

lightllm/models/deepseek2/model.py

@@ -22,6 +22,8 @@ class Deepseek2TpPartModel(LlamaTpPartModel):
 
     def __init__(self, kvargs):
         super().__init__(kvargs)
+        self.disable_qk_absorb = int(os.getenv("DISABLE_QK_ABSORB", 0))
+        self.disable_vo_absorb = int(os.getenv("DISABLE_VO_ABSORB", 0))
         return
 
     def _init_some_value(self):

lightllm/models/gemma_2b/layer_weights/transformer_layer_weight.py

@@ -2,7 +2,7 @@
 import math
 import numpy as np
 from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
-from lightllm.common.basemodel.layer_weights.meta_weights import GEMMANormWeight, ROWMMWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import GEMMANormWeight, ROWMMWeight, MultiROWMMWeightNoTP
 
 
 class Gemma_2bTransformerLayerWeight(LlamaTransformerLayerWeight):
@@ -14,21 +14,11 @@ def _init_qkv(self):
         q_split_n_embed = self.head_dim * self.n_head // self.world_size_
         kv_split_n_embed = self.head_dim * self.n_kv_head
         self.q_proj = ROWMMWeight(self._q_weight_name, self.data_type_, q_split_n_embed, bias_name=self._q_bias_name)
-        self.k_proj = ROWMMWeight(
-            self._k_weight_name,
+        self.kv_proj = MultiROWMMWeightNoTP(
+            [self._k_weight_name, self._v_weight_name],
             self.data_type_,
             kv_split_n_embed,
-            bias_name=self._k_bias_name,
-            wait_fuse=True,
-            disable_tp=True,
-        )
-        self.v_proj = ROWMMWeight(
-            self._v_weight_name,
-            self.data_type_,
-            kv_split_n_embed,
-            bias_name=self._v_bias_name,
-            wait_fuse=True,
-            disable_tp=True,
+            bias_names=[self._k_bias_name, self._v_bias_name],
         )
 
     def _init_norm(self):

lightllm/models/gemma_2b/model.py

@@ -5,6 +5,7 @@
 from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
 from lightllm.models.llama.model import LlamaTpPartModel
+from lightllm.common.mem_utils import select_mem_manager_class
 
 
 class Gemma_2bTpPartModel(LlamaTpPartModel):
@@ -33,3 +34,14 @@ def _verify_params(self):
         # assert self.config["num_key_value_heads"] % self.world_size_ == 0
         assert self.config["num_attention_heads"] % self.world_size_ == 0
         return
+
+    def _init_mem_manager(self):
+        self.mem_manager = select_mem_manager_class(self.mode)(
+            self.max_total_token_num,
+            dtype=self.data_type,
+            head_num=self.config["num_key_value_heads"],
+            head_dim=self.config["hidden_size"] // self.config["num_attention_heads"],
+            layer_num=self.config["num_hidden_layers"],
+            mem_fraction=self.mem_fraction,
+        )
+        return

lightllm/server/api_cli.py

@@ -234,16 +234,4 @@ def make_argument_parser() -> argparse.ArgumentParser:
         help="""Path of quantization config. It can be used for mixed quantization.
             Examples can be found in lightllm/common/quantization/configs.""",
     )
-    parser.add_argument(
-        "--disable_qk_absorb",
-        default=False,
-        action="store_true",
-        help="Disable mla qk weight absorption",
-    )
-    parser.add_argument(
-        "--disable_vo_absorb",
-        default=False,
-        action="store_true",
-        help="Disable mla vo weight absorption",
-    )
     return parser

lightllm/server/router/manager.py

@@ -133,8 +133,6 @@ async def wait_to_model_ready(self):
                 "batch_max_tokens": self.args.batch_max_tokens,
                 "quant_type": self.args.quant_type,
                 "quant_cfg": self.args.quant_cfg,
-                "disable_qk_absorb": self.args.disable_qk_absorb,
-                "disable_vo_absorb": self.args.disable_vo_absorb,
                 "pd_rpyc_port": self.args.pd_tp_infer_rpyc_ports[rank_id],  # 非 pd 模式可以不设置
             }
             init_model_ret.append(self.model_rpcs[rank_id].init_model(kvargs))

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -109,8 +109,6 @@ def init_model(self, kvargs):
             "batch_max_tokens": kvargs.get("batch_max_tokens", None),
             "quant_type": kvargs.get("quant_type", None),
             "quant_cfg": kvargs.get("quant_cfg", None),
-            "disable_qk_absorb": kvargs.get("disable_qk_absorb", False),
-            "disable_vo_absorb": kvargs.get("disable_vo_absorb", False),
             "run_mode": self.run_mode,
         }