Finetune ChatGLM with Deepspeed Zero3 LoRA (#11314)

* Fintune ChatGLM with Deepspeed Zero3 LoRA

* add deepspeed zero3 config

* rename config

* remove offload_param

* add save_checkpoint parameter

* Update lora_deepspeed_zero3_finetune_chatglm3_6b_arc_2_card.sh

* refine

python/llm/example/GPU/LLM-Finetuning/LoRA/README.md

@@ -33,6 +33,12 @@ Here, we provide example usages on different hardware. Please refer to the appro
 bash lora_finetune_llama2_7b_arc_1_card.sh
 ```
 
+##### Finetuning ChatGLM3-6B on two Arc A770
+
+```bash
+bash lora_deepspeed_zero3_finetune_chatglm3_6b_arc_2_card.sh
+```
+
 ##### Finetuning LLaMA2-7B on four Intel Data Center GPU Max 1100
 
 ```bash

python/llm/example/GPU/LLM-Finetuning/LoRA/alpaca_lora_finetuning.py

@@ -107,6 +107,8 @@ def train(
     gradient_checkpointing: bool = False,
     deepspeed: str = None,
     training_mode: str = "lora",
+    deepspeed_zero3: bool = False,
+    save_checkpoint: bool = True,
 ):
     invalidInputError(training_mode == "lora",
                       f"This example is for lora training mode, but got training_mode={training_mode}.")
@@ -136,6 +138,8 @@ def train(
             f"resume_from_checkpoint: {resume_from_checkpoint or False}\n"
             f"prompt template: {prompt_template_name}\n"
             f"training_mode: {training_mode}\n"
+            f"deepspeed_zero3: {deepspeed_zero3}\n"
+            f"save_checkpoint: {save_checkpoint}\n"
         )
     assert (
         base_model
@@ -154,28 +158,54 @@ def train(
     # Check if parameter passed or if set within environ
     use_wandb = wandb_check(wandb_project, wandb_watch, wandb_log_model)
 
+    if deepspeed_zero3:
+        deepspeed = deepspeed if deepspeed is not None else "./deepspeed_zero3_config.json"
+
     if saved_low_bit_model is not None:
         # Load the low bit optimized model if provide the saved path
-        model = AutoModelForCausalLM.load_low_bit(
+        if deepspeed_zero3:
-            saved_low_bit_model,
+            import deepspeed as ds
-            optimize_model=False,
+            with ds.zero.Init(config_dict_or_path=deepspeed):
-            torch_dtype=torch.bfloat16,
+                model = AutoModelForCausalLM.load_low_bit(
-            modules_to_not_convert=["lm_head"],
+                    saved_low_bit_model,
-            trust_remote_code=True,
+                    optimize_model=False,
-        )
+                    torch_dtype=torch.bfloat16,
+                    modules_to_not_convert=["lm_head"],
+                    trust_remote_code=True,
+                )
+        else:
+            model = AutoModelForCausalLM.load_low_bit(
+                saved_low_bit_model,
+                optimize_model=False,
+                torch_dtype=torch.bfloat16,
+                modules_to_not_convert=["lm_head"],
+                trust_remote_code=True,
+            )
     else:
-        model = AutoModelForCausalLM.from_pretrained(
+        if deepspeed_zero3:
-            base_model,
+            import deepspeed as ds
-            load_in_low_bit="bf16",
+            with ds.zero.Init(config_dict_or_path=deepspeed):
-            optimize_model=False,
+                model = AutoModelForCausalLM.from_pretrained(
-            torch_dtype=torch.bfloat16,
+                    base_model,
-            modules_to_not_convert=["lm_head"],
+                    load_in_low_bit="bf16",
-            trust_remote_code=True,
+                    optimize_model=False,
-        )
+                    torch_dtype=torch.bfloat16,
-
+                    modules_to_not_convert=["lm_head"],
-    print(f"Model loaded on rank {os.environ.get('LOCAL_RANK')}")
+                    trust_remote_code=True,
-    model = model.to(f'xpu:{os.environ.get("LOCAL_RANK", 0)}')
+                )
-    print(f"Model moved to rank {os.environ.get('LOCAL_RANK')}")
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                base_model,
+                load_in_low_bit="bf16",
+                optimize_model=False,
+                torch_dtype=torch.bfloat16,
+                modules_to_not_convert=["lm_head"],
+                trust_remote_code=True,
+            )
+    if not deepspeed_zero3:
+        print(f"Model loaded on rank {os.environ.get('LOCAL_RANK')}")
+        model = model.to(f'xpu:{os.environ.get("LOCAL_RANK", 0)}')
+        print(f"Model moved to rank {os.environ.get('LOCAL_RANK')}")
 
     tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
     print(f"Tokenizer loaded on rank {os.environ.get('LOCAL_RANK')}")
@@ -234,12 +264,12 @@ def train(
             logging_steps=1,
             optim="adamw_torch",
             evaluation_strategy="steps" if val_set_size > 0 else "no",
-            save_strategy="steps",
+            save_strategy="steps" if save_checkpoint else "no",
             eval_steps=100 if val_set_size > 0 else None,
             save_steps=100,
             output_dir=output_dir,
             save_total_limit=100,
-            load_best_model_at_end=True if val_set_size > 0 else False,
+            load_best_model_at_end=True if val_set_size > 0 and save_checkpoint else False,
             ddp_find_unused_parameters=False if ddp else None,
             group_by_length=group_by_length,
             report_to="wandb" if use_wandb else None,

python/llm/example/GPU/LLM-Finetuning/LoRA/deepspeed_zero3_config.json

@@ -0,0 +1,16 @@
+{
+    "zero_optimization": {
+      "stage": 3,
+      "contiguous_gradients": true,
+      "overlap_comm": true,
+      "offload_optimizer": {"device": "cpu"}
+    },
+    "bf16": {
+      "enabled": true
+    },
+    "world_size":2,
+    "train_batch_size": 2,
+    "train_micro_batch_size_per_gpu": 1,
+    "gradient_accumulation_steps": 1,
+    "stage3_gather_16bit_weights_on_model_save":true
+}

python/llm/example/GPU/LLM-Finetuning/LoRA/lora_deepspeed_zero3_finetune_chatglm3_6b_arc_2_card.sh

@@ -0,0 +1,43 @@
+#
+# 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.
+#
+
+export MASTER_ADDR=127.0.0.1
+export MASTER_PORT=29503
+export FI_PROVIDER=tcp
+export CCL_ATL_TRANSPORT=ofi
+export CCL_ZE_IPC_EXCHANGE=sockets
+
+basekit_root=/opt/intel/oneapi
+source $basekit_root/setvars.sh --force
+source $basekit_root/ccl/latest/env/vars.sh --force
+
+NUM_GPUS=2 # number of used GPU
+export USE_XETLA=OFF
+export SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS=2
+export TORCH_LLM_ALLREDUCE=0 # Different from PVC
+
+mpirun -n $NUM_GPUS \
+    python ./alpaca_lora_finetuning.py \
+       --base_model "THUDM/chatglm3-6b" \
+       --data_path "yahma/alpaca-cleaned" \
+       --output_dir "./ipex-llm-lora-alpaca" \
+       --gradient_checkpointing True \
+       --lora_target_modules "['query_key_value', 'dense', 'dense_h_to_4h', 'dense_4h_to_h']" \
+       --micro_batch_size 1 \
+       --batch_size 2 \
+       --save_checkpoint False \
+       --deepspeed_zero3 True > lora_deepspeed_zero3_finetune_chatglm3_6b.log
+