diff --git a/fastvideo/utils/checkpoint.py b/fastvideo/utils/checkpoint.py index ef661a190..b624ffee2 100644 --- a/fastvideo/utils/checkpoint.py +++ b/fastvideo/utils/checkpoint.py @@ -11,7 +11,6 @@ from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_ from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import StateDictType -import dataclasses from fastvideo.utils.logging_ import main_print @@ -45,50 +44,13 @@ def save_checkpoint_optimizer(model, optimizer, rank, output_dir, step, discrimi optimizer_path = os.path.join(save_dir, "optimizer.pt") torch.save(optim_state, optimizer_path) else: - weight_path = os.path.join(save_dstate_dictir, "discriminator_pytorch_model.safetensors") + weight_path = os.path.join(save_dir, "discriminator_pytorch_model.safetensors") save_file(cpu_state, weight_path) optimizer_path = os.path.join(save_dir, "discriminator_optimizer.pt") torch.save(optim_state, optimizer_path) main_print(f"--> checkpoint saved at step {step}") -def save_checkpoint_v1(transformer, rank, output_dir, step): - # from torch.distributed.fsdp import FullyShardedDataParallel as FSDP -# from torch.distributed.fsdp import StateDictType, FullStateDictConfig - - # 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 torch.distributed.get_rank() == 0: - # torch.save(state_dict, "model_checkpoint.pt") - if rank <= 0: - save_dir = os.path.join(output_dir, f"checkpoint-{step}") - os.makedirs(save_dir, exist_ok=True) - # save using safetensors - # weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors") - weight_path = os.path.join(save_dir, "diffusion_pytorch_model.pt") - print(weight_path) - # save_file(cpu_state, weight_path) - torch.save(cpu_state, weight_path) - config_dict = transformer.hf_config - if "dtype" in config_dict: - del config_dict["dtype"] # TODO - config_path = os.path.join(save_dir, "config.json") - # save dict as json - with open(config_path, "w") as f: - json.dump(config_dict, f, indent=4) - main_print(f"--> checkpoint saved at step {step}") - - - def save_checkpoint(transformer, rank, output_dir, step): main_print(f"--> saving checkpoint at step {step}") with FSDP.state_dict_type( diff --git a/fastvideo/v1/pipelines/training_pipeline.py b/fastvideo/v1/pipelines/training_pipeline.py index b303a8503..b3846ee0e 100644 --- a/fastvideo/v1/pipelines/training_pipeline.py +++ b/fastvideo/v1/pipelines/training_pipeline.py @@ -19,7 +19,6 @@ from tqdm.auto import tqdm # import torch.distributed as dist import wandb from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input -from fastvideo.utils.checkpoint import save_checkpoint_v1 from fastvideo.v1.configs.sample import SamplingParam from fastvideo.v1.dataset.parquet_datasets import ParquetVideoTextDataset from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group @@ -30,7 +29,7 @@ from fastvideo.v1.pipelines import ComposedPipelineBase from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.pipelines.training_utils import ( _clip_grad_norm_while_handling_failing_dtensor_cases, - compute_density_for_timestep_sampling, get_sigmas) + compute_density_for_timestep_sampling, get_sigmas, save_checkpoint) from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline logger = init_logger(__name__) @@ -558,12 +557,14 @@ class WanTrainingPipeline(TrainingPipeline): total_loss = 0.0 optimizer.zero_grad() for _ in range(gradient_accumulation_steps): + logger.info(f"Rank {self.rank}: Training step {_}") ( latents, encoder_hidden_states, encoder_attention_mask, infos, ) = next(loader_iter) + logger.info(f"Rank {self.rank}: Training step {_} loaded data") latents = latents.to(self.fastvideo_args.device, dtype=torch.bfloat16) encoder_hidden_states = encoder_hidden_states.to( @@ -772,26 +773,16 @@ class WanTrainingPipeline(TrainingPipeline): step=step, ) if step % args.checkpointing_steps == 0: - if args.use_lora: - raise NotImplementedError("LoRA is not supported now") - # Save LoRA weights - save_lora_checkpoint(transformer, optimizer, rank, - args.output_dir, step, pipe) - else: - # Your existing checkpoint saving code - save_checkpoint_v1(self.transformer, self.rank, - args.output_dir, step) - self.transformer.train() + # Your existing checkpoint saving code + save_checkpoint(self.transformer, self.rank, args.output_dir, + step) + self.transformer.train() self.sp_group.barrier() if args.log_validation and step % args.validation_steps == 0: self.log_validation(self.transformer, args, step) - if args.use_lora: - raise NotImplementedError("LoRA is not supported now") - # save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps, pipe) - else: - save_checkpoint_v1(self.transformer, self.rank, args.output_dir, - args.max_train_steps) + save_checkpoint(self.transformer, self.rank, args.output_dir, + args.max_train_steps) if get_sp_group(): cleanup_dist_env_and_memory() diff --git a/fastvideo/v1/pipelines/training_utils.py b/fastvideo/v1/pipelines/training_utils.py index 8fdd25923..4421d4fdc 100644 --- a/fastvideo/v1/pipelines/training_utils.py +++ b/fastvideo/v1/pipelines/training_utils.py @@ -1,16 +1,14 @@ import json -import os import math +import os from typing import List, Optional, Tuple, Union import torch import torch.distributed as dist - -import torch +import torch.distributed.tensor from torch.distributed.fsdp import FullStateDictConfig from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import StateDictType -import torch.distributed.tensor from fastvideo.v1.logger import init_logger @@ -94,9 +92,7 @@ def save_checkpoint(transformer, rank, output_dir, step): # save dict as json with open(config_path, "w") as f: json.dump(config_dict, f, indent=4) - logger.info("--> checkpoint saved at step {step} to {weight_path}", - step=step, - weight_path=weight_path) + logger.info("--> checkpoint saved at step %s to %s", step, weight_path) def _clip_grad_norm_while_handling_failing_dtensor_cases( diff --git a/train_v1.sh b/train_v1.sh index 74dafe025..45999c4be 100644 --- a/train_v1.sh +++ b/train_v1.sh @@ -25,13 +25,13 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\ --sp_size $NUM_GPUS \ --tp_size $NUM_GPUS \ --train_sp_batch_size 1\ - --dataloader_num_workers 1\ + --dataloader_num_workers 5\ --gradient_accumulation_steps=1\ --max_train_steps=120 \ --learning_rate=1e-6\ --mixed_precision="bf16"\ - --checkpointing_steps=60 \ - --validation_steps 20\ + --checkpointing_steps=5 \ + --validation_steps 10\ --validation_sampling_steps "2,4,8" \ --log_validation \ --checkpoints_total_limit 3\