Support deepspeed zero2 and zero2_offload for latte training.

Former-commit-id: c87e95dd357e6fa6d6fb2dcfee781bc49c039e13

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

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

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

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

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

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