Adaptation of InternLM to GQA (#243)

Co-authored-by: baishihao <[email protected]>
ModelTC · Dec 7, 2023 · a9cf015 · a9cf015
1 parent d55c561
commit a9cf015
Showing 1 changed file with 16 additions and 15 deletions.
diff --git a/lightllm/models/internlm/layer_weights/transformer_layer_weight.py b/lightllm/models/internlm/layer_weights/transformer_layer_weight.py
@@ -45,36 +45,37 @@ def _load_qkvo_weights(self, weights):
             self.att_norm_weight_ = self._cuda(weights[f"model.layers.{self.layer_num_}.input_layernorm.weight"])
 
         n_embed = self.network_config_["hidden_size"]
-        split_n_embed = n_embed // self.world_size_
+        q_split_n_embed = n_embed // self.world_size_
+        kv_split_n_embed = n_embed // self.network_config_["num_attention_heads"] * self.network_config_["num_key_value_heads"] // self.world_size_
         # q k v weights for llama
         if f"model.layers.{self.layer_num_}.self_attn.q_proj.weight" in weights:
-            self.q_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.q_proj.weight"][split_n_embed *
-                                                                                                self.tp_rank_: split_n_embed * (self.tp_rank_ + 1), :]
+            self.q_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.q_proj.weight"]
+            self.q_weight_ = self.q_weight_[q_split_n_embed * self.tp_rank_: q_split_n_embed * (self.tp_rank_ + 1), :]
             self.q_weight_ = self._cuda(self.q_weight_.transpose(0, 1))
         if f"model.layers.{self.layer_num_}.self_attn.q_proj.bias" in weights:
-            self.q_bias_ = weights[f"model.layers.{self.layer_num_}.self_attn.q_proj.bias"][split_n_embed *
-                                                                                                self.tp_rank_: split_n_embed * (self.tp_rank_ + 1)]
+            self.q_bias_ = weights[f"model.layers.{self.layer_num_}.self_attn.q_proj.bias"][q_split_n_embed *
+                                                                                                self.tp_rank_: q_split_n_embed * (self.tp_rank_ + 1)]
             self.q_bias_ = self._cuda(self.q_bias_)
         if f"model.layers.{self.layer_num_}.self_attn.k_proj.weight" in weights:
-            self.k_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.k_proj.weight"][split_n_embed *
-                                                                                                self.tp_rank_: split_n_embed * (self.tp_rank_ + 1), :]
+            self.k_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.k_proj.weight"]
+            self.k_weight_ = self.k_weight_[kv_split_n_embed * self.tp_rank_: kv_split_n_embed * (self.tp_rank_ + 1), :]
             self.k_weight_ = self._cuda(self.k_weight_.transpose(0, 1))
         if f"model.layers.{self.layer_num_}.self_attn.k_proj.bias" in weights:
-            self.k_bias_ = weights[f"model.layers.{self.layer_num_}.self_attn.k_proj.bias"][split_n_embed *
-                                                                                                self.tp_rank_: split_n_embed * (self.tp_rank_ + 1)]
+            self.k_bias_ = weights[f"model.layers.{self.layer_num_}.self_attn.k_proj.bias"][kv_split_n_embed *
+                                                                                                self.tp_rank_: kv_split_n_embed * (self.tp_rank_ + 1)]
             self.k_bias_ = self._cuda(self.k_bias_)
         if f"model.layers.{self.layer_num_}.self_attn.v_proj.weight" in weights:
-            self.v_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.v_proj.weight"][split_n_embed *
-                                                                                                self.tp_rank_: split_n_embed * (self.tp_rank_ + 1), :]
+            self.v_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.v_proj.weight"]
+            self.v_weight_ = self.v_weight_[kv_split_n_embed * self.tp_rank_: kv_split_n_embed * (self.tp_rank_ + 1), :]
             self.v_weight_ = self._cuda(self.v_weight_.transpose(0, 1))
         if f"model.layers.{self.layer_num_}.self_attn.v_proj.bias" in weights:
-            self.v_bias_ = weights[f"model.layers.{self.layer_num_}.self_attn.v_proj.bias"][split_n_embed *
-                                                                                                self.tp_rank_: split_n_embed * (self.tp_rank_ + 1)]
+            self.v_bias_ = weights[f"model.layers.{self.layer_num_}.self_attn.v_proj.bias"][kv_split_n_embed *
+                                                                                                self.tp_rank_: kv_split_n_embed * (self.tp_rank_ + 1)]
             self.v_bias_ = self._cuda(self.v_bias_)        
         # attention output dense params
         if f"model.layers.{self.layer_num_}.self_attn.o_proj.weight" in weights:
-            self.o_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.o_proj.weight"][:,
-                                                                                                            split_n_embed * self.tp_rank_: split_n_embed * (self.tp_rank_ + 1)]
+            self.o_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.o_proj.weight"]
+            self.o_weight_ = self.o_weight_[:, q_split_n_embed * self.tp_rank_: q_split_n_embed * (self.tp_rank_ + 1)]
             self.o_weight_ = self._cuda(self.o_weight_.transpose(0, 1))
         if f"model.layers.{self.layer_num_}.self_attn.o_proj.bias" in weights:
             self.o_bias_ = weights[f"model.layers.{self.layer_num_}.self_attn.o_proj.bias"]