From c050639da7aa089d94e8f8e77ce80af19c84d87d Mon Sep 17 00:00:00 2001
From: Dominic789654 <47224289+Dominic789654@users.noreply.github.com>
Date: Thu, 28 Mar 2024 01:54:29 +0800
Subject: [PATCH] Fix convergence problem caused by deepspeed/hf

- Also add different settings of deepspeed/hf for single-gpu and multi-gpu
---
 configs/ds_config_zero0_no_offload.json | 31 ++++++++++++++++++++++
 configs/ds_config_zero3_no_offload.json | 34 +++++++++++++++++++++++++
 scripts/run_finetune.sh                 |  2 +-
 scripts/run_finetune_with_lisa.sh       | 11 ++++++--
 src/lmflow/pipeline/finetuner.py        |  5 ++--
 5 files changed, 77 insertions(+), 6 deletions(-)
 create mode 100644 configs/ds_config_zero0_no_offload.json
 create mode 100644 configs/ds_config_zero3_no_offload.json

diff --git a/configs/ds_config_zero0_no_offload.json b/configs/ds_config_zero0_no_offload.json
new file mode 100644
index 000000000..4ba4833e4
--- /dev/null
+++ b/configs/ds_config_zero0_no_offload.json
@@ -0,0 +1,31 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 0,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
diff --git a/configs/ds_config_zero3_no_offload.json b/configs/ds_config_zero3_no_offload.json
new file mode 100644
index 000000000..e0c74b24e
--- /dev/null
+++ b/configs/ds_config_zero3_no_offload.json
@@ -0,0 +1,34 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "stage3_max_live_parameters": 2e10,
+        "stage3_max_reuse_distance": 2e10,
+        "stage3_gather_16bit_weights_on_model_save": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
diff --git a/scripts/run_finetune.sh b/scripts/run_finetune.sh
index bafff31bb..02d8613c4 100755
--- a/scripts/run_finetune.sh
+++ b/scripts/run_finetune.sh
@@ -49,7 +49,7 @@ deepspeed ${deepspeed_args} \
     --num_train_epochs 0.01 \
     --learning_rate 2e-5 \
     --disable_group_texts 1 \
-    --block_size 512 \
+    --block_size 256 \
     --per_device_train_batch_size 1 \
     --deepspeed configs/ds_config_zero3.json \
     --fp16 \
diff --git a/scripts/run_finetune_with_lisa.sh b/scripts/run_finetune_with_lisa.sh
index a6456dc58..63a04512d 100755
--- a/scripts/run_finetune_with_lisa.sh
+++ b/scripts/run_finetune_with_lisa.sh
@@ -13,10 +13,17 @@ deepspeed_args="--master_port=11000"
 
 # Other optional arguments that can improve memory saving
 gradient_checkpointing=True
-ds_config_file="configs/ds_config_zero2_no_offload.json"
 use_flash_attention=0
 gradient_accumulation_steps=1
-block_size=512
+block_size=256
+
+# Enable model parallelism for multiple gpus, modify this if you prefer
+# customized deepspeed zero-redundancy optimization settings
+num_gpu=$(python -c "import torch; print(torch.cuda.device_count())")
+ds_config_file=configs/ds_config_zero0_no_offload.json
+if [ ${num_gpu} -ge 2 ]; then
+  ds_config_file=configs/ds_config_zero2_no_offload.json
+fi
 
 while [[ $# -ge 1 ]]; do
   key="$1"
diff --git a/src/lmflow/pipeline/finetuner.py b/src/lmflow/pipeline/finetuner.py
index 98bce5f32..840422e6e 100644
--- a/src/lmflow/pipeline/finetuner.py
+++ b/src/lmflow/pipeline/finetuner.py
@@ -317,8 +317,7 @@ def __init__(self, n_layers, interval_steps, model):
                         self.layers_attribute = 'model.transformer.h'  # General access path
                     self.total_layers = len(eval('self.' + self.layers_attribute))  # Dynamically execute to get the number of layers
 
-                    # Freeze all layers upon initialization
-                    self.freeze_all_layers()
+                    self.switch_active_layers()
                     self.active_layers_indices = []
 
                 def freeze_all_layers(self):
@@ -329,7 +328,7 @@ def freeze_all_layers(self):
 
                 def on_step_begin(self, args, state, control, **kwargs):
                     # Check if it's time to switch active layers, including at step 0
-                    if state.global_step % self.interval_steps == 0 or state.global_step == 1:
+                    if state.global_step % self.interval_steps == 0:
                         self.switch_active_layers()
 
                 def switch_active_layers(self):