Support minicpm-1B in level0 pipeline (#12297)

Author: plusbang

python/llm/example/NPU/HF-Transformers-AutoModels/LLM/Pipeline-Models/README.md

@@ -9,6 +9,7 @@ In this directory, you will find examples on how to directly run HuggingFace `tr
 | Llama2 | [meta-llama/Llama-2-7b-chat-hf](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf) |
 | Llama3 | [meta-llama/Meta-Llama-3-8B-Instruct](https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct) |
 | Baichuan2 | [baichuan-inc/Baichuan2-7B-Chat](https://huggingface.co/baichuan-inc/Baichuan-7B-Chat) |
+| MiniCPM | [openbmb/MiniCPM-1B-sft-bf16](https://huggingface.co/openbmb/MiniCPM-1B-sft-bf16) |
 
 ## 0. Requirements
 To run these examples with IPEX-LLM on Intel NPUs, make sure to install the newest driver version of Intel NPU.
@@ -47,6 +48,9 @@ python llama3.py
 
 :: to run Baichuan2-7B-Chat
 python baichuan2.py
+
+:: to run MiniCPM-1B-sft-bf16
+python minicpm.py
 ```
 
 Arguments info:

python/llm/example/NPU/HF-Transformers-AutoModels/LLM/Pipeline-Models/minicpm.py

@@ -0,0 +1,105 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+import torch
+import time
+import argparse
+from ipex_llm.transformers.npu_model import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from transformers.utils import logging
+import os
+
+logger = logging.get_logger(__name__)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Predict Tokens using `generate()` API for npu model"
+    )
+    parser.add_argument(
+        "--repo-id-or-model-path",
+        type=str,
+        default="openbmb/MiniCPM-1B-sft-bf16",
+        help="The huggingface repo id for the MiniCPM model to be downloaded"
+        ", or the path to the huggingface checkpoint folder",
+    )
+    parser.add_argument("--lowbit-path", type=str,
+        default="",
+        help="The path to the lowbit model folder, leave blank if you do not want to save. \
+             If path not exists, lowbit model will be saved there. \
+             Else, lowbit model will be loaded.",
+    )
+    parser.add_argument('--prompt', type=str, default="What is AI?",
+                        help='Prompt to infer')
+    parser.add_argument("--n-predict", type=int, default=32, help="Max tokens to predict")
+    parser.add_argument("--max-context-len", type=int, default=1024)
+    parser.add_argument("--max-prompt-len", type=int, default=512)
+    parser.add_argument("--disable-transpose-value-cache", action="store_true", default=False)
+
+    args = parser.parse_args()
+    model_path = args.repo_id_or_model_path
+
+    if not args.lowbit_path or not os.path.exists(args.lowbit_path):
+        model = AutoModelForCausalLM.from_pretrained(model_path,
+                                                     optimize_model=True,
+                                                     pipeline=True,
+                                                     max_context_len=args.max_context_len,
+                                                     max_prompt_len=args.max_prompt_len,
+                                                     torch_dtype=torch.float16,
+                                                     attn_implementation="eager",
+                                                     transpose_value_cache=not args.disable_transpose_value_cache,
+                                                     trust_remote_code=True)
+    else:
+        model = AutoModelForCausalLM.load_low_bit(
+            args.lowbit_path,
+            attn_implementation="eager",
+            torch_dtype=torch.float16,
+            max_context_len=args.max_context_len,
+            max_prompt_len=args.max_prompt_len,
+            pipeline=True,
+            transpose_value_cache=not args.disable_transpose_value_cache,
+            trust_remote_code=True
+        )
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+    if args.lowbit_path and not os.path.exists(args.lowbit_path):
+        model.save_low_bit(args.lowbit_path)
+
+    print("-" * 80)
+    print("done")
+    with torch.inference_mode():
+        print("finish to load")
+        for i in range(5):
+            prompt = "<用户>{}<AI>".format(args.prompt)
+            _input_ids = tokenizer.encode(prompt, return_tensors="pt")
+            print("input length:", len(_input_ids[0]))
+            st = time.time()
+            output = model.generate(
+                _input_ids, max_new_tokens=args.n_predict, do_print=True
+            )
+            end = time.time()
+            print(f"Inference time: {end-st} s")
+            input_str = tokenizer.decode(_input_ids[0], skip_special_tokens=False)
+            print("-" * 20, "Input", "-" * 20)
+            print(input_str)
+            output_str = tokenizer.decode(output[0], skip_special_tokens=False)
+            print("-" * 20, "Output", "-" * 20)
+            print(output_str)
+
+    print("-" * 80)
+    print("done")
+    print("success shut down")

python/llm/example/NPU/HF-Transformers-AutoModels/LLM/minicpm.py

@@ -92,7 +92,7 @@
 
         print("finish to load")
         for i in range(5):
-            _input_ids = tokenizer.encode("<用户>{}".format(args.prompt), return_tensors="pt")
+            _input_ids = tokenizer.encode("<用户>{}<AI>".format(args.prompt), return_tensors="pt")
             print("input length:", len(_input_ids[0]))
             st = time.time()
             output = model.generate(

python/llm/src/ipex_llm/transformers/npu_models/convert_mp.py

@@ -227,6 +227,46 @@ def convert_baichuan(
     convert_forward(model, module.BaichuanModel, baichuan_model_forward)
 
 
+def convert_minicpm(
+    model: torch.nn.Module,
+    max_output_len=1024,
+    max_prompt_len=1024,
+    decoder=False,
+    inter_pp=None,
+    intra_pp=None,
+    transpose_value_cache=True,
+):
+    from ipex_llm.transformers.npu_models.minicpm_mp import gen_minicpm_fused_model_forward
+    from ipex_llm.transformers.npu_models.minicpm_mp import DecodeRunner, PrefillRunner
+    modeling_module_name = model.__class__.__module__
+    module = importlib.import_module(modeling_module_name)
+
+    if decoder:
+        decode_runner = DecodeRunner(
+            model,
+            max_seq_len=max_output_len,
+            inter_pp=inter_pp,
+            intra_pp=intra_pp,
+            transpose_value_cache=transpose_value_cache,
+        )
+    else:
+        decode_runner = None
+    prefill_runner = PrefillRunner(
+        model,
+        max_output_len=max_output_len,
+        max_prompt_len=max_prompt_len,
+        transpose_value_cache=transpose_value_cache,
+    )
+    minicpm_model_forward = gen_minicpm_fused_model_forward(
+        prefill_runner=prefill_runner, decode_runner=decode_runner
+    )
+    convert_forward(model, module.MiniCPMModel, minicpm_model_forward)
+    if model.config.num_hidden_layers == 40:
+        # for minicpm-2b
+        from ipex_llm.transformers.npu_models.minicpm_mp import minicpm_casullm_forward
+        convert_forward(model, module.MiniCPMForCausalLM, minicpm_casullm_forward)
+
+
 def optimize_llm(
     model: torch.nn.Module,
     max_context_len=1024,
@@ -291,41 +331,13 @@ def optimize_llm(
             intra_pp = 2
         if inter_pp is None:
             inter_pp = 2
-
-        from ipex_llm.transformers.npu_models.minicpm_mp import gen_minicpm_fused_model_forward
-        from ipex_llm.transformers.npu_models.minicpm_mp import DecodeRunner, PrefillRunner
-
-        modeling_module_name = model.__class__.__module__
-        module = importlib.import_module(modeling_module_name)
-
-        if model.config.num_hidden_layers == 52:
-            # for minicpm-1b
-            transpose_cache = transpose_value_cache
-        elif model.config.num_hidden_layers == 40:
-            # for minicpm-2b
-            transpose_cache = False
-
-        decode_runner = DecodeRunner(
-            model,
-            max_seq_len=max_context_len,
-            inter_pp=inter_pp,
-            intra_pp=intra_pp,
-            transpose_value_cache=transpose_cache,
-        )
-        prefill_runner = PrefillRunner(
-            model,
-            max_output_len=max_context_len,
-            max_prompt_len=max_prompt_len,
-            transpose_value_cache=transpose_cache,
-        )
-        minicpm_model_forward = gen_minicpm_fused_model_forward(
-            prefill_runner=prefill_runner, decode_runner=decode_runner
-        )
-        convert_forward(model, module.MiniCPMModel, minicpm_model_forward)
-        if model.config.num_hidden_layers == 40:
-            # for minicpm-2b
-            from ipex_llm.transformers.npu_models.minicpm_mp import minicpm_casullm_forward
-            convert_forward(model, module.MiniCPMForCausalLM, minicpm_casullm_forward)
+        convert_minicpm(model,
+                        max_output_len=max_context_len,
+                        max_prompt_len=max_prompt_len,
+                        inter_pp=inter_pp,
+                        intra_pp=intra_pp,
+                        decoder=True,
+                        transpose_value_cache=transpose_value_cache)
     elif model.config.model_type == "baichuan" and model.config.num_hidden_layers == 32:
         # for Baichuan2-7B
         if intra_pp is None:
@@ -339,7 +351,7 @@ def optimize_llm(
                          intra_pp=intra_pp,
                          decoder=True,
                          transpose_value_cache=transpose_value_cache)
-    if isinstance(model.lm_head, SlicedLMHead):
+    if hasattr(model, 'lm_head') and isinstance(model.lm_head, SlicedLMHead):
         model.lm_head.get_fused_lm_head()
 
 

python/llm/src/ipex_llm/transformers/npu_models/minicpm_mp.py

@@ -54,7 +54,7 @@
 from torch.nn import CrossEntropyLoss
 
 
-class LowBitLlamaMultiDecoderlayer(LLMBaseNNFactory):
+class LowBitMinicpmMultiDecoderlayer(LLMBaseNNFactory):
     def __init__(
         self,
         # batch_size: int,
@@ -118,31 +118,13 @@ def __init__(
 
         # Self Attention
         if mode == "decode":
-            attention_mask = self.create_input_op((self.batch_size, 1, 1, self.max_seq_len + 1))
+            attention_mask = self.create_input_op((self.batch_size, 1, 1, self.max_seq_len + 1),
+                                                  dtype=np.int64)
         else:
-            attention_mask = self.create_input_op((self.batch_size, 1, self.seq_len, self.seq_len))
+            attention_mask = self.create_input_op((self.batch_size, 1, self.seq_len, self.seq_len),
+                                                  dtype=np.int64)
 
-        position_ids = self.create_input_op((self.batch_size, self.seq_len))
-        past_keys = []
-        past_values = []
-        if mode == "decode":
-            for i in range(num_layers):
-                past_key = self.create_cache_op(
-                    (self.batch_size, self.num_key_value_heads, self.max_seq_len, self.head_dim)
-                )
-                if transpose_value:
-                    past_value = self.create_cache_op(
-                        (self.batch_size, self.num_key_value_heads, self.head_dim, self.max_seq_len)
-                    )
-                else:
-                    past_value = self.create_cache_op(
-                        (self.batch_size, self.num_key_value_heads, self.max_seq_len, self.head_dim)
-                    )
-                past_keys.append(past_key)
-                past_values.append(past_value)
-        else:
-            past_keys = [None] * num_layers
-            past_values = [None] * num_layers
+        position_ids = self.create_input_op((self.batch_size, self.seq_len), dtype=np.int64)
 
         if input_layernorm_weights is None:
             input_layernorm_weights = []
@@ -168,6 +150,27 @@ def __init__(
             input_layernorm_weights = [self.constant(w) for w in input_layernorm_weights]
             post_attn_layernorm_weights = [self.constant(w) for w in post_attn_layernorm_weights]
 
+        past_keys = []
+        past_values = []
+        if mode == "decode":
+            for i in range(num_layers):
+                past_key = self.create_cache_op(
+                    (self.batch_size, self.num_key_value_heads, self.max_seq_len, self.head_dim)
+                )
+                if transpose_value:
+                    past_value = self.create_cache_op(
+                        (self.batch_size, self.num_key_value_heads, self.head_dim, self.max_seq_len)
+                    )
+                else:
+                    past_value = self.create_cache_op(
+                        (self.batch_size, self.num_key_value_heads, self.max_seq_len, self.head_dim)
+                    )
+                past_keys.append(past_key)
+                past_values.append(past_value)
+        else:
+            past_keys = [None] * num_layers
+            past_values = [None] * num_layers
+
         hidden_states = input
 
         curr_key_values = []
@@ -297,7 +300,7 @@ def __init__(
             start, end = self.layer_ranges[i]
             lm_0 = input_laynorm_weights[start:end]
             lm_1 = post_attn_layernorm_weights[start:end]
-            decoder = LowBitLlamaMultiDecoderlayer(
+            decoder = LowBitMinicpmMultiDecoderlayer(
                 [1, 1, num_heads * head_dim],
                 input_layernorm_weights=lm_0,
                 post_attn_layernorm_weights=lm_1,
@@ -334,15 +337,15 @@ def forward(
 
         inputs = (
             hidden_states.to(torch.float16),
-            attention_mask,
-            position_ids.to(torch.float16),
+            attention_mask.to(torch.int64),
+            position_ids.to(torch.int64),
         )
 
         for i in range(self.intra_stages):
             start, end = self.layer_ranges[i]
             self.backend_decoders[i].update_cache(past_key_value, self.layer_indexes[start:end])
 
-        hidden_states, new_keys, new_values = LowBitLlamaMultiDecoderlayer.run_decoders(
+        hidden_states, new_keys, new_values = LowBitMinicpmMultiDecoderlayer.run_decoders(
             inputs,
             decoders=self.backend_decoders)
 
@@ -403,7 +406,7 @@ def __init__(
             np_dtype = np.float16
 
         self.backend_cls_prefill = partial(
-            LowBitLlamaMultiDecoderlayer,
+            LowBitMinicpmMultiDecoderlayer,
             num_heads=num_heads,
             num_key_value_heads=num_key_value_heads,
             num_layers=1,
@@ -445,7 +448,9 @@ def forward(
         seq_len = hidden_states.shape[1]
 
         backend_cls = self.backend_cls_prefill
-        inputs = (hidden_states.to(torch.float16), attention_mask, position_ids.to(torch.float16))
+        inputs = (hidden_states.to(torch.float16),
+                  attention_mask.to(torch.int64),
+                  position_ids.to(torch.int64))
         inputs += (self.layer_norm_0, self.layer_norm_1)
         hidden_states, past_key, past_value = run_model(
             inputs, self.op_parameters, backend_cls, self.op_id, replica=2
@@ -578,9 +583,9 @@ def run_decode(
 
                 pad_mask = (0, pad_len)
                 padded_causal_mask = F.pad(
-                    causal_mask.to(torch.float16), pad_mask, value=torch.finfo(torch.float16).min
+                    causal_mask.to(torch.int64), pad_mask, value=torch.iinfo(torch.int64).min
                 )
-                padded_causal_mask[:, :, :, -1] = 0.0
+                padded_causal_mask[:, :, :, -1] = 0
                 dist.recv(hidden_states, src=rank - 1)
                 layer_outputs = multi_decoder(
                     hidden_states,
@@ -831,9 +836,9 @@ def forward(
         hidden_states = F.pad(hidden_states.to(torch.float16), (0, 0, 0, pad_len), value=0.0)
         position_ids = F.pad(position_ids, (0, pad_len), value=0)
         attention_mask = F.pad(
-            attention_mask.to(torch.float16),
+            attention_mask.to(torch.int64),
             (0, pad_len, 0, pad_len),
-            value=torch.finfo(torch.float16).min,
+            value=torch.iinfo(torch.int64).min,
         )
 
         args = (hidden_states, position_ids, attention_mask, past_key_value)