refactor weight

Author: shihaobai

lightllm/models/__init__.py

@@ -29,7 +29,8 @@
 from lightllm.models.internvl.model import InternVLInternlm2TpPartModel
 from lightllm.models.qwen2_vl.model import Qwen2VLTpPartModel
 from lightllm.models.qwen2_reward.model import Qwen2RewardTpPartModel
-from lightllm.models.qwen3_vl.model import Qwen3VLTpPartModel, Qwen3VLMOETpPartModel
+from lightllm.models.qwen3_vl.model import Qwen3VLTpPartModel
+from lightllm.models.qwen3_vl_moe.model import Qwen3VLMOETpPartModel
 from lightllm.models.gemma3.model import Gemma3TpPartModel
 from lightllm.models.tarsier2.model import (
     Tarsier2Qwen2TpPartModel,

lightllm/models/qwen3_vl/__init__.py

@@ -1,36 +1,22 @@
 import numpy as np
-from lightllm.common.basemodel import PreAndPostLayerWeight
+from lightllm.models.qwen2.layer_weights.pre_and_post_layer_weight import Qwen2PreAndPostLayerWeight
 
+# add key: language_model.xxx -> xxx
+# only change keys at PreAndPostLayerWeight load, TransformLayerWeight is correct now
+def rename_weight_keys(weights):
+    prefix = "model.language_model."
+    keys = list(weights.keys())
+    for k in keys:
+        if prefix in k:
+            weights[k.replace(prefix, "model.")] = weights.pop(k)
 
-class Qwen3VLPreAndPostLayerWeight(PreAndPostLayerWeight):
+
+class Qwen3VLPreAndPostLayerWeight(Qwen2PreAndPostLayerWeight):
     def __init__(self, data_type, network_config, mode):
         super().__init__(data_type, network_config, mode)
         return
 
     def load_hf_weights(self, weights):
-        vob_size = self.network_config_["vocab_size"]
-        split_indexes = np.linspace(0, vob_size, self.tp_world_size_ + 1, dtype=np.int64)
-        split_start = split_indexes[self.tp_rank_]
-        split_end = split_indexes[self.tp_rank_ + 1]
-        if "model.language_model.embed_tokens.weight" in weights:
-            self.wte_weight_ = self._cuda(weights["model.language_model.embed_tokens.weight"][split_start:split_end, :])
-            tie_word_embeddings = self.network_config_.get("tie_word_embeddings", False)
-            if tie_word_embeddings:
-                self.lm_head_weight_ = self.wte_weight_
-        if "lm_head.weight" in weights:
-            self.lm_head_weight_ = self._cuda(weights["lm_head.weight"][split_start:split_end, :])
-        if "model.language_model.norm.weight" in weights:
-            self.final_norm_weight_ = self._cuda(weights["model.language_model.norm.weight"])
-
-        return
-
-    def verify_load(self):
-        errors = "weights load not ok"
-        weights = [
-            self.wte_weight_,
-            self.lm_head_weight_,
-            self.final_norm_weight_,
-        ]
-        for i in range(len(weights)):
-            assert weights[i] is not None, "index:" + str(i) + " " + errors
+        rename_weight_keys(weights)
+        super().load_hf_weights(weights)
         return

lightllm/models/qwen3_vl/layer_infer/__init__.py

@@ -1,18 +1,13 @@
 import os
 import json
-import numpy as np
 from lightllm.common.build_utils import repair_config
 from lightllm.models.registry import ModelRegistry
 from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo
 from lightllm.models.qwen3_vl.layer_infer.pre_layer_infer import Qwen3VLMultimodalPreLayerInfer
 from lightllm.models.qwen3_vl.layer_infer.transformer_layer_infer import Qwen3VLTransformerLayerInfer
 from lightllm.models.qwen3_vl.layer_weights.pre_and_post_layer_weight import Qwen3VLPreAndPostLayerWeight
-from lightllm.models.qwen3_vl.layer_weights.transformers_layer_weight import Qwen3VLTransformerLayerWeight
-from lightllm.models.qwen3_vl_moe.layer_weights.transformers_layer_weight import Qwen3VLMOETransformerLayerWeight
-from lightllm.models.qwen3_vl_moe.layer_infer.transformer_layer_infer import Qwen3VLMOETransformerLayerInfer
 from lightllm.models.qwen2_vl.model import QWen2VLTokenizer
 from lightllm.models.qwen3.model import Qwen3TpPartModel
-from lightllm.models.qwen3_moe.model import Qwen3MOEModel
 
 
 class QWen3VLTokenizer(QWen2VLTokenizer):
@@ -35,38 +30,6 @@ class Qwen3VLTpPartModel(Qwen3TpPartModel):
     transformer_layer_infer_class = Qwen3VLTransformerLayerInfer
 
     pre_and_post_weight_class = Qwen3VLPreAndPostLayerWeight
-    transformer_weight_class = Qwen3VLTransformerLayerWeight
-
-    infer_state_class = Qwen3VLInferStateInfo
-
-    def __init__(self, kvargs):
-        super().__init__(kvargs)
-        return
-
-    def _init_inferstate_cls(self):
-        pass
-
-    def _init_config(self):
-        with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file:
-            all_config = json.load(json_file)
-            self.config = all_config["text_config"]
-        # rename keys
-        repair_config(self.config, same_names=["num_attention_heads", "n_head"])
-        repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"])
-        repair_config(self.config, same_names=["num_hidden_layers", "n_layer"])
-        if self.finetune_config:
-            self.config["vocab_size"] = self.finetune_config.vocab_size
-        return
-
-
-@ModelRegistry(["qwen3_vl_moe"], is_multimodal=True)
-class Qwen3VLMOETpPartModel(Qwen3MOEModel):
-
-    pre_layer_infer_class = Qwen3VLMultimodalPreLayerInfer
-    transformer_layer_infer_class = Qwen3VLMOETransformerLayerInfer
-
-    pre_and_post_weight_class = Qwen3VLPreAndPostLayerWeight
-    transformer_weight_class = Qwen3VLMOETransformerLayerWeight
 
     infer_state_class = Qwen3VLInferStateInfo