From fabf54e0521d727dde86eca05e584a3848494045 Mon Sep 17 00:00:00 2001 From: binbin Deng <108676127+plusbang@users.noreply.github.com> Date: Wed, 24 Apr 2024 09:28:52 +0800 Subject: [PATCH] LLM: make pipeline parallel inference example more common (#10786) --- .../GPU/Pipeline-Parallel-Inference/README.md | 7 ++-- .../Pipeline-Parallel-Inference/generate.py | 32 +++++++++++-------- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/python/llm/example/GPU/Pipeline-Parallel-Inference/README.md b/python/llm/example/GPU/Pipeline-Parallel-Inference/README.md index 58379184..8974afdd 100644 --- a/python/llm/example/GPU/Pipeline-Parallel-Inference/README.md +++ b/python/llm/example/GPU/Pipeline-Parallel-Inference/README.md @@ -38,7 +38,7 @@ python setup.py install > **Important**: IPEX 2.1.10+xpu requires Intel® oneAPI Base Toolkit's version == 2024.0. Please make sure you have installed the correct version. -### 2. Run tensor parallel inference on multiple GPUs +### 2. Run pipeline parallel inference on multiple GPUs Here, we provide example usages on different models and different hardwares. Please refer to the appropriate script based on your model and device: ### 3. Run @@ -51,13 +51,14 @@ export SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS=1 ``` ``` -python ./generate.py --repo-id-or-model-path REPO_ID_OR_MODEL_PATH --prompt PROMPT --n-predict N_PREDICT +python ./generate.py --repo-id-or-model-path REPO_ID_OR_MODEL_PATH --prompt PROMPT --n-predict N_PREDICT --gpu-num GPU_NUM ``` Arguments info: -- `--repo-id-or-model-path REPO_ID_OR_MODEL_PATH`: argument defining the huggingface repo id for the Llama2 model (e.g. `meta-llama/Llama-2-7b-chat-hf`) to be downloaded, or the path to the huggingface checkpoint folder. It is default to be `'meta-llama/Llama-2-7b-chat-hf'`. +- `--repo-id-or-model-path REPO_ID_OR_MODEL_PATH`: argument defining the huggingface repo id for the Llama2 model (e.g. `meta-llama/Llama-2-7b-chat-hf` and `meta-llama/Llama-2-13b-chat-hf`) to be downloaded, or the path to the huggingface checkpoint folder. It is default to be `'meta-llama/Llama-2-7b-chat-hf'`. - `--prompt PROMPT`: argument defining the prompt to be infered (with integrated prompt format for chat). It is default to be `'What is AI?'`. - `--n-predict N_PREDICT`: argument defining the max number of tokens to predict. It is default to be `32`. +- `--gpu-num GPU_NUM`: argument defining the number of GPU to use. It is default to be `2`. #### Sample Output #### [meta-llama/Llama-2-7b-chat-hf](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf) diff --git a/python/llm/example/GPU/Pipeline-Parallel-Inference/generate.py b/python/llm/example/GPU/Pipeline-Parallel-Inference/generate.py index e54cf881..badba4a9 100644 --- a/python/llm/example/GPU/Pipeline-Parallel-Inference/generate.py +++ b/python/llm/example/GPU/Pipeline-Parallel-Inference/generate.py @@ -21,7 +21,7 @@ import time import argparse from ipex_llm.transformers import AutoModelForCausalLM -from transformers import LlamaTokenizer +from transformers import AutoTokenizer # you could tune the prompt based on your own model, # here the prompt tuning refers to https://huggingface.co/georgesung/llama2_7b_chat_uncensored#prompt-style @@ -51,6 +51,7 @@ if __name__ == '__main__': help='Prompt to infer') parser.add_argument('--n-predict', type=int, default=32, help='Max tokens to predict') + parser.add_argument('--gpu-num', type=int, default=2, help='GPU number to use') args = parser.parse_args() model_path = args.repo_id_or_model_path @@ -62,19 +63,19 @@ if __name__ == '__main__': optimize_model=True, trust_remote_code=True, use_cache=True) - first_half = ['model.embed_tokens', 'model.layers.0', 'model.layers.1', 'model.layers.2', - 'model.layers.3', 'model.layers.4', 'model.layers.5', 'model.layers.6', - 'model.layers.7', 'model.layers.8', 'model.layers.9', 'model.layers.10', - 'model.layers.11', 'model.layers.12', 'model.layers.13', 'model.layers.14', - 'model.layers.15'] - second_half = ['model.layers.16', 'model.layers.17', 'model.layers.18', 'model.layers.19', - 'model.layers.20', 'model.layers.21', 'model.layers.22', 'model.layers.23', - 'model.layers.24', 'model.layers.25', 'model.layers.26', 'model.layers.27', - 'model.layers.28', 'model.layers.29', 'model.layers.30', 'model.layers.31', - 'model.norm', 'lm_head'] - device_map=({key: 'xpu:0' for key in first_half}) - device_map.update({key: 'xpu:1' for key in second_half}) + model_layers = ['model.embed_tokens'] + for i in range(model.config.num_hidden_layers): + model_layers.append(f'model.layers.{i}') + model_layers = model_layers + ['model.norm', 'lm_head'] + + device_map = {} + split_len = len(model_layers) // args.gpu_num + for i in range(args.gpu_num): + device_map.update({key: f'xpu:{i}' for key in model_layers[split_len * i: split_len * (i + 1)]}) + if i == args.gpu_num - 1: + device_map.update({key: f'xpu:{i}' for key in model_layers[split_len * (i + 1): ]}) + from accelerate import dispatch_model model = dispatch_model( model, @@ -84,7 +85,7 @@ if __name__ == '__main__': ) # Load tokenizer - tokenizer = LlamaTokenizer.from_pretrained(model_path, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) # Generate predicted tokens with torch.inference_mode(): @@ -92,8 +93,10 @@ if __name__ == '__main__': input_ids = tokenizer.encode(prompt, return_tensors="pt").to('xpu:0') # ipex_llm model needs a warmup, then inference time can be accurate output = model.generate(input_ids, + do_sample=False, max_new_tokens=args.n_predict) output = model.generate(input_ids, + do_sample=False, max_new_tokens=args.n_predict) # start inference @@ -103,6 +106,7 @@ if __name__ == '__main__': # it is important to set `use_cache=True` explicitly in the `generate` function # to obtain optimal performance with IPEX-LLM INT4 optimizations output = model.generate(input_ids, + do_sample=False, max_new_tokens=args.n_predict) torch.xpu.synchronize() end = time.time()