Qwen3MOE for tp8 (repeat kv) and qwen3 dense fix (#877)

Co-authored-by: baishihao <baishihao@sensetime.com>
Co-authored-by: wangzaijun <wzjhelloworld@qq.com>

Author: shihaobai

lightllm/models/qwen3/layer_weights/transformer_layer_weight.py

@@ -19,12 +19,6 @@
 
 class Qwen3TransformerLayerWeight(LlamaTransformerLayerWeight):
     def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        self.n_routed_experts = network_config["num_experts"]
-        self.is_moe = (
-            network_config["num_experts"] > 0
-            and layer_num not in network_config["mlp_only_layers"]
-            and (layer_num + 1) % network_config["decoder_sparse_step"] == 0
-        )
         super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
         return
 

lightllm/models/qwen3/model.py

@@ -4,6 +4,7 @@
 from lightllm.models.qwen3.layer_weights.transformer_layer_weight import Qwen3TransformerLayerWeight
 from lightllm.models.llama.model import LlamaTpPartModel
 from lightllm.utils.log_utils import init_logger
+from lightllm.common.mem_utils import select_mem_manager_class
 
 
 logger = init_logger(__name__)
@@ -19,3 +20,17 @@ 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_infer/transformer_layer_infer.py

@@ -29,6 +29,8 @@ def __init__(self, layer_num, network_config, mode=[]):
         self.norm_topk_prob = network_config["norm_topk_prob"]
         super().__init__(layer_num, network_config, mode)
         self.head_dim_ = network_config["head_dim"]
+        self.tp_k_head_num_ = max(self.tp_k_head_num_, 1)
+        self.tp_v_head_num_ = max(self.tp_v_head_num_, 1)
         return
 
     def _bind_func(self):

lightllm/models/qwen3_moe/layer_weights/transformer_layer_weight.py

@@ -46,6 +46,31 @@ 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"], self.head_dim, -1)
+                .unsqueeze(1)
+                .repeat(repeat_params)
+                .reshape(self.network_config_["num_key_value_heads"] * self.head_dim * repeat_size, -1)
+            )
+
+    def load_hf_weights(self, weights):
+        self._repeat_weight(self._k_weight_name, weights)
+        self._repeat_weight(self._v_weight_name, weights)
+        return super().load_hf_weights(weights)
+
     def _init_weight(self):
         self._init_qkv()
         self._init_o()
@@ -99,6 +124,5 @@ def _init_moe(self):
 
     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,6 +3,7 @@
 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.utils.log_utils import init_logger
 
 
@@ -19,3 +20,22 @@ 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_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/server/api_models.py

@@ -76,6 +76,7 @@ class ChatCompletionRequest(BaseModel):
     ignore_eos: Optional[bool] = False
     role_settings: Optional[Dict[str, str]] = None
     character_settings: Optional[List[Dict[str, str]]] = None
+    chat_template_kwargs: Optional[Dict[str, bool]] = None
 
 
 class FunctionResponse(BaseModel):

lightllm/server/build_prompt.py

@@ -16,6 +16,10 @@ async def build_prompt(request, tools) -> str:
         kwargs["character_settings"] = request.character_settings
     if request.role_settings:
         kwargs["role_setting"] = request.role_settings
+
+    if request.chat_template_kwargs:
+        kwargs.update(request.chat_template_kwargs)
+
     try:
         input_str = tokenizer.apply_chat_template(**kwargs, tokenize=False, add_generation_prompt=True, tools=tools)
     except: