nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[Training] Add distributed checkpointing (#458)

Browse Source
This commit is contained in:
Kevin Lin
2025-06-02 15:15:45 -07:00
committed by GitHub
parent d156461785
commit bdfdf1dfee
6 changed files with 320 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/v1/fastvideo_args.py
+4 -1
View File
@@ -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",
fastvideo/v1/training/checkpointing_utils.py
+107
View File
@@ -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
+8 -2
View File
@@ -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
fastvideo/v1/training/training_utils.py
+159 -20
View File
@@ -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
fastvideo/v1/training/wan_training_pipeline.py
+37 -11
View File
@@ -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()
scripts/finetune/finetune_v1.sh
+5 -1
View File
@@ -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"