Speculative Ziya on CPU (#10160)
* Speculative Ziya on CPU * Without part of Accelerate with BIGDL_OPT_IPEX
This commit is contained in:
Ziteng Zhang
GitHub
parent
4fbf449c2d
commit
276ef0e885
2 changed files with 144 additions and 0 deletions
57
python/llm/example/CPU/Speculative-Decoding/ziya/README.md
Normal file
57
python/llm/example/CPU/Speculative-Decoding/ziya/README.md
Normal file
|
|
@ -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
|
||||
<human>:
|
||||
写一段快速排序
|
||||
<bot>:
|
||||
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
|
||||
```
|
||||
|
|
@ -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 = "<human>: \n{prompt}\n<bot>: \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")
|
||||
Loading…
Reference in a new issue