From f3274235d8d3dcb704fe6736aaa29bb37167b0f0 Mon Sep 17 00:00:00 2001
From: sennnnn <cyanlaser@stu.pku.edu.cn>
Date: Sun, 10 Mar 2024 11:27:33 +0800
Subject: [PATCH] Support deepspeed zero2 and zero2_offload for latte training.

Former-commit-id: c87e95dd357e6fa6d6fb2dcfee781bc49c039e13
---
 .../deepspeed_zero2_config.yaml               | 13 ++++++++++
 .../deepspeed_zero2_offload_config.yaml       | 13 ++++++++++
 scripts/accelerate_configs/zero2.json         | 23 ++++++++++++++++
 scripts/accelerate_configs/zero2_offload.json | 26 +++++++++++++++++++
 scripts/train_256_dsz2.sh                     | 20 ++++++++++++++
 scripts/train_256_dsz2_off.sh                 | 20 ++++++++++++++
 6 files changed, 115 insertions(+)
 create mode 100644 scripts/accelerate_configs/deepspeed_zero2_config.yaml
 create mode 100644 scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml
 create mode 100644 scripts/accelerate_configs/zero2.json
 create mode 100644 scripts/accelerate_configs/zero2_offload.json
 create mode 100644 scripts/train_256_dsz2.sh
 create mode 100644 scripts/train_256_dsz2_off.sh

diff --git a/scripts/accelerate_configs/deepspeed_zero2_config.yaml b/scripts/accelerate_configs/deepspeed_zero2_config.yaml
new file mode 100644
index 000000000..43ec1e80c
--- /dev/null
+++ b/scripts/accelerate_configs/deepspeed_zero2_config.yaml
@@ -0,0 +1,13 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: DEEPSPEED
+deepspeed_config:
+ deepspeed_config_file: scripts/accelerate_configs/zero2.json
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: 29501
+main_training_function: main
+num_machines: 1
+num_processes: 8
+gpu_ids: 0,1,2,3,4,5,6,7
+use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml b/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml
new file mode 100644
index 000000000..9c96bae67
--- /dev/null
+++ b/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml
@@ -0,0 +1,13 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: DEEPSPEED
+deepspeed_config:
+ deepspeed_config_file: scripts/accelerate_configs/zero2_offload.json
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: 29501
+main_training_function: main
+num_machines: 1
+num_processes: 8
+gpu_ids: 0,1,2,3,4,5,6,7
+use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/zero2.json b/scripts/accelerate_configs/zero2.json
new file mode 100644
index 000000000..c95ebefe0
--- /dev/null
+++ b/scripts/accelerate_configs/zero2.json
@@ -0,0 +1,23 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "train_micro_batch_size_per_gpu": "auto",
+    "train_batch_size": "auto",
+    "gradient_accumulation_steps": "auto",
+    "zero_optimization": {
+        "stage": 2,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto"
+    }
+}
\ No newline at end of file
diff --git a/scripts/accelerate_configs/zero2_offload.json b/scripts/accelerate_configs/zero2_offload.json
new file mode 100644
index 000000000..4ea31ddc8
--- /dev/null
+++ b/scripts/accelerate_configs/zero2_offload.json
@@ -0,0 +1,26 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "train_micro_batch_size_per_gpu": "auto",
+    "train_batch_size": "auto",
+    "gradient_accumulation_steps": "auto",
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu"
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto"
+    }
+}
\ No newline at end of file
diff --git a/scripts/train_256_dsz2.sh b/scripts/train_256_dsz2.sh
new file mode 100644
index 000000000..120bf7547
--- /dev/null
+++ b/scripts/train_256_dsz2.sh
@@ -0,0 +1,20 @@
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml \
+    opensora/train/train.py \
+    --model Latte-XL/122 \
+    --dataset sky \
+    --ae stabilityai/sd-vae-ft-mse \
+    --data-path datasets/sky_timelapse/sky_train \
+    --extras 1 \
+    --sample-rate 3 \
+    --num-frames 16 \
+    --max-image-size 256 \
+    --max-train-steps 1000000 \
+    --local-batch-size 5 \
+    --lr 1e-4 \
+    --ckpt-every 500 \
+    --log-every 50 \
+    --gradient-checkpointing \
+    --attention_mode math \
+    --mixed-precision fp16
diff --git a/scripts/train_256_dsz2_off.sh b/scripts/train_256_dsz2_off.sh
new file mode 100644
index 000000000..120bf7547
--- /dev/null
+++ b/scripts/train_256_dsz2_off.sh
@@ -0,0 +1,20 @@
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml \
+    opensora/train/train.py \
+    --model Latte-XL/122 \
+    --dataset sky \
+    --ae stabilityai/sd-vae-ft-mse \
+    --data-path datasets/sky_timelapse/sky_train \
+    --extras 1 \
+    --sample-rate 3 \
+    --num-frames 16 \
+    --max-image-size 256 \
+    --max-train-steps 1000000 \
+    --local-batch-size 5 \
+    --lr 1e-4 \
+    --ckpt-every 500 \
+    --log-every 50 \
+    --gradient-checkpointing \
+    --attention_mode math \
+    --mixed-precision fp16