qwen weight repeat for tp_size > kv_head_num

Author: shihaobai

lightllm/models/qwen2/layer_weights/transformer_layer_weight.py

@@ -13,3 +13,42 @@ def _init_weight_names(self):
         self._q_bias_name = f"model.layers.{self.layer_num_}.self_attn.q_proj.bias"
         self._k_bias_name = f"model.layers.{self.layer_num_}.self_attn.k_proj.bias"
         self._v_bias_name = f"model.layers.{self.layer_num_}.self_attn.v_proj.bias"
+
+    def _parse_config(self):
+        self.tp_q_head_num_ = self.network_config_["num_attention_heads"] // self.tp_world_size_
+        self.tp_k_head_num_ = max(self.network_config_["num_key_value_heads"] // self.tp_world_size_, 1)
+        self.tp_v_head_num_ = self.tp_k_head_num_
+        self.tp_o_head_num_ = self.tp_q_head_num_
+        head_dim = self.network_config_["hidden_size"] // self.network_config_["num_attention_heads"]
+        self.head_dim = self.network_config_.get("head_dim", head_dim)
+        assert self.tp_k_head_num_ * self.tp_world_size_ % self.network_config_["num_key_value_heads"] == 0
+
+    def _repeat_weight(self, name, weights):
+        # for tp_world_size_ > num_key_value_heads
+        if name not in weights:
+            return
+
+        tensor = weights[name]
+        num_kv_heads = self.network_config_["num_key_value_heads"]
+        repeat_size = self.tp_k_head_num_ * self.tp_world_size_ // num_kv_heads
+
+        if tensor.ndim == 1:
+            # Bias (1D tensor)
+            tensor = tensor.reshape(num_kv_heads, -1).unsqueeze(1).repeat(1, repeat_size, 1).reshape(-1)
+        else:
+            # Weight (2D tensor)
+            tensor = (
+                tensor.reshape(num_kv_heads, -1, tensor.shape[-1])
+                .unsqueeze(1)
+                .repeat(1, repeat_size, 1, 1)
+                .reshape(-1, tensor.shape[-1])
+            )
+        weights[name] = tensor
+
+    def load_hf_weights(self, weights):
+        self._repeat_weight(self._k_weight_name, weights)
+        self._repeat_weight(self._v_weight_name, weights)
+        if self._k_bias_name is not None and self._v_bias_name is not None:
+            self._repeat_weight(self._k_bias_name, weights)
+            self._repeat_weight(self._v_bias_name, weights)
+        return super().load_hf_weights(weights)

lightllm/models/qwen2/model.py

@@ -2,6 +2,7 @@
 from lightllm.models.qwen2.layer_weights.pre_and_post_layer_weight import Qwen2PreAndPostLayerWeight
 from lightllm.models.qwen2.layer_weights.transformer_layer_weight import Qwen2TransformerLayerWeight
 from lightllm.models.llama.model import LlamaTpPartModel
+from lightllm.common.mem_utils import select_mem_manager_class
 
 
 @ModelRegistry("qwen2")
@@ -22,7 +23,30 @@ def _init_config(self):
         return
 
     def _verify_params(self):
-        assert self.load_way in ["HF"], "mistral only supports HF format to load Now!"
-        assert self.config["num_key_value_heads"] % self.tp_world_size_ == 0
+        assert self.load_way in ["HF", "DS"], "llama only supports HF and DS format to load Now!"
         assert self.config["num_attention_heads"] % self.tp_world_size_ == 0
         return
+
+    def _init_some_value(self):
+        # Dealing with head_dim_!=n_embed // num_attention_heads scenarios, such as mistral 13B
+        head_dim_ = self.config["n_embed"] // self.config["num_attention_heads"]
+        self.head_dim_ = self.config.get("head_dim", head_dim_)
+        self.tp_k_head_num_ = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1)
+        self.tp_v_head_num_ = self.tp_k_head_num_
+        self.layers_num = self.config["n_layer"]
+        self.vocab_size = self.config["vocab_size"]
+        return
+
+    def _init_mem_manager(self):
+        head_dim_ = self.config["hidden_size"] // self.config["num_attention_heads"]
+        head_dim_ = self.config.get("head_dim", head_dim_)
+        tp_k_head_num_ = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1)
+        self.mem_manager = select_mem_manager_class(self.mode)(
+            self.max_total_token_num,
+            dtype=self.data_type,
+            head_num=tp_k_head_num_,
+            head_dim=head_dim_,
+            layer_num=self.config["num_hidden_layers"],
+            mem_fraction=self.mem_fraction,
+        )
+        return

lightllm/models/qwen3/layer_weights/transformer_layer_weight.py

@@ -2,9 +2,8 @@
 import torch
 import math
 import numpy as np
-from lightllm.common.basemodel import TransformerLayerWeight
 from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
-from lightllm.utils.envs_utils import enable_env_vars
+from lightllm.models.qwen2.layer_weights.transformer_layer_weight import Qwen2TransformerLayerWeight
 from lightllm.common.basemodel.layer_weights.meta_weights import (
     ROWMMWeight,
     MultiROWMMWeight,
@@ -17,7 +16,7 @@
 from functools import partial
 
 
-class Qwen3TransformerLayerWeight(LlamaTransformerLayerWeight):
+class Qwen3TransformerLayerWeight(Qwen2TransformerLayerWeight):
     def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
         super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
         return
@@ -26,6 +25,9 @@ def _init_weight_names(self):
         super()._init_weight_names()
         self._q_norm_name = f"model.layers.{self.layer_num_}.self_attn.q_norm.weight"
         self._k_norm_name = f"model.layers.{self.layer_num_}.self_attn.k_norm.weight"
+        self._q_bias_name = None
+        self._k_bias_name = None
+        self._v_bias_name = None
 
     def _init_norm(self):
         super()._init_norm()

lightllm/models/qwen3/model.py

@@ -3,16 +3,15 @@
 from lightllm.models.registry import ModelRegistry
 from lightllm.models.qwen3.layer_infer.transformer_layer_infer import Qwen3TransformerLayerInfer
 from lightllm.models.qwen3.layer_weights.transformer_layer_weight import Qwen3TransformerLayerWeight
-from lightllm.models.llama.model import LlamaTpPartModel
+from lightllm.models.qwen2.model import Qwen2TpPartModel
 from lightllm.utils.log_utils import init_logger
-from lightllm.common.mem_utils import select_mem_manager_class
 
 
 logger = init_logger(__name__)
 
 
 @ModelRegistry("qwen3")
-class Qwen3TpPartModel(LlamaTpPartModel):
+class Qwen3TpPartModel(Qwen2TpPartModel):
     # weight class
     transformer_weight_class = Qwen3TransformerLayerWeight
 
@@ -22,17 +21,3 @@ class Qwen3TpPartModel(LlamaTpPartModel):
     def __init__(self, kvargs):
         super().__init__(kvargs)
         return
-
-    def _init_mem_manager(self):
-        head_dim_ = self.config["hidden_size"] // self.config["num_attention_heads"]
-        head_dim_ = self.config.get("head_dim", head_dim_)
-        tp_k_head_num_ = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1)
-        self.mem_manager = select_mem_manager_class(self.mode)(
-            self.max_total_token_num,
-            dtype=self.data_type,
-            head_num=tp_k_head_num_,
-            head_dim=head_dim_,
-            layer_num=self.config["num_hidden_layers"],
-            mem_fraction=self.mem_fraction,
-        )
-        return

lightllm/models/qwen3_moe/layer_weights/transformer_layer_weight.py

@@ -3,7 +3,7 @@
 import math
 import numpy as np
 from lightllm.common.basemodel import TransformerLayerWeight
-from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
+from lightllm.models.qwen3.layer_weights.transformer_layer_weight import Qwen3TransformerLayerWeight
 from lightllm.utils.envs_utils import enable_env_vars
 from lightllm.common.basemodel.layer_weights.meta_weights import (
     ROWMMWeight,
@@ -17,7 +17,7 @@
 from functools import partial
 
 
-class Qwen3MOETransformerLayerWeight(LlamaTransformerLayerWeight):
+class Qwen3MOETransformerLayerWeight(Qwen3TransformerLayerWeight):
     def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
         self.n_routed_experts = network_config["num_experts"]
         self.is_moe = (
@@ -46,36 +46,15 @@ def _init_weight_names(self):
         self._ffn_norm_weight_name = f"model.layers.{self.layer_num_}.post_attention_layernorm.weight"
         self._ffn_norm_bias_name = None
 
-    def _parse_config(self):
-        self.tp_q_head_num_ = self.network_config_["num_attention_heads"] // self.tp_world_size_
-        self.tp_k_head_num_ = max(self.network_config_["num_key_value_heads"] // self.tp_world_size_, 1)
-        self.tp_v_head_num_ = self.tp_k_head_num_
-        self.tp_o_head_num_ = self.tp_q_head_num_
-        self.head_dim = self.network_config_["head_dim"]
-        assert self.tp_k_head_num_ * self.tp_world_size_ % self.network_config_["num_key_value_heads"] == 0
-
-    def _repeat_weight(self, name, weights):
-        repeat_size = self.tp_k_head_num_ * self.tp_world_size_ // self.network_config_["num_key_value_heads"]
-        repeat_params = (1, repeat_size, 1, 1)
-        if name in weights:
-            weights[name] = (
-                weights[name]
-                .reshape(self.network_config_["num_key_value_heads"], -1, weights[name].shape[1])
-                .unsqueeze(1)
-                .repeat(repeat_params)
-                .reshape(-1, weights[name].shape[1])
-            )
-
     def load_hf_weights(self, weights):
-        self._repeat_weight(self._k_weight_name, weights)
-        self._repeat_weight(self._v_weight_name, weights)
+        super().load_hf_weights(weights)
         kv_b_quant_method = self.quant_cfg.get_quant_method(self.layer_num_, "kv_b_proj")
         if self.quant_cfg.quantized_weight:
             _k_scale_weight_name = self._k_weight_name.replace("weight", kv_b_quant_method.weight_scale_suffix)
             self._repeat_weight(_k_scale_weight_name, weights)
             _v_scale_weight_name = self._v_weight_name.replace("weight", kv_b_quant_method.weight_scale_suffix)
             self._repeat_weight(_v_scale_weight_name, weights)
-        return super().load_hf_weights(weights)
+        return
 
     def _init_weight(self):
         self._init_qkv()
@@ -127,8 +106,3 @@ def _init_moe(self):
             )
         else:
             raise ValueError(f"Unsupported moe mode: {moe_mode}")
-
-    def _init_norm(self):
-        super()._init_norm()
-        self.q_norm_weight_ = NormWeight(weight_name=self._q_norm_name, data_type=self.data_type_)
-        self.k_norm_weight_ = NormWeight(weight_name=self._k_norm_name, data_type=self.data_type_)

lightllm/models/qwen3_moe/model.py

@@ -3,16 +3,15 @@
 from lightllm.models.registry import ModelRegistry
 from lightllm.models.qwen3_moe.layer_infer.transformer_layer_infer import Qwen3MOETransformerLayerInfer
 from lightllm.models.qwen3_moe.layer_weights.transformer_layer_weight import Qwen3MOETransformerLayerWeight
-from lightllm.models.llama.model import LlamaTpPartModel
-from lightllm.common.mem_utils import select_mem_manager_class
+from lightllm.models.qwen3.model import Qwen3TpPartModel
 from lightllm.utils.log_utils import init_logger
 
 
 logger = init_logger(__name__)
 
 
 @ModelRegistry("qwen3_moe")
-class Qwen3MOEModel(LlamaTpPartModel):
+class Qwen3MOEModel(Qwen3TpPartModel):
     # weight class
     transformer_weight_class = Qwen3MOETransformerLayerWeight
 
@@ -22,32 +21,3 @@ class Qwen3MOEModel(LlamaTpPartModel):
     def __init__(self, kvargs):
         super().__init__(kvargs)
         return
-
-    def _verify_params(self):
-        assert self.load_way in ["HF", "DS"], "llama only supports HF and DS format to load Now!"
-        assert self.config["num_attention_heads"] % self.tp_world_size_ == 0
-        return
-
-    def _init_some_value(self):
-        # Dealing with head_dim_!=n_embed // num_attention_heads scenarios, such as mistral 13B
-        head_dim_ = self.config["n_embed"] // self.config["num_attention_heads"]
-        self.head_dim_ = self.config.get("head_dim", head_dim_)
-        self.tp_k_head_num_ = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1)
-        self.tp_v_head_num_ = self.tp_k_head_num_
-        self.layers_num = self.config["n_layer"]
-        self.vocab_size = self.config["vocab_size"]
-        return
-
-    def _init_mem_manager(self):
-        head_dim_ = self.config["hidden_size"] // self.config["num_attention_heads"]
-        head_dim_ = self.config.get("head_dim", head_dim_)
-        tp_k_head_num_ = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1)
-        self.mem_manager = select_mem_manager_class(self.mode)(
-            self.max_total_token_num,
-            dtype=self.data_type,
-            head_num=tp_k_head_num_,
-            head_dim=head_dim_,
-            layer_num=self.config["num_hidden_layers"],
-            mem_fraction=self.mem_fraction,
-        )
-        return