[LLM] Add support for low_low_bit benchmark on Windows GPU (#10167)

* Add support for low_low_bit performance test on Windows GPU

* Small fix

* Small fix

* Save memory during converting model process

* Drop the results for first time when loading in low bit on mtl igpu for better performance

* Small fix

python/llm/dev/benchmark/all-in-one/README.md

@@ -43,13 +43,21 @@ test_api:
   - "pytorch_autocast_bf16"
   # - "transformer_autocast_bf16"
   # - "ipex_fp16_gpu" # on Intel GPU
+  # - "bigdl_fp16_gpu" # on Intel GPU
   # - "transformer_int4_gpu"  # on Intel GPU
   # - "optimize_model_gpu"  # on Intel GPU
   # - "deepspeed_transformer_int4_cpu" # on Intel SPR Server
-  # - "transformer_int4_gpu_win" # on Intel GPU for Windows (catch GPU peak memory)
+  # - "transformer_int4_gpu_win" # on Intel GPU for Windows
+  # - "transformer_int4_loadlowbit_gpu_win" # on Intel GPU for Windows using load_low_bit API. Please make sure you have used the save.py to save the converted low bit model
 cpu_embedding: False # whether put embedding to CPU (only avaiable now for gpu win related test_api)
+
 ```
 
+## (Optional) Save model in low bit
+If you choose the `transformer_int4_loadlowbit_gpu_win` test API, you will need to save the model in low bit first.
+
+Run `python save.py` will save all models declared in `repo_id` list into low bit models under `local_model_hub` folder.
+
 ## Run
 
 run `python run.py`, this will output results to `results.csv`.

python/llm/dev/benchmark/all-in-one/config.yaml

@@ -23,5 +23,6 @@ test_api:
   # - "transformer_int4_gpu"  # on Intel GPU
   # - "optimize_model_gpu"  # on Intel GPU
   # - "deepspeed_transformer_int4_cpu" # on Intel SPR Server
-  # - "transformer_int4_gpu_win" # on Intel GPU for Windows (catch GPU peak memory)
+  # - "transformer_int4_gpu_win" # on Intel GPU for Windows
+  # - "transformer_int4_loadlowbit_gpu_win" # on Intel GPU for Windows using load_low_bit API. Please make sure you have used the save.py to save the converted low bit model
 cpu_embedding: False # whether put embedding to CPU (only avaiable now for gpu win related test_api)

python/llm/dev/benchmark/all-in-one/run.py

@@ -86,6 +86,10 @@ def run_model(repo_id, test_api, in_out_pairs, local_model_hub=None, warm_up=1,
         result = run_deepspeed_transformer_int4_cpu(repo_id, local_model_hub, in_out_pairs, warm_up, num_trials, num_beams, low_bit, batch_size)
     elif test_api == 'transformer_int4_gpu_win':
         result = run_transformer_int4_gpu_win(repo_id, local_model_hub, in_out_pairs, warm_up, num_trials, num_beams, low_bit, cpu_embedding, batch_size)
+    elif test_api == 'transformer_int4_loadlowbit_gpu_win':
+        # drop the results of the first time for better performance
+        run_transformer_int4_loadlowbit_gpu_win(repo_id, local_model_hub, in_out_pairs, warm_up, num_trials, num_beams, low_bit, cpu_embedding, batch_size)
+        result = run_transformer_int4_loadlowbit_gpu_win(repo_id, local_model_hub, in_out_pairs, warm_up, num_trials, num_beams, low_bit, cpu_embedding, batch_size)
     elif test_api == 'transformer_autocast_bf16':
         result = run_transformer_autocast_bf16(repo_id, local_model_hub, in_out_pairs, warm_up, num_trials, num_beams, batch_size)
 
@@ -102,7 +106,7 @@ def run_model(repo_id, test_api, in_out_pairs, local_model_hub=None, warm_up=1,
                             num_beams,
                             low_bit,
                             cpu_embedding if 'win' in test_api else 'N/A',
-                            result[in_out_pair][-1][5] if 'int4_gpu' in test_api else 'N/A']) # currently only peak mem for transformer_int4_gpu is caught here
+                            result[in_out_pair][-1][5] if 'int4_gpu' in test_api or 'int4_loadlowbit_gpu' in test_api else 'N/A']) # currently only peak mem for transformer_int4_gpu is caught here
 
 
 def get_model_path(repo_id, local_model_hub):
@@ -800,8 +804,8 @@ def run_transformer_int4_gpu_win(repo_id,
         tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
         model = model.to('xpu')
     elif repo_id in LLAMA_IDS:
-        model = AutoModelForCausalLM.from_pretrained(model_path, load_in_low_bit=low_bit, trust_remote_code=True,
+        model = AutoModelForCausalLM.from_pretrained(model_path, load_in_low_bit=low_bit, optimize_model=True,
-                                                     use_cache=True, cpu_embedding=cpu_embedding).eval()
+                                                     trust_remote_code=True, use_cache=True, cpu_embedding=cpu_embedding).eval()
         tokenizer = LlamaTokenizer.from_pretrained(model_path, trust_remote_code=True)
         model = model.to('xpu')
     elif repo_id in LLAVA_IDS:
@@ -873,6 +877,102 @@ def run_transformer_int4_gpu_win(repo_id,
     gc.collect()
     return result
 
+
+def run_transformer_int4_loadlowbit_gpu_win(repo_id,
+                                            local_model_hub,
+                                            in_out_pairs,
+                                            warm_up,
+                                            num_trials,
+                                            num_beams,
+                                            low_bit,
+                                            cpu_embedding,
+                                            batch_size):
+    from bigdl.llm.transformers import AutoModel, AutoModelForCausalLM
+    from transformers import AutoTokenizer, GPTJForCausalLM, LlamaTokenizer
+    import intel_extension_for_pytorch as ipex
+    model_path = get_model_path(repo_id, local_model_hub)
+    # Load BigDL-LLM optimized low bit model
+    st = time.perf_counter()
+    if repo_id in CHATGLM_IDS:
+        model = AutoModel.load_low_bit(model_path+'-'+low_bit, optimize_model=True, trust_remote_code=True,
+                                       use_cache=True, cpu_embedding=cpu_embedding).eval()
+        tokenizer = AutoTokenizer.from_pretrained(model_path+'-'+low_bit, trust_remote_code=True)
+        model = model.to('xpu')
+    elif repo_id in LLAMA_IDS:
+        model = AutoModelForCausalLM.load_low_bit(model_path+'-'+low_bit, optimize_model=True, trust_remote_code=True,
+                                                  use_cache=True, cpu_embedding=cpu_embedding).eval()
+        tokenizer = LlamaTokenizer.from_pretrained(model_path+'-'+low_bit, trust_remote_code=True)
+        model = model.to('xpu')
+    elif repo_id in LLAVA_IDS:
+        llava_repo_dir = os.environ.get('LLAVA_REPO_DIR')
+        sys.path.append(rf"{llava_repo_dir}")
+        from llava.model.language_model.llava_llama import LlavaLlamaForCausalLM
+        model = AutoModelForCausalLM.load_low_bit(model_path+'-'+low_bit, optimize_model=True, trust_remote_code=True,
+                                                  use_cache=True, cpu_embedding=cpu_embedding).eval()
+        tokenizer = AutoTokenizer.from_pretrained(model_path+'-'+low_bit, trust_remote_code=True)
+        model = model.to('xpu')
+    else:
+        model = AutoModelForCausalLM.load_low_bit(model_path+'-'+low_bit, optimize_model=True, trust_remote_code=True,
+                                                  use_cache=True, cpu_embedding=cpu_embedding).eval()
+        tokenizer = AutoTokenizer.from_pretrained(model_path+'-'+low_bit, trust_remote_code=True)
+        model = model.to('xpu')
+        if isinstance(model, GPTJForCausalLM):
+            # For gpt-j model family, this optimization can provide a better performance.
+            model = ipex.optimize(model.eval(), inplace=True)
+    end = time.perf_counter()
+    print(">> loading of model costs {}s and {}GB".format(end - st, torch.xpu.memory.memory_reserved()/(1024**3)))
+
+    model = BenchmarkWrapper(model)
+
+    result = {}
+    with torch.inference_mode():
+        for in_out in in_out_pairs:
+            try:
+                in_out_len = in_out.split("-")
+                in_len = int(in_out_len[0])
+                out_len = int(in_out_len[1])
+                # As different tokenizer has different encodings,
+                # in_len.txt maybe shorter than we need,
+                # use much longer context to make sure input length
+                test_length = min(in_len*2, 8192)
+                while test_length not in [32, 256, 1024, 2048, 8192]:
+                    test_length = test_length * 2
+                input_str = open(f"prompt/{test_length}.txt", 'r').read()
+                # As different tokenizer has different encodings,
+                # slice the input_ids to ensure the prompt length is required length.
+                input_ids = tokenizer.encode(input_str, return_tensors="pt")
+                input_ids = input_ids[:, :in_len]
+                true_str = tokenizer.batch_decode(input_ids)[0]
+                input_list = [true_str] * batch_size
+                input_ids = tokenizer(input_list, return_tensors="pt").input_ids.to('xpu')
+                actual_in_len = input_ids.shape[1]
+                result[in_out] = []
+                for i in range(num_trials + warm_up):
+                    st = time.perf_counter()
+                    output_ids = model.generate(input_ids, do_sample=False, max_new_tokens=out_len,
+                                                num_beams=num_beams)
+                    torch.xpu.synchronize()
+                    end = time.perf_counter()
+                    output_ids = output_ids.cpu()
+                    print("model generate cost: " + str(end - st))
+                    output = tokenizer.batch_decode(output_ids)
+                    print(output[0])
+                    actual_out_len = output_ids.shape[1] - actual_in_len
+                    if i >= warm_up:
+                        result[in_out].append([model.first_cost, model.rest_cost_mean, model.encoder_time,
+                                               actual_in_len, actual_out_len, model.peak_memory])
+                    # torch.xpu.empty_cache() # this may make first token slower
+            except RuntimeError:
+                traceback.print_exc()
+                pass
+    model.to('cpu')
+    torch.xpu.synchronize()
+    torch.xpu.empty_cache()
+    del model
+    gc.collect()
+    return result
+
+
 def run_transformer_autocast_bf16( repo_id,
                     local_model_hub,
                     in_out_pairs,

python/llm/dev/benchmark/all-in-one/save.py

@@ -0,0 +1,74 @@
+#
+# 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.
+#
+
+# this code is to support converting of model in load bit
+# for performance tests using load_low_bit
+
+import omegaconf
+import time
+import os
+import sys
+import gc
+
+from run import LLAMA_IDS, CHATGLM_IDS, LLAVA_IDS, get_model_path
+
+current_dir = os.path.dirname(os.path.realpath(__file__))
+
+def save_model_in_low_bit(repo_id,
+                          local_model_hub,
+                          low_bit):
+    from bigdl.llm.transformers import AutoModel, AutoModelForCausalLM
+    from transformers import AutoTokenizer, LlamaTokenizer
+    model_path = get_model_path(repo_id, local_model_hub)
+    # Load model in 4 bit,
+    # which convert the relevant layers in the model into INT4 format
+    st = time.perf_counter()
+    if repo_id in CHATGLM_IDS:
+        model = AutoModel.from_pretrained(model_path, load_in_low_bit=low_bit, optimize_model=True,
+                                          trust_remote_code=True, use_cache=True).eval()
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    elif repo_id in LLAMA_IDS:
+        model = AutoModelForCausalLM.from_pretrained(model_path, load_in_low_bit=low_bit, trust_remote_code=True,
+                                                     use_cache=True).eval()
+        tokenizer = LlamaTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    elif repo_id in LLAVA_IDS:
+        llava_repo_dir = os.environ.get('LLAVA_REPO_DIR')
+        sys.path.append(rf"{llava_repo_dir}")
+        from llava.model.language_model.llava_llama import LlavaLlamaForCausalLM
+        model = AutoModelForCausalLM.from_pretrained(model_path, load_in_low_bit=low_bit, optimize_model=True,
+                                          trust_remote_code=True, use_cache=True).eval()
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(model_path, optimize_model=True, load_in_low_bit=low_bit,
+                                                     trust_remote_code=True, use_cache=True).eval()
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    end = time.perf_counter()
+    print(">> loading of and converting of model costs {}s".format(end - st))
+
+    model.save_low_bit(model_path+'-'+low_bit)
+    tokenizer.save_pretrained(model_path+'-'+low_bit)
+
+    del model
+    gc.collect()
+
+if __name__ == '__main__':
+    from omegaconf import OmegaConf
+    conf = OmegaConf.load(f'{current_dir}/config.yaml')
+
+    for model in conf.repo_id:
+        save_model_in_low_bit(repo_id=model,
+                              local_model_hub=conf['local_model_hub'],
+                              low_bit=conf['low_bit'])