fix: resuming from a checkpoint when using deepspeed. (#38)

* fix: resuming from a checkpoint when using deepspeed.

* remove changes to prepare_dataset.py

* propagate to others.

* tackle gradnorm.

---------

Co-authored-by: Aryan <aryan@huggingface.co>

Makefile

@@ -4,7 +4,7 @@ check_dirs := training tests
 
 quality:
 	ruff check $(check_dirs)
-	ruff format --check $(check_dirs) setup.py
+	ruff format --check $(check_dirs)
 
 style:
 	ruff check $(check_dirs) --fix

training/cogvideox_image_to_video_lora.py

@@ -26,6 +26,7 @@ from typing import Any, Dict
 import diffusers
 import torch
 import transformers
+import wandb
 from accelerate import Accelerator, DistributedType
 from accelerate.logging import get_logger
 from accelerate.utils import (
@@ -52,8 +53,6 @@ from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import AutoTokenizer, T5EncoderModel
 
-import wandb
-
 
 from args import get_args  # isort:skip
 from dataset import BucketSampler, VideoDatasetWithResizing, VideoDatasetWithResizeAndRectangleCrop  # isort:skip
@@ -385,13 +384,15 @@ def main(args):
             transformer_lora_layers_to_save = None
 
             for model in models:
-                if isinstance(model, type(unwrap_model(transformer))):
+                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
+                    model = unwrap_model(model)
                     transformer_lora_layers_to_save = get_peft_model_state_dict(model)
                 else:
                     raise ValueError(f"unexpected save model: {model.__class__}")
 
                 # make sure to pop weight so that corresponding model is not saved again
-                weights.pop()
+                if weights:
+                    weights.pop()
 
             CogVideoXImageToVideoPipeline.save_lora_weights(
                 output_dir,
@@ -401,13 +402,20 @@ def main(args):
     def load_model_hook(models, input_dir):
         transformer_ = None
 
-        while len(models) > 0:
+        # This is a bit of a hack but I don't know any other solution.
-            model = models.pop()
+        if not accelerator.distributed_type == DistributedType.DEEPSPEED:
+            while len(models) > 0:
+                model = models.pop()
 
-            if isinstance(model, type(unwrap_model(transformer))):
+                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
-                transformer_ = model
+                    transformer_ = unwrap_model(model)
-            else:
+                else:
-                raise ValueError(f"Unexpected save model: {model.__class__}")
+                    raise ValueError(f"Unexpected save model: {unwrap_model(model).__class__}")
+        else:
+            transformer_ = CogVideoXTransformer3DModel.from_pretrained(
+                args.pretrained_model_name_or_path, subfolder="transformer"
+            )
+            transformer_.add_adapter(transformer_lora_config)
 
         lora_state_dict = CogVideoXImageToVideoPipeline.lora_state_dict(input_dir)
 
@@ -795,12 +803,15 @@ def main(args):
                         logger.info(f"Saved state to {save_path}")
 
             last_lr = lr_scheduler.get_last_lr()[0] if lr_scheduler is not None else args.learning_rate
-            logs = {
+            logs = {"loss": loss.detach().item(), "lr": last_lr}
-                "loss": loss.detach().item(),
+            # gradnorm + deepspeed: https://github.com/microsoft/DeepSpeed/issues/4555
-                "lr": last_lr,
+            if accelerator.distributed_type != DistributedType.DEEPSPEED:
-                "gradient_norm_before_clip": gradient_norm_before_clip,
+                logs.update(
-                "gradient_norm_after_clip": gradient_norm_after_clip,
+                    {
-            }
+                        "gradient_norm_before_clip": gradient_norm_before_clip,
+                        "gradient_norm_after_clip": gradient_norm_after_clip,
+                    }
+                )
             progress_bar.set_postfix(**logs)
             accelerator.log(logs, step=global_step)
 

training/cogvideox_text_to_video_lora.py

@@ -25,6 +25,7 @@ from typing import Any, Dict
 import diffusers
 import torch
 import transformers
+import wandb
 from accelerate import Accelerator, DistributedType
 from accelerate.logging import get_logger
 from accelerate.utils import (
@@ -51,8 +52,6 @@ from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import AutoTokenizer, T5EncoderModel
 
-import wandb
-
 
 from args import get_args  # isort:skip
 from dataset import BucketSampler, VideoDatasetWithResizing, VideoDatasetWithResizeAndRectangleCrop  # isort:skip
@@ -358,13 +357,15 @@ def main(args):
             transformer_lora_layers_to_save = None
 
             for model in models:
-                if isinstance(model, type(unwrap_model(transformer))):
+                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
+                    model = unwrap_model(model)
                     transformer_lora_layers_to_save = get_peft_model_state_dict(model)
                 else:
                     raise ValueError(f"unexpected save model: {model.__class__}")
 
                 # make sure to pop weight so that corresponding model is not saved again
-                weights.pop()
+                if weights:
+                    weights.pop()
 
             CogVideoXPipeline.save_lora_weights(
                 output_dir,
@@ -374,13 +375,20 @@ def main(args):
     def load_model_hook(models, input_dir):
         transformer_ = None
 
-        while len(models) > 0:
+        # This is a bit of a hack but I don't know any other solution.
-            model = models.pop()
+        if not accelerator.distributed_type == DistributedType.DEEPSPEED:
+            while len(models) > 0:
+                model = models.pop()
 
-            if isinstance(model, type(unwrap_model(transformer))):
+                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
-                transformer_ = model
+                    transformer_ = unwrap_model(model)
-            else:
+                else:
-                raise ValueError(f"Unexpected save model: {model.__class__}")
+                    raise ValueError(f"Unexpected save model: {unwrap_model(model).__class__}")
+        else:
+            transformer_ = CogVideoXTransformer3DModel.from_pretrained(
+                args.pretrained_model_name_or_path, subfolder="transformer"
+            )
+            transformer_.add_adapter(transformer_lora_config)
 
         lora_state_dict = CogVideoXPipeline.lora_state_dict(input_dir)
 
@@ -553,7 +561,7 @@ def main(args):
 
     # We need to initialize the trackers we use, and also store our configuration.
     # The trackers initializes automatically on the main process.
-    if accelerator.is_main_process:
+    if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
         tracker_name = args.tracker_name or "cogvideox-lora"
         accelerator.init_trackers(tracker_name, config=vars(args))
 
@@ -731,7 +739,7 @@ def main(args):
                 progress_bar.update(1)
                 global_step += 1
 
-                if accelerator.is_main_process or accelerator.distributed_type == DistributedType.DEEPSPEED:
+                if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
                     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:
@@ -758,12 +766,15 @@ def main(args):
                         logger.info(f"Saved state to {save_path}")
 
             last_lr = lr_scheduler.get_last_lr()[0] if lr_scheduler is not None else args.learning_rate
-            logs = {
+            logs = {"loss": loss.detach().item(), "lr": last_lr}
-                "loss": loss.detach().item(),
+            # gradnorm + deepspeed: https://github.com/microsoft/DeepSpeed/issues/4555
-                "lr": last_lr,
+            if accelerator.distributed_type != DistributedType.DEEPSPEED:
-                "gradient_norm_before_clip": gradient_norm_before_clip,
+                logs.update(
-                "gradient_norm_after_clip": gradient_norm_after_clip,
+                    {
-            }
+                        "gradient_norm_before_clip": gradient_norm_before_clip,
+                        "gradient_norm_after_clip": gradient_norm_after_clip,
+                    }
+                )
             progress_bar.set_postfix(**logs)
             accelerator.log(logs, step=global_step)
 

training/cogvideox_text_to_video_sft.py

@@ -25,7 +25,8 @@ from typing import Any, Dict
 import diffusers
 import torch
 import transformers
-from accelerate import Accelerator, DistributedType
+import wandb
+from accelerate import Accelerator, DistributedType, init_empty_weights
 from accelerate.logging import get_logger
 from accelerate.utils import (
     DistributedDataParallelKwargs,
@@ -50,8 +51,6 @@ from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import AutoTokenizer, T5EncoderModel
 
-import wandb
-
 
 from args import get_args  # isort:skip
 from dataset import BucketSampler, VideoDatasetWithResizing, VideoDatasetWithResizeAndRectangleCrop  # isort:skip
@@ -336,8 +335,9 @@ def main(args):
     def save_model_hook(models, weights, output_dir):
         if accelerator.is_main_process:
             for model in models:
-                if isinstance(model, type(unwrap_model(transformer))):
+                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
                     model: CogVideoXTransformer3DModel
+                    model = unwrap_model(model)
                     model.save_pretrained(
                         os.path.join(output_dir, "transformer"), safe_serialization=True, max_shard_size="5GB"
                     )
@@ -345,22 +345,32 @@ def main(args):
                     raise ValueError(f"Unexpected save model: {model.__class__}")
 
                 # make sure to pop weight so that corresponding model is not saved again
-                weights.pop()
+                if weights:
+                    weights.pop()
 
     def load_model_hook(models, input_dir):
         transformer_ = None
+        init_under_meta = False
 
-        while len(models) > 0:
+        # This is a bit of a hack but I don't know any other solution.
-            model = models.pop()
+        if not accelerator.distributed_type == DistributedType.DEEPSPEED:
+            while len(models) > 0:
+                model = models.pop()
 
-            if isinstance(model, type(unwrap_model(transformer))):
+                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
-                transformer_: CogVideoXTransformer3DModel = model
+                    transformer_ = unwrap_model(model)
-            else:
+                else:
-                raise ValueError(f"Unexpected save model: {model.__class__.__name__}")
+                    raise ValueError(f"Unexpected save model: {unwrap_model(model).__class__}")
+        else:
+            with init_empty_weights():
+                transformer_ = CogVideoXTransformer3DModel.from_config(
+                    args.pretrained_model_name_or_path, subfolder="transformer"
+                )
+                init_under_meta = True
 
         load_model = CogVideoXTransformer3DModel.from_pretrained(os.path.join(input_dir, "transformer"))
         transformer_.register_to_config(**load_model.config)
-        transformer_.load_state_dict(load_model.state_dict())
+        transformer_.load_state_dict(load_model.state_dict(), assign=init_under_meta)
         del load_model
 
         # Make sure the trainable params are in float32. This is again needed since the base models
@@ -722,12 +732,15 @@ def main(args):
                         logger.info(f"Saved state to {save_path}")
 
             last_lr = lr_scheduler.get_last_lr()[0] if lr_scheduler is not None else args.learning_rate
-            logs = {
+            logs = {"loss": loss.detach().item(), "lr": last_lr}
-                "loss": loss.detach().item(),
+            # gradnorm + deepspeed: https://github.com/microsoft/DeepSpeed/issues/4555
-                "lr": last_lr,
+            if accelerator.distributed_type != DistributedType.DEEPSPEED:
-                "gradient_norm_before_clip": gradient_norm_before_clip,
+                logs.update(
-                "gradient_norm_after_clip": gradient_norm_after_clip,
+                    {
-            }
+                        "gradient_norm_before_clip": gradient_norm_before_clip,
+                        "gradient_norm_after_clip": gradient_norm_after_clip,
+                    }
+                )
             progress_bar.set_postfix(**logs)
             accelerator.log(logs, step=global_step)