fix odd vocab size for qwen2 tp>1 (#460)

Co-authored-by: baishihao <baishihao@sensetime.com>
ModelTC · Jul 9, 2024 · c600d14 · c600d14
1 parent 50f0e2e
commit c600d14
Showing 1 changed file with 6 additions and 12 deletions.
diff --git a/lightllm/models/qwen2/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen2/layer_weights/pre_and_post_layer_weight.py
@@ -9,23 +9,17 @@ def __init__(self, tp_rank, world_size, data_type, network_config, mode):
         return
 
     def load_hf_weights(self, weights):
-
         vob_size = self.network_config_["vocab_size"]
-        split_vob_size = vob_size // self.world_size_
-
+        split_indexes = np.linspace(0, vob_size, self.world_size_ + 1, dtype=np.int64)
+        split_start = split_indexes[self.tp_rank_]
+        split_end = split_indexes[self.tp_rank_ + 1]
         if "model.embed_tokens.weight" in weights:
-            self.wte_weight_ = self._cuda(
-                weights["model.embed_tokens.weight"][
-                    split_vob_size * self.tp_rank_ : split_vob_size * (self.tp_rank_ + 1), :
-                ]
-            )
-            tie_word_embeddings = self.network_config_.get('tie_word_embeddings', False)
+            self.wte_weight_ = self._cuda(weights["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_vob_size * self.tp_rank_ : split_vob_size * (self.tp_rank_ + 1), :]
-            )
+            self.lm_head_weight_ = self._cuda(weights["lm_head.weight"][split_start:split_end, :])
         if "model.norm.weight" in weights:
             self.final_norm_weight_ = self._cuda(weights["model.norm.weight"])