support DeepSpeed.

finetrainers/state.py

@@ -15,6 +15,7 @@ class State:
     learning_rate: float = None
     train_batch_size: int = None
     generator: torch.Generator = None
+    num_update_steps_per_epoch: int = None
 
     # Hub state
     repo_id: str = None

finetrainers/trainer.py

@@ -1,10 +1,8 @@
-import inspect
 import json
 import logging
 import math
 import os
 import random
-import shutil
 from datetime import timedelta
 from typing import Any, Dict
 from pathlib import Path
@@ -46,10 +44,11 @@ from .dataset import BucketSampler, PrecomputedDataset, VideoDatasetWithResizing
 from .models import get_config_from_model_name
 from .state import State
 from .utils.data_utils import should_perform_precomputation
-from .utils.file_utils import find_files, delete_files, string_to_filename
-from .utils.optimizer_utils import get_optimizer, gradient_norm
+from .utils.file_utils import string_to_filename
+from .utils.optimizer_utils import get_optimizer
 from .utils.memory_utils import get_memory_statistics, free_memory, make_contiguous
 from .utils.torch_utils import unwrap_model, align_device_and_dtype
+from .utils.checkpointing import sort_out_and_load_latest_ckpt_states, save_intermediate_ckpt_states
 
 
 logger = get_logger("finetrainers")
@@ -389,7 +388,13 @@ class Trainer:
         )
         self.transformer.add_adapter(transformer_lora_config)
 
-        # TODO: refactor
+        # Enable TF32 for faster training on Ampere GPUs: https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+        if self.args.allow_tf32 and torch.cuda.is_available():
+            torch.backends.cuda.matmul.allow_tf32 = True
+
+        self.register_saving_loading_hooks(transformer_lora_config)
+
+    def register_saving_loading_hooks(self, transformer_lora_config):
         # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
         def save_model_hook(models, weights, output_dir):
             if self.state.accelerator.is_main_process:
@@ -415,13 +420,20 @@ class Trainer:
                 )
 
         def load_model_hook(models, input_dir):
-            transformer_ = self.model_config["pipeline_cls"].from_pretrained(
-                self.args.pretrained_model_name_or_path, subfolder="transformer"
-            )
-            transformer_.add_adapter(transformer_lora_config)
+            if not self.state.accelerator.distributed_type == DistributedType.DEEPSPEED:
+                while len(models) > 0:
+                    model = models.pop()
+                    if isinstance(unwrap_model(self.state.accelerator, model), type(unwrap_model(self.state.accelerator, self.transformer))):
+                        transformer_ = unwrap_model(self.state.accelerator, model)
+                    else:
+                        raise ValueError(f"Unexpected save model: {unwrap_model(self.state.accelerator, model).__class__}")
+            else:
+                transformer_ = unwrap_model(self.state.accelerator, self.transformer).__class__.from_pretrained(
+                    self.args.pretrained_model_name_or_path, subfolder="transformer"
+                )
+                transformer_.add_adapter(transformer_lora_config)
             
             lora_state_dict = self.model_config["pipeline_cls"].lora_state_dict(input_dir)
-
             transformer_state_dict = {
                 f'{k.replace("transformer.", "")}': v
                 for k, v in lora_state_dict.items()
@@ -447,10 +459,6 @@ class Trainer:
         self.state.accelerator.register_save_state_pre_hook(save_model_hook)
         self.state.accelerator.register_load_state_pre_hook(load_model_hook)
 
-        # Enable TF32 for faster training on Ampere GPUs: https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
-        if self.args.allow_tf32 and torch.cuda.is_available():
-            torch.backends.cuda.matmul.allow_tf32 = True
-
     def prepare_optimizer(self) -> None:
         logger.info("Initializing optimizer and lr scheduler")
 
@@ -519,6 +527,7 @@ class Trainer:
             self.state.train_steps = self.state.train_epochs * num_update_steps_per_epoch
         # Afterwards we recalculate our number of training epochs
         self.state.train_epochs = math.ceil(self.state.train_steps / num_update_steps_per_epoch)
+        self.state.num_update_steps_per_epoch = num_update_steps_per_epoch
 
     def prepare_trackers(self) -> None:
         logger.info("Initializing trackers")
@@ -547,10 +556,18 @@ class Trainer:
         }
         logger.info(f"Training configuration: {json.dumps(info, indent=4)}")
 
-        # TODO(aryan): handle resume from checkpoint
         global_step = 0
         first_epoch = 0
         initial_global_step = 0
+        
+        # Potentially load in the weights and states from a previous save
+        initial_global_step, global_step, first_epoch = sort_out_and_load_latest_ckpt_states(
+            accelerator=self.state.accelerator, 
+            resume_from_checkpoint=self.args.resume_from_checkpoint, 
+            num_update_steps_per_epoch=self.state.num_update_steps_per_epoch, 
+            output_dir=self.args.output_dir
+        )
+
         progress_bar = tqdm(
             range(0, self.state.train_steps),
             initial=initial_global_step,
@@ -682,20 +699,9 @@ class Trainer:
                     # Checkpointing
                     if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
                         if global_step % self.args.checkpointing_steps == 0:
-                            # before saving state, check if this save would set us over the `checkpointing_limit`
-                            if self.args.checkpointing_limit is not None:
-                                checkpoints = find_files(self.args.output_dir, prefix="checkpoint")
-
-                                # before we save the new checkpoint, we need to have at_most `checkpoints_total_limit - 1` checkpoints
-                                if len(checkpoints) >= self.args.checkpointing_limit:
-                                    num_to_remove = len(checkpoints) - self.args.checkpointing_limit + 1
-                                    checkpoints_to_remove = checkpoints[0:num_to_remove]
-                                    delete_files(checkpoints_to_remove)
-
-                            logger.info(f"Checkpointing at step {global_step}")
-                            save_path = os.path.join(self.args.output_dir, f"checkpoint-{global_step}")
-                            accelerator.save_state(save_path)
-                            logger.info(f"Saved state to {save_path}")
+                            save_intermediate_ckpt_states(
+                                accelerator=accelerator, checkpointing_limit=self.args.checkpointing_limit, step=global_step, output_dir=self.args.output_dir
+                            )
 
                 # Maybe run validation
                 should_run_validation = (
@@ -725,6 +731,7 @@ class Trainer:
 
         accelerator.wait_for_everyone()
         if accelerator.is_main_process:
+            # TODO: consider factoring this out when supporting other types of training algos.
             self.transformer = unwrap_model(accelerator, self.transformer)
             dtype = (
                 torch.float16
@@ -741,6 +748,13 @@ class Trainer:
                 transformer_lora_layers=transformer_lora_layers,
             )
 
+            if self.args.push_to_hub:
+                upload_folder(
+                    repo_id=self.state.repo_id, folder_path=self.args.output_dir, ignore_patterns=["checkpoint-*"]
+                )
+            
+            self.validate(step=global_step, final_validation=True)
+
         del self.tokenizer, self.text_encoder, self.transformer, self.vae, self.scheduler
         free_memory()
         memory_statistics = get_memory_statistics()
@@ -748,7 +762,7 @@ class Trainer:
 
         accelerator.end_training()
 
-    def validate(self, step: int) -> None:
+    def validate(self, step: int, final_validation: bool = False) -> None:
         logger.info("Starting validation")
 
         accelerator = self.state.accelerator
@@ -763,21 +777,34 @@ class Trainer:
         memory_statistics = get_memory_statistics()
         logger.info(f"Memory before validation start: {json.dumps(memory_statistics, indent=4)}")
 
-        pipeline = self.model_config["initialize_pipeline"](
-            model_id=self.args.pretrained_model_name_or_path,
-            tokenizer=self.tokenizer,
-            text_encoder=self.text_encoder,
-            tokenizer_2=self.tokenizer_2,
-            text_encoder_2=self.text_encoder_2,
-            transformer=unwrap_model(accelerator, self.transformer),
-            vae=self.vae,
-            device=accelerator.device,
-            revision=self.args.revision,
-            cache_dir=self.args.cache_dir,
-            enable_slicing=self.args.enable_slicing,
-            enable_tiling=self.args.enable_tiling,
-            enable_model_cpu_offload=self.args.enable_model_cpu_offload,
-        )
+        if not final_validation:
+            pipeline = self.model_config["initialize_pipeline"](
+                model_id=self.args.pretrained_model_name_or_path,
+                tokenizer=self.tokenizer,
+                text_encoder=self.text_encoder,
+                tokenizer_2=self.tokenizer_2,
+                text_encoder_2=self.text_encoder_2,
+                transformer=unwrap_model(accelerator, self.transformer),
+                vae=self.vae,
+                device=accelerator.device,
+                revision=self.args.revision,
+                cache_dir=self.args.cache_dir,
+                enable_slicing=self.args.enable_slicing,
+                enable_tiling=self.args.enable_tiling,
+                enable_model_cpu_offload=self.args.enable_model_cpu_offload,
+            )
+        else:
+            # `torch_dtype` is manually set within `initialize_pipeline()`.
+            pipeline = self.model_config["initialize_pipeline"](
+                model_id=self.args.pretrained_model_name_or_path, 
+                device=accelerator.device,
+                revision=self.args.revision,
+                cache_dir=self.args.cache_dir,
+                enable_slicing=self.args.enable_slicing,
+                enable_tiling=self.args.enable_tiling,
+                enable_model_cpu_offload=self.args.enable_model_cpu_offload,
+            )
+            pipeline.load_lora_weights(self.args.output_dir)
 
         all_processes_artifacts = []
         for i in range(num_validation_samples):
@@ -811,6 +838,7 @@ class Trainer:
                 num_frames=num_frames,
                 num_videos_per_prompt=self.args.num_validation_videos_per_prompt,
                 generator=self.state.generator,
+                # todo support passing `fps` for supported pipelines.
             )
 
             # Remove all hooks that might have been added during pipeline initialization to the models
@@ -845,6 +873,7 @@ class Trainer:
                     artifact_value = wandb.Image(filename)
                 elif artifact_type == "video":
                     logger.debug(f"Saving video to {filename}")
+                    # TODO: this should be configurable here as well as in validation runs where we call the pipeline that has `fps`.
                     export_to_video(artifact_value, filename, fps=15) 
                     artifact_value = wandb.Video(filename, caption=prompt)
 
@@ -853,9 +882,10 @@ class Trainer:
         all_artifacts = gather_object(all_processes_artifacts)
 
         if accelerator.is_main_process:
+            tracker_key = "final" if final_validation else "validation"
             for tracker in accelerator.trackers:
                 if tracker.name == "wandb":
-                    tracker.log({"validation": all_artifacts}, step=step)
+                    tracker.log({tracker_key: all_artifacts}, step=step)
 
         accelerator.wait_for_everyone()
 
@@ -864,11 +894,11 @@ class Trainer:
         logger.info(f"Memory after validation end: {json.dumps(memory_statistics, indent=4)}")
         torch.cuda.reset_peak_memory_stats(accelerator.device)
 
-        self.transformer.train()
+        if not final_validation:
+            self.transformer.train()
 
     def evaluate(self) -> None:
-        logger.info("Starting evaluation")
-        # TODO: implement metrics for evaluation
+        raise NotImplementedError
 
     def _init_distributed(self) -> None:
         logging_dir = Path(self.args.output_dir, self.args.logging_dir)

finetrainers/utils/checkpointing.py

@@ -0,0 +1,57 @@
+import os
+from accelerate.logging import get_logger
+from ..constants import FINETRAINERS_LOG_LEVEL
+from ..utils.file_utils import find_files, delete_files
+
+logger = get_logger("finetrainers")
+logger.setLevel(FINETRAINERS_LOG_LEVEL)
+
+def sort_out_and_load_latest_ckpt_states(
+        accelerator, resume_from_checkpoint, num_update_steps_per_epoch, output_dir
+    ):
+    if not resume_from_checkpoint:
+        initial_global_step = 0
+        global_step = 0
+        first_epoch = 0
+    else:
+        if resume_from_checkpoint != "latest":
+            path = os.path.basename(resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = os.listdir(output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+
+        if path is None:
+            logger.info(
+                f"Checkpoint '{resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+            resume_from_checkpoint = None
+            initial_global_step = 0
+        else:
+            logger.info(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(output_dir, path))
+            global_step = int(path.split("-")[1])
+
+            initial_global_step = global_step
+            first_epoch = global_step // num_update_steps_per_epoch
+
+    return initial_global_step, global_step, first_epoch
+
+
+def save_intermediate_ckpt_states(accelerator, checkpointing_limit, step, output_dir):
+    # before saving state, check if this save would set us over the `checkpointing_limit`
+    if checkpointing_limit is not None:
+        checkpoints = find_files(output_dir, prefix="checkpoint")
+
+        # before we save the new checkpoint, we need to have at_most `checkpoints_total_limit - 1` checkpoints
+        if len(checkpoints) >= checkpointing_limit:
+            num_to_remove = len(checkpoints) - checkpointing_limit + 1
+            checkpoints_to_remove = checkpoints[0:num_to_remove]
+            delete_files(checkpoints_to_remove)
+
+    logger.info(f"Checkpointing at step {step}")
+    save_path = os.path.join(output_dir, f"checkpoint-{step}")
+    accelerator.save_state(save_path)
+    logger.info(f"Saved state to {save_path}")

train.py

@@ -33,7 +33,7 @@ def main():
         trainer.prepare_for_training()
         trainer.prepare_trackers()
         trainer.train()
-        trainer.evaluate()
+        # trainer.evaluate()
 
     except KeyboardInterrupt:
         logger.info("Received keyboard interrupt. Exiting...")