Rpan dev lisa

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
+}

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
+}

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 \

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"

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):