update

Co-Authored-By: yuan-shenghai <[email protected]>

Author: a-r-r-o-w

examples/cogvideo/train_cogvideox_image_to_video_lora.py

@@ -1461,7 +1461,7 @@ def collate_fn(examples):
                 progress_bar.update(1)
                 global_step += 1
 
-                if accelerator.is_main_process:
+                if accelerator.is_main_process or accelerator.distributed_type == DistributedType.DEEPSPEED:
                     if global_step % args.checkpointing_steps == 0:
                         # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
                         if args.checkpoints_total_limit is not None:
@@ -1494,7 +1494,7 @@ def collate_fn(examples):
             if global_step >= args.max_train_steps:
                 break
 
-        if accelerator.is_main_process or accelerator.distributed_type == DistributedType.DEEPSPEED:
+        if accelerator.is_main_process:
             if args.validation_prompt is not None and (epoch + 1) % args.validation_epochs == 0:
                 # Create pipeline
                 pipe = CogVideoXImageToVideoPipeline.from_pretrained(

examples/cogvideo/train_cogvideox_lora.py

@@ -25,7 +25,7 @@
 import torch
 import torchvision.transforms as TT
 import transformers
-from accelerate import Accelerator
+from accelerate import Accelerator, DistributedType
 from accelerate.logging import get_logger
 from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed
 from huggingface_hub import create_repo, upload_folder
@@ -1211,7 +1211,7 @@ def load_model_hook(models, input_dir):
     )
     use_deepspeed_scheduler = (
         accelerator.state.deepspeed_plugin is not None
-        and "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        and "scheduler" in accelerator.state.deepspeed_plugin.deepspeed_config
     )
 
     optimizer = get_optimizer(args, params_to_optimize, use_deepspeed=use_deepspeed_optimizer)
@@ -1456,7 +1456,7 @@ def collate_fn(examples):
                 progress_bar.update(1)
                 global_step += 1
 
-                if accelerator.is_main_process:
+                if accelerator.is_main_process or accelerator.distributed_type == DistributedType.DEEPSPEED:
                     if global_step % args.checkpointing_steps == 0:
                         # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
                         if args.checkpoints_total_limit is not None: