add support for MiniCPM (#338)

Author: shihaobai

README.md

@@ -47,6 +47,7 @@ LightLLM is a Python-based LLM (Large Language Model) inference and serving fram
 - [Llava-13b](https://huggingface.co/liuhaotian/llava-v1.5-13b)  
 - [Mixtral]()
 - [Stablelm](https://huggingface.co/stabilityai/stablelm-2-1_6b)
+- [MiniCPM](https://huggingface.co/openbmb/MiniCPM-2B-sft-bf16)
 
 > When you start Qwen-7b, you need to set the parameter '--eos_id 151643 --trust_remote_code'.
 

lightllm/models/minicpm/__init__.py

@@ -0,0 +1,48 @@
+import torch
+import torch.functional as F
+import torch.distributed as dist
+import numpy as np
+
+from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd
+from lightllm.models.internlm.layer_weights.transformer_layer_weight import InternlmTransformerLayerWeight
+from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+
+
+class InternlmTransformerLayerInfer(LlamaTransformerLayerInfer):
+    def __init__(self, layer_num, tp_rank, world_size, network_config, mode=[]):
+        super().__init__(layer_num, tp_rank, world_size, network_config, mode)
+        return
+
+    def _get_qkv(
+        self, input, cache_kv, infer_state: LlamaInferStateInfo, layer_weight: InternlmTransformerLayerWeight
+    ) -> torch.Tensor:
+        q = torch.addmm(
+            layer_weight.q_bias_, input.view(-1, self.embed_dim_), layer_weight.q_weight_, beta=1.0, alpha=1.0
+        )
+        torch.addmm(
+            layer_weight.kv_bias_,
+            input.view(-1, self.embed_dim_),
+            layer_weight.kv_weight_,
+            beta=1.0,
+            alpha=1.0,
+            out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_),
+        )
+        rotary_emb_fwd(
+            q.view(-1, self.tp_q_head_num_, self.head_dim_),
+            cache_kv[:, 0 : self.tp_k_head_num_, :],
+            infer_state.position_cos,
+            infer_state.position_sin,
+        )
+        return q, cache_kv
+
+    def _get_o(
+        self, input, infer_state: LlamaInferStateInfo, layer_weight: InternlmTransformerLayerWeight
+    ) -> torch.Tensor:
+        o_tensor = torch.addmm(
+            layer_weight.o_bias_,
+            input.view(-1, self.tp_o_head_num_ * self.head_dim_),
+            layer_weight.o_weight_,
+            beta=1.0 / self.world_size_,
+        )
+        return o_tensor

lightllm/models/minicpm/layer_infer/__init__.py

@@ -0,0 +1,39 @@
+import torch
+import numpy as np
+from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
+
+
+class MiniCPMPreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
+    def __init__(self, tp_rank, world_size, data_type, network_config, mode):
+        super().__init__(tp_rank, world_size, data_type, network_config, mode)
+        hidden_size = self.network_config_["hidden_size"]
+        dim_model_base = self.network_config_.get("dim_model_base", hidden_size)
+        self.lm_head_scale = hidden_size / dim_model_base
+        self.scale_emb = self.network_config_.get("scale_emb", 1)
+        return
+
+    def load_hf_weights(self, weights):
+        vob_size = self.network_config_["vocab_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:
+            # print(weights['model.embed_tokens.weight'].shape)
+            self.wte_weight_ = self._cuda(weights["model.embed_tokens.weight"][split_start:split_end, :])
+        if "lm_head.weight" in weights:
+            # print(weights['lm_head.weight'].shape)
+            self.lm_head_weight_ = self._cuda(weights["lm_head.weight"][split_start:split_end, :]) / self.lm_head_scale
+        if "model.norm.weight" in weights:
+            self.final_norm_weight_ = self._cuda(weights["model.norm.weight"])
+
+        return
+
+    def verify_load(self):
+        if not hasattr(self, "lm_head_weight_"):
+            self.lm_head_weight_ =  self.wte_weight_ / self.lm_head_scale
+        self.wte_weight_ = self.wte_weight_ * self.scale_emb
+        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
+        return

lightllm/models/minicpm/layer_infer/transformer_layer_infer.py

@@ -0,0 +1,83 @@
+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
+
+
+class MiniCPMTransformerLayerWeight(LlamaTransformerLayerWeight):
+    def __init__(self, layer_num, tp_rank, world_size, data_type, network_config, mode=[]):
+        super().__init__(layer_num, tp_rank, world_size, data_type, network_config, mode)
+        num_hidden_layers = self.network_config_["num_hidden_layers"]
+        scale_depth = self.network_config_.get("scale_depth", math.sqrt(num_hidden_layers))
+        self.layer_scale  =scale_depth /  math.sqrt(num_hidden_layers)
+        return
+
+    def _load_qkvo_weights(self, weights):
+        # input layernorm params
+        if f"model.layers.{self.layer_num_}.input_layernorm.weight" in weights:
+            self.att_norm_weight_ = self._cuda(weights[f"model.layers.{self.layer_num_}.input_layernorm.weight"])
+
+        n_embed = self.network_config_["hidden_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_
+        )
+        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"]
+            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.k_proj.weight" in weights:
+            k_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.k_proj.weight"]
+            k_weight_ = k_weight_[kv_split_n_embed * self.tp_rank_ : kv_split_n_embed * (self.tp_rank_ + 1), :]
+            self.k_weight_ = k_weight_.transpose(0, 1)
+
+        if f"model.layers.{self.layer_num_}.self_attn.v_proj.weight" in weights:
+            v_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.v_proj.weight"]
+            v_weight_ = v_weight_[kv_split_n_embed * self.tp_rank_ : kv_split_n_embed * (self.tp_rank_ + 1), :]
+            self.v_weight_ = v_weight_.transpose(0, 1)
+
+        # 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"]
+            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)) * self.layer_scale
+
+        self._try_cat_to(["k_weight_", "v_weight_"], "kv_weight_", cat_dim=1)
+
+        return
+
+    def _load_ffn_weights(self, weights):
+        if f"model.layers.{self.layer_num_}.post_attention_layernorm.weight" in weights:
+            self.ffn_norm_weight_ = self._cuda(
+                weights[f"model.layers.{self.layer_num_}.post_attention_layernorm.weight"]
+            )
+
+        inter_size = self.network_config_["intermediate_size"]
+        split_inter_size = inter_size // self.world_size_
+
+        if f"model.layers.{self.layer_num_}.mlp.up_proj.weight" in weights:
+            up_proj = weights[f"model.layers.{self.layer_num_}.mlp.up_proj.weight"][
+                split_inter_size * self.tp_rank_ : split_inter_size * (self.tp_rank_ + 1), :
+            ]
+            self.up_proj = up_proj.transpose(0, 1)
+
+        if f"model.layers.{self.layer_num_}.mlp.gate_proj.weight" in weights:
+            gate_proj = weights[f"model.layers.{self.layer_num_}.mlp.gate_proj.weight"][
+                split_inter_size * self.tp_rank_ : split_inter_size * (self.tp_rank_ + 1), :
+            ]
+            self.gate_proj = gate_proj.transpose(0, 1)
+
+        self._try_cat_to(["gate_proj", "up_proj"], "gate_up_proj", cat_dim=1)
+
+        if f"model.layers.{self.layer_num_}.mlp.down_proj.weight" in weights:
+            self.down_proj = weights[f"model.layers.{self.layer_num_}.mlp.down_proj.weight"][
+                :, split_inter_size * self.tp_rank_ : split_inter_size * (self.tp_rank_ + 1)
+            ]
+            self.down_proj = self._cuda(self.down_proj.transpose(0, 1)) * self.layer_scale
+        return

lightllm/models/minicpm/layer_weights/__init__.py

@@ -0,0 +1,16 @@
+import os
+import json
+import torch
+from lightllm.models.minicpm.layer_weights.transformer_layer_weight import MiniCPMTransformerLayerWeight
+from lightllm.models.minicpm.layer_weights.pre_and_post_layer_weight import MiniCPMPreAndPostLayerWeight
+from lightllm.models.llama.model import LlamaTpPartModel
+
+
+class MiniCPMTpPartModel(LlamaTpPartModel):
+    # weight class
+    transformer_weight_class = MiniCPMTransformerLayerWeight
+    pre_and_post_weight_class = MiniCPMPreAndPostLayerWeight
+
+    def __init__(self, kvargs):
+        super().__init__(kvargs)
+    

lightllm/models/minicpm/layer_weights/pre_and_post_layer_weight.py

@@ -29,6 +29,7 @@
 from lightllm.models.internlm_wquant.model import InternlmTpPartModelWQuant
 from lightllm.models.yi.model import YiTpPartModel
 from lightllm.models.mistral.model import MistralTpPartModel
+from lightllm.models.minicpm.model import MiniCPMTpPartModel
 from lightllm.models.llava.model import LlavaTpPartModel
 from lightllm.models.qwen_vl.model import QWenVLTpPartModel
 from lightllm.models.internlm_xcomposer.model import InternlmComposerTpPartModel
@@ -87,7 +88,7 @@ def exposed_init_model(self, kvargs):
         }
 
         try:
-            self.model_type = model_cfg["model_type"]
+            self.model_type = model_cfg.get("model_type", "")
             if self.model_type == "bloom":
                 self.model = BloomTpPartModel(model_kvargs)
             elif self.model_type == "llama":
@@ -141,6 +142,8 @@ def exposed_init_model(self, kvargs):
                 self.model = StablelmTpPartModel(model_kvargs)
             elif self.model_type == "mixtral":
                 self.model = MixtralTpPartModel(model_kvargs)
+            elif self.model_type == "minicpm" or model_cfg["architectures"][0]=="MiniCPMForCausalLM":
+                self.model = MiniCPMTpPartModel(model_kvargs)
             elif self.model_type == "llava":
                 self.model = LlavaTpPartModel(model_kvargs)
                 self.is_multimodal = True

lightllm/models/minicpm/layer_weights/transformer_layer_weight.py

@@ -69,9 +69,10 @@ def get_tokenizer(
                                                   **kwargs)
 
     model_cfg, _ = PretrainedConfig.get_config_dict(tokenizer_name)
-    if model_cfg["model_type"] == "llava" or model_cfg["model_type"] == "internlmxcomposer2":
+    model_type = model_cfg.get("model_type", "")
+    if model_type == "llava" or model_type == "internlmxcomposer2":
         tokenizer = LlavaTokenizer(tokenizer, model_cfg)
-    elif model_cfg["model_type"] == "qwen" and "visual" in model_cfg:
+    elif model_type == "qwen" and "visual" in model_cfg:
         tokenizer = QWenVLTokenizer(tokenizer, model_cfg)
 
     if not isinstance(tokenizer, PreTrainedTokenizerFast):