diff --git a/python/llm/example/CPU/Speculative-Decoding/ziya/README.md b/python/llm/example/CPU/Speculative-Decoding/ziya/README.md new file mode 100644 index 00000000..e3f95fd1 --- /dev/null +++ b/python/llm/example/CPU/Speculative-Decoding/ziya/README.md @@ -0,0 +1,57 @@ +# Ziya +In this directory, you will find examples on how you could run Ziya BF16 inference with self-speculative decoding using BigDL-LLM on [Intel CPUs](../README.md). For illustration purposes,we utilize the [IDEA-CCNL/Ziya-Coding-34B-v1.0](https://huggingface.co/IDEA-CCNL/Ziya-Coding-34B-v1.0) as reference Ziya model. + +## 0. Requirements +To run the example with BigDL-LLM on Intel CPUs, we have some recommended requirements for your machine, please refer to [here](../README.md#recommended-requirements) for more information. + +## Example: Predict Tokens using `generate()` API +In the example [speculative.py](./speculative.py), we show a basic use case for a Ziya model to predict the next N tokens using `generate()` API, with BigDL-LLM speculative decoding optimizations on Intel CPUs. +### 1. Install +We suggest using conda to manage environment: +```bash +conda create -n llm python=3.9 +conda activate llm +pip install --pre --upgrade bigdl-llm[all] +pip install intel_extension_for_pytorch==2.1.0 +pip install transformers==4.35.2 +``` +### 2. Configures high-performing processor environment variables +```bash +source bigdl-llm-init -t +export OMP_NUM_THREADS=48 # you can change 48 here to #cores of one processor socket +``` +### 3. Run + +We recommend to use `numactl` to bind the program to a specified processor socket: + +```bash +numactl -C 0-47 -m 0 python ./speculative.py --repo-id-or-model-path REPO_ID_OR_MODEL_PATH --prompt PROMPT --n-predict N_PREDICT +``` + +For example, 0-47 means bind the python program to core list 0-47 for a 48-core socket. + +Arguments info: + +- `--repo-id-or-model-path REPO_ID_OR_MODEL_PATH`: argument defining the huggingface repo id for the Ziya model (e.g. `IDEA-CCNL/Ziya-Coding-34B-v1.0`) to be downloaded, or the path to the huggingface checkpoint folder. It is default to be `IDEA-CCNL/Ziya-Coding-34B-v1.0`. +- `--prompt PROMPT`: argument defining the prompt to be infered (with integrated prompt format for chat). A default prompt is provided. +- `--n-predict N_PREDICT`: argument defining the max number of tokens to predict. It is default to be `128`. + +#### Sample Output +#### [IDEA-CCNL/Ziya-Coding-34B-v1.0](https://huggingface.co/IDEA-CCNL/Ziya-Coding-34B-v1.0) + +```log +: +写一段快速排序 +: +def quick_sort(arr): + if len(arr) <= 1: + return arr + pivot = arr[len(arr) // 2] + left = [x for x in arr if x < pivot] + middle = [x for x in arr if x == pivot] + right = [x for x in arr if x > pivot] + return quick_sort(left) + middle + quick_sort(right) +Tokens generated 100 +E2E Generation time xx.xxxxs +First token latency xx.xxxxs +``` diff --git a/python/llm/example/CPU/Speculative-Decoding/ziya/speculative.py b/python/llm/example/CPU/Speculative-Decoding/ziya/speculative.py new file mode 100644 index 00000000..0d35294d --- /dev/null +++ b/python/llm/example/CPU/Speculative-Decoding/ziya/speculative.py @@ -0,0 +1,87 @@ +# +# 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 +from bigdl.llm.transformers import AutoModelForCausalLM +from transformers import AutoTokenizer +import argparse +import time +import numpy as np + + +torch.nn.Linear.reset_parameters = lambda x: None +seed=42 +torch.manual_seed(seed) +np.random.seed(seed) + +ZIYA_PROMPT_FORMAT = ": \n{prompt}\n: \n" +prompt = "写一段快速排序" + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Predict Tokens using `generate()` API for Mistral model') + parser.add_argument('--repo-id-or-model-path', type=str, default="IDEA-CCNL/Ziya-Coding-34B-v1.0", + help='The huggingface repo id for the Ziya (e.g. `IDEA-CCNL/Ziya-Coding-34B-v1.0`) to be downloaded' + ', or the path to the huggingface checkpoint folder') + parser.add_argument('--prompt', type=str, default=prompt, + help='Prompt to infer') + parser.add_argument('--n-predict', type=int, default=128, + help='Max tokens to predict') + + args = parser.parse_args() + model_path = args.repo_id_or_model_path + + # Load model in optimized bf16 here. + # Set `speculative=True`` to enable speculative decoding, + # it only works when load_in_low_bit="fp16" on Intel GPU or load_in_low_bit="bf16" on latest Intel Xeon CPU + model = AutoModelForCausalLM.from_pretrained(model_path, + optimize_model=True, + torch_dtype=torch.bfloat16, + load_in_low_bit="bf16", + speculative=True, + torchscript=True, + trust_remote_code=True, + use_cache=True) + + tokenizer = AutoTokenizer.from_pretrained(model_path) + + with torch.inference_mode(): + prompt = ZIYA_PROMPT_FORMAT.format(prompt=args.prompt) + inputs = tokenizer(prompt, return_tensors='pt') + input_ids = inputs.input_ids.to(model.device) + actual_in_len = input_ids.shape[1] + print("actual input_ids length:" + str(actual_in_len)) + attention_mask = inputs.attention_mask.to(model.device) + + # warmup + output = model.generate(input_ids, + max_new_tokens=args.n_predict, + attention_mask=attention_mask, + do_sample=False) + output_str = tokenizer.decode(output[0]) + + # speculative decoding + st = time.perf_counter() + output = model.generate(input_ids, + max_new_tokens=args.n_predict, + attention_mask=attention_mask, + do_sample=False) + output_str = tokenizer.decode(output[0], skip_special_tokens=True) + end = time.perf_counter() + + print(output_str) + print(f"Tokens generated {model.n_token_generated}") + print(f"E2E Generation time {(end - st):.4f}s") + print(f"First token latency {model.first_token_time:.4f}s")