Merge pull request #716 from OptimalScale/rpan-improve-lisa-script

Add more memory saving options for LISA script

scripts/run_finetune_with_lisa.sh

@@ -11,6 +11,13 @@ lisa_activated_layers=1
 lisa_interval_steps=20
 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
+
 while [[ $# -ge 1 ]]; do
   key="$1"
   case ${key} in
@@ -38,6 +45,26 @@ while [[ $# -ge 1 ]]; do
       lisa_interval_steps="$2"
       shift
       ;;
+    --gradient_checkpointing)
+      gradient_checkpointing="$2"
+      shift
+      ;;
+    --deepspeed)
+      ds_config_file="$2"
+      shift
+      ;;
+    --use_flash_attention)
+      use_flash_attention="$2"
+      shift
+      ;;
+    --gradient_accumulation_steps)
+      gradient_accumulation_steps="$2"
+      shift
+      ;;
+    --block_size)
+      block_size="$2"
+      shift
+      ;;
     *)
       echo "error: unknown option \"${key}\"" 1>&2
       exit 1
@@ -58,9 +85,9 @@ deepspeed ${deepspeed_args} \
     --output_dir ${output_dir} --overwrite_output_dir \
     --num_train_epochs 1 \
     --learning_rate 2e-5 \
-    --block_size 512 \
+    --block_size ${block_size} \
     --per_device_train_batch_size 1 \
-    --deepspeed configs/ds_config_zero2_no_offload.json \
+    --deepspeed ${ds_config_file} \
     --fp16 \
     --run_name finetune \
     --validation_split_percentage 0 \
@@ -69,7 +96,9 @@ deepspeed ${deepspeed_args} \
     --ddp_timeout 72000 \
     --save_steps 5000 \
     --dataloader_num_workers 1 \
-    --gradient_checkpointing True \
+    --gradient_checkpointing ${gradient_checkpointing} \
+    --use_flash_attention ${use_flash_attention} \
+    --gradient_accumulation_steps ${gradient_accumulation_steps} \
     --use_lisa 1 \
     --lisa_activated_layers ${lisa_activated_layers} \
     --lisa_interval_steps ${lisa_interval_steps} \