diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index fcb9030f0..20c1cc2ff 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/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 @@ class TrainingArgs(FastVideoArgs): 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", diff --git a/fastvideo/v1/training/checkpointing_utils.py b/fastvideo/v1/training/checkpointing_utils.py new file mode 100644 index 000000000..4edb48727 --- /dev/null +++ b/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"]) diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index 059704dd9..69d2f74ee 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -135,6 +135,13 @@ class TrainingPipeline(ComposedPipelineBase, ABC): # 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 @@ class TrainingPipeline(ComposedPipelineBase, ABC): 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 @@ class TrainingPipeline(ComposedPipelineBase, ABC): n_tokens=n_tokens, do_classifier_free_guidance=False, eta=0.0, - extra={}, ) # Run validation inference diff --git a/fastvideo/v1/training/training_utils.py b/fastvideo/v1/training/training_utils.py index df74567ce..6ccf94b6c 100644 --- a/fastvideo/v1/training/training_utils.py +++ b/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), - ) +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) - # Now get the state dict - cpu_state = transformer.state_dict() + states = { + "model": ModelWrapper(transformer), + "random_state": RandomStateWrapper(noise_generator), + } - # 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) + 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 diff --git a/fastvideo/v1/training/wan_training_pipeline.py b/fastvideo/v1/training/wan_training_pipeline.py index 46442efda..0bbf1ca75 100644 --- a/fastvideo/v1/training/wan_training_pipeline.py +++ b/fastvideo/v1/training/wan_training_pipeline.py @@ -1,8 +1,10 @@ +import random import sys import time from collections import deque from copy import deepcopy +import numpy as np import torch from diffusers import FlowMatchEulerDiscreteScheduler from tqdm.auto import tqdm @@ -18,8 +20,8 @@ from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline from fastvideo.v1.training.training_pipeline import TrainingPipeline from fastvideo.v1.training.training_utils import ( clip_grad_norm_while_handling_failing_dtensor_cases, - compute_density_for_timestep_sampling, get_sigmas, normalize_dit_input, - save_checkpoint) + compute_density_for_timestep_sampling, get_sigmas, load_checkpoint, + normalize_dit_input, save_checkpoint) import wandb # isort: skip @@ -172,7 +174,19 @@ class WanTrainingPipeline(TrainingPipeline): fastvideo_args: FastVideoArgs, ): assert self.training_args is not None - noise_random_generator = None + + # Set random seeds for deterministic training + seed = self.training_args.seed if self.training_args.seed is not None else 42 + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + noise_random_generator = torch.Generator(device="cpu") + noise_random_generator.manual_seed(seed) + + logger.info("Initialized random seeds with seed: %s", seed) noise_scheduler = FlowMatchEulerDiscreteScheduler() @@ -187,7 +201,6 @@ class WanTrainingPipeline(TrainingPipeline): # logger.info(f" Num examples = {len(train_dataset)}") # logger.info(f" Dataloader size = {len(train_dataloader)}") # logger.info(f" Num Epochs = {args.num_train_epochs}") - logger.info(" Resume training from step %s", self.init_steps) logger.info(" Instantaneous batch size per device = %s", self.training_args.train_batch_size) logger.info( @@ -206,11 +219,21 @@ class WanTrainingPipeline(TrainingPipeline): logger.info(" Master weight dtype: %s", self.transformer.parameters().__next__().dtype) - # Potentially load in the weights and states from a previous save if self.training_args.resume_from_checkpoint: - assert NotImplementedError( - "resume_from_checkpoint is not supported now.") - # TODO + logger.info("Loading checkpoint from %s", + self.training_args.resume_from_checkpoint) + resumed_step = load_checkpoint( + self.transformer, self.rank, + self.training_args.resume_from_checkpoint, self.optimizer, + self.train_dataloader, self.lr_scheduler, + noise_random_generator) + if resumed_step > 0: + self.init_steps = resumed_step + logger.info("Successfully resumed from step %s", resumed_step) + else: + logger.warning( + "Failed to load checkpoint, starting from step 0") + self.init_steps = 0 progress_bar = tqdm( range(0, self.training_args.max_train_steps), @@ -285,9 +308,10 @@ class WanTrainingPipeline(TrainingPipeline): step=step, ) if step % self.training_args.checkpointing_steps == 0: - # Your existing checkpoint saving code save_checkpoint(self.transformer, self.rank, - self.training_args.output_dir, step) + self.training_args.output_dir, step, + self.optimizer, self.train_dataloader, + self.lr_scheduler, noise_random_generator) self.transformer.train() self.sp_group.barrier() if self.training_args.log_validation and step % self.training_args.validation_steps == 0: @@ -295,7 +319,9 @@ class WanTrainingPipeline(TrainingPipeline): save_checkpoint(self.transformer, self.rank, self.training_args.output_dir, - self.training_args.max_train_steps) + self.training_args.max_train_steps, self.optimizer, + self.train_dataloader, self.lr_scheduler, + noise_random_generator) if get_sp_group(): cleanup_dist_env_and_memory() diff --git a/scripts/finetune/finetune_v1.sh b/scripts/finetune/finetune_v1.sh index ff88be3f5..c5accdd5d 100644 --- a/scripts/finetune/finetune_v1.sh +++ b/scripts/finetune/finetune_v1.sh @@ -8,6 +8,8 @@ NUM_GPUS=1 # export CUDA_VISIBLE_DEVICES=4,5 # IP=[MASTER NODE IP] +CHECKPOINT_PATH="$DATA_DIR/outputs/wan_finetune/checkpoint-5" + # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\ fastvideo/v1/training/wan_training_pipeline.py\ @@ -47,4 +49,6 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\ --weight_decay 0.01 \ --not_apply_cfg_solver \ --master_weight_type "fp32" \ - --max_grad_norm 1.0 + --max_grad_norm 1.0 \ + +# --resume_from_checkpoint "$CHECKPOINT_PATH"