[Training] Add distributed checkpointing (#458)

Author: kevin314

fastvideo/v1/fastvideo_args.py

@@ -472,7 +472,7 @@ class TrainingArgs(FastVideoArgs):
     validation_steps: float = 0.0
     log_validation: bool = False
     tracker_project_name: str = ""
-    # seed: int
+    seed: Optional[int] = None
 
     # output
     output_dir: str = ""
@@ -630,6 +630,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parser.add_argument("--tracker-project-name",
                             type=str,
                             help="Project name for tracking")
+        parser.add_argument("--seed",
+                            type=int,
+                            help="Seed for deterministic training")
 
         # Output configuration
         parser.add_argument("--output-dir",

fastvideo/v1/training/checkpointing_utils.py

@@ -0,0 +1,107 @@
+import random
+from typing import Any, Dict, Optional
+
+import numpy as np
+import torch
+import torch.distributed.checkpoint.stateful
+from torch.distributed.checkpoint.state_dict import (StateDictOptions,
+                                                     get_model_state_dict,
+                                                     get_optimizer_state_dict,
+                                                     set_model_state_dict,
+                                                     set_optimizer_state_dict)
+
+
+class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful):
+
+    def __init__(self, model: torch.nn.Module) -> None:
+        self.model = model
+
+    def state_dict(self) -> Dict[str, Any]:
+        return get_model_state_dict(self.model)  # type: ignore[no-any-return]
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        set_model_state_dict(
+            self.model,
+            model_state_dict=state_dict,
+            options=StateDictOptions(strict=False),
+        )
+
+
+class OptimizerWrapper(torch.distributed.checkpoint.stateful.Stateful):
+
+    def __init__(self, model: torch.nn.Module,
+                 optimizer: torch.optim.Optimizer) -> None:
+        self.model = model
+        self.optimizer = optimizer
+
+    def state_dict(self) -> Dict[str, Any]:
+        return get_optimizer_state_dict(  # type: ignore[no-any-return]
+            self.model,
+            self.optimizer,
+            options=StateDictOptions(flatten_optimizer_state_dict=True),
+        )
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        set_optimizer_state_dict(
+            self.model,
+            self.optimizer,
+            optim_state_dict=state_dict,
+            options=StateDictOptions(flatten_optimizer_state_dict=True),
+        )
+
+
+class SchedulerWrapper(torch.distributed.checkpoint.stateful.Stateful):
+
+    def __init__(self, scheduler) -> None:
+        self.scheduler = scheduler
+
+    def state_dict(self) -> Dict[str, Any]:
+        return {"scheduler": self.scheduler.state_dict()}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        self.scheduler.load_state_dict(state_dict["scheduler"])
+
+
+class RandomStateWrapper(torch.distributed.checkpoint.stateful.Stateful):
+
+    def __init__(self,
+                 noise_generator: Optional[torch.Generator] = None) -> None:
+        self.noise_generator = noise_generator
+
+    def state_dict(self) -> Dict[str, Any]:
+        state = {
+            "torch_rng_state": torch.get_rng_state(),
+            "numpy_rng_state": np.random.get_state(),
+            "python_rng_state": random.getstate(),
+        }
+
+        if torch.cuda.is_available():
+            state["cuda_rng_state"] = torch.cuda.get_rng_state()
+            if torch.cuda.device_count() > 1:
+                state["cuda_rng_state_all"] = torch.cuda.get_rng_state_all()
+
+        if self.noise_generator is not None:
+            state["noise_generator_state"] = self.noise_generator.get_state()
+
+        return state
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        if "torch_rng_state" in state_dict:
+            torch.set_rng_state(state_dict["torch_rng_state"])
+
+        if "numpy_rng_state" in state_dict:
+            np.random.set_state(state_dict["numpy_rng_state"])
+
+        if "python_rng_state" in state_dict:
+            random.setstate(state_dict["python_rng_state"])
+
+        # Restore CUDA random state
+        if torch.cuda.is_available():
+            if "cuda_rng_state" in state_dict:
+                torch.cuda.set_rng_state(state_dict["cuda_rng_state"])
+            if "cuda_rng_state_all" in state_dict:
+                torch.cuda.set_rng_state_all(state_dict["cuda_rng_state_all"])
+
+        # Restore noise generator state
+        if "noise_generator_state" in state_dict and self.noise_generator is not None:
+            self.noise_generator.set_state(state_dict["noise_generator_state"])

fastvideo/v1/training/training_pipeline.py

@@ -135,6 +135,13 @@ def log_validation(self, transformer, training_args, global_step) -> None:
         # Create sampling parameters if not provided
         sampling_param = SamplingParam.from_pretrained(training_args.model_path)
 
+        # Set deterministic seed for validation
+        validation_seed = training_args.seed if training_args.seed is not None else 42
+        torch.manual_seed(validation_seed)
+        torch.cuda.manual_seed_all(validation_seed)
+
+        logger.info("Using validation seed: %s", validation_seed)
+
         # Prepare validation prompts
         logger.info('fastvideo_args.validation_prompt_dir: %s',
                     training_args.validation_prompt_dir)
@@ -192,7 +199,7 @@ def log_validation(self, transformer, training_args, global_step) -> None:
             batch = ForwardBatch(
                 data_type="video",
                 latents=None,
-                # seed=sampling_param.seed,
+                seed=validation_seed,  # Use deterministic seed
                 prompt_embeds=[prompt_embeds],
                 prompt_attention_mask=[prompt_attention_mask],
                 # make sure we use the same height, width, and num_frames as the training pipeline
@@ -206,7 +213,6 @@ def log_validation(self, transformer, training_args, global_step) -> None:
                 n_tokens=n_tokens,
                 do_classifier_free_guidance=False,
                 eta=0.0,
-                extra={},
             )
 
             # Run validation inference

fastvideo/v1/training/training_utils.py

@@ -1,21 +1,57 @@
 import json
 import math
 import os
-from typing import List, Optional, Tuple, Union
+import time
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.distributed as dist
-from torch.distributed.fsdp import FullStateDictConfig
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp import StateDictType
+import torch.distributed.checkpoint as dcp
+import torch.distributed.checkpoint.stateful
+from safetensors.torch import save_file
 
 from fastvideo.v1.logger import init_logger
+from fastvideo.v1.training.checkpointing_utils import (ModelWrapper,
+                                                       OptimizerWrapper,
+                                                       RandomStateWrapper,
+                                                       SchedulerWrapper)
 
 logger = init_logger(__name__)
 
 _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False
 
 
+def gather_state_dict_on_cpu_rank0(
+    model,
+    device: Optional[torch.device] = None,
+) -> Dict[str, Any]:
+    rank = dist.get_rank()
+    cpu_state_dict = {}
+    sharded_sd = model.state_dict()
+    for param_name, param in sharded_sd.items():
+        if hasattr(param, "_local_tensor"):
+            # DTensor case
+            if param.is_cpu:
+                # Gather directly on CPU
+                param = param.full_tensor()
+            else:
+                if device is not None:
+                    param = param.to(device)
+                param = param.full_tensor()
+        else:
+            # Regular tensor case
+            if param.is_cpu:
+                pass
+            else:
+                if device is not None:
+                    param = param.to(device)
+
+        if rank == 0:
+            cpu_state_dict[param_name] = param.cpu()
+
+    return cpu_state_dict
+
+
 def compute_density_for_timestep_sampling(
     weighting_scheme: str,
     batch_size: int,
@@ -66,24 +102,67 @@ def get_sigmas(noise_scheduler,
     return sigma
 
 
-def save_checkpoint(transformer, rank, output_dir, step) -> None:
-    # Configure FSDP to save full state dict
-    FSDP.set_state_dict_type(
-        transformer,
-        state_dict_type=StateDictType.FULL_STATE_DICT,
-        state_dict_config=FullStateDictConfig(offload_to_cpu=True,
-                                              rank0_only=True),
-    )
-
-    # Now get the state dict
-    cpu_state = transformer.state_dict()
-
-    # Save it (only on rank 0 since we used rank0_only=True)
-    if rank <= 0:
-        save_dir = os.path.join(output_dir, f"checkpoint-{step}")
-        os.makedirs(save_dir, exist_ok=True)
-        weight_path = os.path.join(save_dir, "diffusion_pytorch_model.pt")
-        torch.save(cpu_state, weight_path)
+def save_checkpoint(transformer,
+                    rank,
+                    output_dir,
+                    step,
+                    optimizer=None,
+                    dataloader=None,
+                    scheduler=None,
+                    noise_generator=None) -> None:
+    """
+    Save checkpoint following finetrainer's distributed checkpoint approach.
+    Saves both distributed checkpoint and consolidated model weights.
+    """
+    save_dir = os.path.join(output_dir, f"checkpoint-{step}")
+    os.makedirs(save_dir, exist_ok=True)
+
+    states = {
+        "model": ModelWrapper(transformer),
+        "random_state": RandomStateWrapper(noise_generator),
+    }
+
+    if optimizer is not None:
+        states["optimizer"] = OptimizerWrapper(transformer, optimizer)
+
+    if dataloader is not None:
+        states["dataloader"] = dataloader
+
+    if scheduler is not None:
+        states["scheduler"] = SchedulerWrapper(scheduler)
+
+    dcp_dir = os.path.join(save_dir, "distributed_checkpoint")
+    logger.info("rank: %s, saving distributed checkpoint to %s",
+                rank,
+                dcp_dir,
+                local_main_process_only=False)
+
+    begin_time = time.perf_counter()
+    dcp.save(states, checkpoint_id=dcp_dir)
+    end_time = time.perf_counter()
+
+    logger.info("rank: %s, distributed checkpoint saved in %.2f seconds",
+                rank,
+                end_time - begin_time,
+                local_main_process_only=False)
+
+    cpu_state = gather_state_dict_on_cpu_rank0(transformer, device=None)
+
+    if rank == 0:
+        # Save model weights (consolidated)
+        weight_path = os.path.join(save_dir,
+                                   "diffusion_pytorch_model.safetensors")
+        logger.info("rank: %s, saving consolidated checkpoint to %s",
+                    rank,
+                    weight_path,
+                    local_main_process_only=False)
+        save_file(cpu_state, weight_path)
+        logger.info("rank: %s, consolidated checkpoint saved to %s",
+                    rank,
+                    weight_path,
+                    local_main_process_only=False)
+
+        # Save model config
         config_dict = transformer.hf_config
         if "dtype" in config_dict:
             del config_dict["dtype"]  # TODO
@@ -94,6 +173,66 @@ def save_checkpoint(transformer, rank, output_dir, step) -> None:
         logger.info("--> checkpoint saved at step %s to %s", step, weight_path)
 
 
+def load_checkpoint(transformer,
+                    rank,
+                    checkpoint_path,
+                    optimizer=None,
+                    dataloader=None,
+                    scheduler=None,
+                    noise_generator=None) -> int:
+    """
+    Load checkpoint following finetrainer's distributed checkpoint approach.
+    Returns the step number from which training should resume.
+    """
+    if not os.path.exists(checkpoint_path):
+        logger.warning("Checkpoint path %s does not exist", checkpoint_path)
+        return 0
+
+    # Extract step number from checkpoint path
+    step = int(os.path.basename(checkpoint_path).split('-')[-1])
+
+    if rank == 0:
+        logger.info("Loading checkpoint from step %s", step)
+
+    dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint")
+
+    if not os.path.exists(dcp_dir):
+        logger.warning("Distributed checkpoint directory %s does not exist",
+                       dcp_dir)
+        return 0
+
+    states = {
+        "model": ModelWrapper(transformer),
+        "random_state": RandomStateWrapper(noise_generator),
+    }
+
+    if optimizer is not None:
+        states["optimizer"] = OptimizerWrapper(transformer, optimizer)
+
+    if dataloader is not None:
+        states["dataloader"] = dataloader
+
+    if scheduler is not None:
+        states["scheduler"] = SchedulerWrapper(scheduler)
+
+    logger.info("rank: %s, loading distributed checkpoint from %s",
+                rank,
+                dcp_dir,
+                local_main_process_only=False)
+
+    begin_time = time.perf_counter()
+    dcp.load(states, checkpoint_id=dcp_dir)
+    end_time = time.perf_counter()
+
+    logger.info("rank: %s, distributed checkpoint loaded in %.2f seconds",
+                rank,
+                end_time - begin_time,
+                local_main_process_only=False)
+    logger.info("--> checkpoint loaded from step %s", step)
+
+    return step
+
+
 def normalize_dit_input(model_type, latents, args=None) -> torch.Tensor:
     if model_type == "hunyuan_hf" or model_type == "hunyuan":
         return latents * 0.476986