LLM: make pipeline parallel inference example more common (#10786)

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)

python/llm/example/GPU/Pipeline-Parallel-Inference/generate.py

@@ -21,7 +21,7 @@
 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 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
                         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 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
                                                  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,16 +85,18 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
     )
 
     # 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():
         prompt = get_prompt(args.prompt, [], system_prompt=DEFAULT_SYSTEM_PROMPT)
         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 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
         # 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()