import os import sys sys.path.append(os.path.abspath('.')) import argparse import datetime import numpy as np import time import torch import logging import json import math import random import diffusers import transformers from pathlib import Path from packaging import version from copy import deepcopy from dataset import ( ImageTextDataset, LengthGroupedVideoTextDataset, create_image_text_dataloaders, create_length_grouped_video_text_dataloader ) from pyramid_dit import ( PyramidDiTForVideoGeneration, JointTransformerBlock, FluxSingleTransformerBlock, FluxTransformerBlock, ) from trainer_misc import ( init_distributed_mode, setup_for_distributed, create_optimizer, train_one_epoch_with_fsdp, constant_scheduler, cosine_scheduler, ) from trainer_misc import ( is_sequence_parallel_initialized, init_sequence_parallel_group, get_sequence_parallel_proc_num, init_sync_input_group, get_sync_input_group, ) from collections import OrderedDict from PIL import Image from PIL import ImageFile ImageFile.LOAD_TRUNCATED_IMAGES = True from torch.distributed.fsdp.fully_sharded_data_parallel import ( FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig, ShardedStateDictConfig, ShardingStrategy, BackwardPrefetch, MixedPrecision, CPUOffload, StateDictType, ) from torch.distributed.fsdp.wrap import ModuleWrapPolicy, size_based_auto_wrap_policy from transformers.models.clip.modeling_clip import CLIPEncoderLayer from transformers.models.t5.modeling_t5 import T5Block import accelerate from accelerate import Accelerator from accelerate.utils import DistributedType, ProjectConfiguration, set_seed from accelerate import FullyShardedDataParallelPlugin from diffusers.utils import is_wandb_available from accelerate.logging import get_logger from accelerate.state import AcceleratorState from diffusers.optimization import get_scheduler logger = get_logger(__name__) def get_args(): parser = argparse.ArgumentParser('Pyramid-Flow Multi-process Training script', add_help=False) parser.add_argument('--task', default='t2v', type=str, choices=["t2v", "t2i"], help="Training image generation or video generation") parser.add_argument('--batch_size', default=4, type=int, help="The per device batch size") parser.add_argument('--epochs', default=100, type=int) parser.add_argument('--print_freq', default=20, type=int) parser.add_argument('--iters_per_epoch', default=2000, type=int) parser.add_argument('--save_ckpt_freq', default=20, type=int) # Model parameters parser.add_argument('--ema_update', action='store_true') parser.add_argument('--ema_decay', default=0.9999, type=float, metavar='MODEL', help='ema decay rate') parser.add_argument('--load_ema_model', default='', type=str, help='The ema model checkpoint loading') parser.add_argument('--model_name', default='pyramid_flux', type=str, help="The Model Architecture Name", choices=["pyramid_flux", "pyramid_mmdit"]) parser.add_argument('--model_path', default='', type=str, help='The pre-trained dit weight path') parser.add_argument('--model_variant', default='diffusion_transformer_384p', type=str, help='The dit model variant', choices=['diffusion_transformer_768p', 'diffusion_transformer_384p', 'diffusion_transformer_image']) parser.add_argument('--model_dtype', default='bf16', type=str, help="The Model Dtype: bf16 or fp16", choices=['bf16', 'fp16']) parser.add_argument('--load_model_ema_to_cpu', action='store_true') # FSDP condig parser.add_argument('--use_fsdp', action='store_true') parser.add_argument('--fsdp_shard_strategy', default='zero2', type=str, choices=['zero2', 'zero3']) # The training manner config parser.add_argument('--use_flash_attn', action='store_true') parser.add_argument('--use_temporal_causal', action='store_true', default=True) parser.add_argument('--interp_condition_pos', action='store_true', default=True) parser.add_argument('--sync_video_input', action='store_true', help="whether to sync the video input") parser.add_argument('--load_text_encoder', action='store_true', help="whether to load the text encoder during training") parser.add_argument('--load_vae', action='store_true', help="whether to load the video vae during training") # Sequence Parallel config parser.add_argument('--use_sequence_parallel', action='store_true') parser.add_argument('--sp_group_size', default=1, type=int, help="The group size of sequence parallel") parser.add_argument('--sp_proc_num', default=-1, type=int, help="The number of process used for video training, default=-1 means using all process. This args indicated using how many processes for video training") # Model input config parser.add_argument('--max_frames', default=16, type=int, help='number of max video frames') parser.add_argument('--frame_per_unit', default=1, type=int, help="The number of frames per training unit") parser.add_argument('--schedule_shift', default=1.0, type=float, help="The flow matching schedule shift") parser.add_argument('--corrupt_ratio', default=1/3, type=float, help="The corruption ratio for the clean history in AR training") # Dataset Cconfig parser.add_argument('--anno_file', default='', type=str, help="The annotation jsonl file") parser.add_argument('--resolution', default='384p', type=str, help="The input resolution", choices=['384p', '768p']) # Training set config parser.add_argument('--dit_pretrained_weight', default='', type=str, help='The pretrained dit checkpoint') parser.add_argument('--vae_pretrained_weight', default='', type=str,) parser.add_argument('--not_add_normalize', action='store_true') parser.add_argument('--use_temporal_pyramid', action='store_true', help="Whether to use the AR temporal pyramid training for video generation") parser.add_argument('--gradient_checkpointing', action='store_true') parser.add_argument('--gradient_checkpointing_ratio', type=float, default=0.75, help="The ratio of transformer blocks used for gradient_checkpointing") parser.add_argument('--gradient_accumulation_steps', default=1, type=int, help="Number of updates steps to accumulate before performing a backward/update pass.") parser.add_argument('--video_sync_group', default=8, type=int, help="The number of process that accepts the same input video, used for temporal pyramid AR training. \ This contributes to stable AR training. We recommend to set this value to 4, 8 or 16. If you have enough GPUs, set it equals to max_frames (16 for 5s, 32 for 10s), \ make sure to satisfy `max_frames % video_sync_group == 0`") # Optimizer parameters parser.add_argument('--opt', default='adamw', type=str, metavar='OPTIMIZER', help='Optimizer (default: "adamw"') parser.add_argument('--opt_eps', default=1e-8, type=float, metavar='EPSILON', help='Optimizer Epsilon (default: 1e-8)') parser.add_argument('--opt_beta1', default=0.9, type=float, metavar='BETA1', help='Optimizer Betas (default: None, use opt default)') parser.add_argument('--opt_beta2', default=0.999, type=float, metavar='BETA2', help='Optimizer Betas (default: None, use opt default)') parser.add_argument('--clip_grad', type=float, default=None, metavar='NORM', help='Clip gradient norm (default: None, no clipping)') parser.add_argument('--weight_decay', type=float, default=1e-4, help='weight decay (default: 1e-4)') parser.add_argument('--lr', type=float, default=5e-5, metavar='LR', help='learning rate (default: 5e-5)') parser.add_argument('--warmup_lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)') parser.add_argument('--min_lr', type=float, default=1e-5, metavar='LR', help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') parser.add_argument( "--lr_scheduler", type=str, default="constant_with_warmup", help=( 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' ' "constant", "constant_with_warmup"]' ), ) parser.add_argument('--warmup_epochs', type=int, default=1, metavar='N', help='epochs to warmup LR, if scheduler supports') parser.add_argument('--warmup_steps', type=int, default=-1, metavar='N', help='epochs to warmup LR, if scheduler supports') # Dataset parameters parser.add_argument('--output_dir', type=str, default='', help='path where to save, empty for no saving') parser.add_argument('--logging_dir', type=str, default='log', help='path where to tensorboard log') parser.add_argument( "--report_to", type=str, default="tensorboard", help=( 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' ), ) # Distributed Training parameters parser.add_argument('--device', default='cuda', type=str, help='device to use for training / testing') parser.add_argument('--seed', default=0, type=int) parser.add_argument('--resume', default='', help='resume from checkpoint') parser.add_argument('--auto_resume', action='store_true') parser.set_defaults(auto_resume=True) parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch') parser.add_argument('--global_step', default=0, type=int, metavar='N', help='The global optimization step') parser.add_argument('--num_workers', default=10, type=int) parser.add_argument('--pin_mem', action='store_true', help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.') parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem', help='') parser.set_defaults(pin_mem=True) # distributed training parameters parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes') parser.add_argument('--local_rank', default=-1, type=int) parser.add_argument('--dist_on_itp', action='store_true') parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training', type=str) return parser.parse_args() def build_model_runner(args): model_dtype = args.model_dtype model_path = args.model_path model_name = args.model_name model_variant = args.model_variant print(f"Load the {model_name} model checkpoint from path: {model_path}, using dtype {model_dtype}") sample_ratios = [1, 2, 1] # The sample_ratios of each stage assert args.batch_size % int(sum(sample_ratios)) == 0, "The batchsize should be diivided by sum(sample_ratios)" runner = PyramidDiTForVideoGeneration( model_path, model_dtype, model_name=model_name, use_gradient_checkpointing=args.gradient_checkpointing, gradient_checkpointing_ratio=args.gradient_checkpointing_ratio, return_log=True, model_variant=model_variant, timestep_shift=args.schedule_shift, stages=[1, 2, 4], # using 3 stages stage_range=[0, 1/3, 2/3, 1], sample_ratios=sample_ratios, # The sample proportion in a training batch use_mixed_training=True, use_flash_attn=args.use_flash_attn, load_text_encoder=args.load_text_encoder, load_vae=args.load_vae, max_temporal_length=args.max_frames, frame_per_unit=args.frame_per_unit, use_temporal_causal=args.use_temporal_causal, corrupt_ratio=args.corrupt_ratio, interp_condition_pos=args.interp_condition_pos, video_sync_group=args.video_sync_group, ) if args.dit_pretrained_weight: dit_pretrained_weight = args.dit_pretrained_weight print(f"Loading the pre-trained DiT checkpoint from {dit_pretrained_weight}") runner.load_checkpoint(dit_pretrained_weight) if args.vae_pretrained_weight: vae_pretrained_weight = args.vae_pretrained_weight print(f"Loading the pre-trained VAE checkpoint from {vae_pretrained_weight}") runner.load_vae_checkpoint(vae_pretrained_weight) return runner def auto_resume(args, accelerator): if len(args.resume) > 0: path = args.resume else: # Get the most recent checkpoint dirs = os.listdir(args.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: accelerator.print( f"Checkpoint does not exist. Starting a new training run." ) initial_global_step = 0 else: accelerator.print(f"Resuming from checkpoint {path}") accelerator.load_state(os.path.join(args.output_dir, path)) global_step = int(path.split("-")[1]) initial_global_step = global_step return initial_global_step def build_fsdp_plugin(args): fsdp_plugin = FullyShardedDataParallelPlugin( sharding_strategy=ShardingStrategy.SHARD_GRAD_OP if args.fsdp_shard_strategy == 'zero2' else ShardingStrategy.FULL_SHARD, backward_prefetch=BackwardPrefetch.BACKWARD_PRE, auto_wrap_policy=ModuleWrapPolicy([FluxSingleTransformerBlock, FluxTransformerBlock, JointTransformerBlock, T5Block, CLIPEncoderLayer]), cpu_offload=CPUOffload(offload_params=False), state_dict_type=StateDictType.FULL_STATE_DICT, state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=True), optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True), ) return fsdp_plugin def main(args): logging_dir = Path(args.output_dir, args.logging_dir) accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) # Initialize the Environment variables throught MPI run init_distributed_mode(args, init_pytorch_ddp=False) # set `init_pytorch_ddp` to False, since the accelerate will do later if args.use_fsdp: fsdp_plugin = build_fsdp_plugin(args) else: fsdp_plugin = None accelerator = Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.model_dtype, log_with=args.report_to, project_config=accelerator_project_config, fsdp_plugin=fsdp_plugin, ) # To block the print on non main process setup_for_distributed(accelerator.is_main_process) # If uses the sequence parallel if args.use_sequence_parallel: assert args.sp_group_size > 1, "Sequence Parallel needs group size > 1" init_sequence_parallel_group(args) print(f"Using sequence parallel, the parallel size is {args.sp_group_size}") if args.sp_proc_num == -1: args.sp_proc_num = accelerator.num_processes # if not specified, all processes are used for video training if args.report_to == "wandb": if not is_wandb_available(): raise ImportError("Make sure to install wandb if you want to use it for logging during training.") import wandb logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO, ) logger.info(accelerator.state, main_process_only=False) if accelerator.is_local_main_process: transformers.utils.logging.set_verbosity_warning() diffusers.utils.logging.set_verbosity_info() else: transformers.utils.logging.set_verbosity_error() diffusers.utils.logging.set_verbosity_error() if args.seed is not None: set_seed(args.seed, device_specific=True) device = accelerator.device # building model runner = build_model_runner(args) # For mixed precision training we cast all non-trainable weights to half-precision # as these weights are only used for inference, keeping weights in full precision is not required. weight_dtype = torch.float32 if accelerator.mixed_precision == "fp16": weight_dtype = torch.float16 elif accelerator.mixed_precision == "bf16": weight_dtype = torch.bfloat16 # Move unet, vae and text_encoder to device and cast to weight_dtype # The VAE is in float32 to avoid NaN losses. if runner.vae: logger.info(f"Rank {args.rank}: Casting VAE to {weight_dtype}", main_process_only=False) runner.vae.to(dtype=weight_dtype) if runner.text_encoder: logger.info(f"Rank {args.rank}: Casting TextEncoder to {weight_dtype}", main_process_only=False) runner.text_encoder.to(dtype=weight_dtype) # building dataloader global_rank = accelerator.process_index anno_file = args.anno_file if args.task == 't2i': # For image generation training if args.resolution == '384p': image_ratios = [1/1, 3/5, 5/3] image_sizes = [(512, 512), (384, 640), (640, 384)] else: assert args.resolution == '768p' image_ratios = [1/1, 3/5, 5/3] image_sizes = [(1024, 1024), (768, 1280), (1280, 768)] image_text_dataset = ImageTextDataset( anno_file, add_normalize=not args.not_add_normalize, ratios=image_ratios, sizes=image_sizes, ) train_dataloader = create_image_text_dataloaders( image_text_dataset, batch_size=args.batch_size, num_workers=args.num_workers, multi_aspect_ratio=True, epoch=args.seed, sizes=image_sizes, use_distributed=True, world_size=accelerator.num_processes, rank=global_rank, ) else: assert args.task == 't2v' # For video generation training video_text_dataset = LengthGroupedVideoTextDataset( anno_file, max_frames=args.max_frames, resolution=args.resolution, load_vae_latent=not args.load_vae, load_text_fea=not args.load_text_encoder, ) if args.sync_video_input: assert args.sp_proc_num % args.video_sync_group == 0, "The video_sync_group should be divided by world size" assert args.max_frames % args.video_sync_group == 0, "The video_sync_group should be divided by num_frames" train_dataloader = create_length_grouped_video_text_dataloader( video_text_dataset, batch_size=args.batch_size, num_workers=args.num_workers, max_frames=args.max_frames, world_size=args.sp_proc_num // args.video_sync_group, rank=global_rank // args.video_sync_group, epoch=args.seed, use_distributed=True, ) else: train_dataloader = create_length_grouped_video_text_dataloader( video_text_dataset, batch_size=args.batch_size, num_workers=args.num_workers, max_frames=args.max_frames, world_size=args.sp_proc_num, rank=global_rank, epoch=args.seed, use_distributed=True, ) accelerator.wait_for_everyone() logger.info("Building dataset finished") # building ema model model_ema = deepcopy(runner.dit) if args.ema_update else None if model_ema: model_ema.eval() # set the ema model not update by gradient if model_ema: model_ema.to(dtype=weight_dtype) for param in model_ema.parameters(): param.requires_grad = False # report model details n_learnable_parameters = sum(p.numel() for p in runner.dit.parameters() if p.requires_grad) n_fix_parameters = sum(p.numel() for p in runner.dit.parameters() if not p.requires_grad) logger.info(f'total number of learnable params: {n_learnable_parameters / 1e6} M') logger.info(f'total number of fixed params in : {n_fix_parameters / 1e6} M') # `accelerate` 0.16.0 will have better support for customized saving # Register Hook to load and save model_ema if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): if accelerator.is_main_process: if model_ema: model_ema_state = model_ema.state_dict() torch.save(model_ema_state, os.path.join(output_dir, 'pytorch_model_ema.bin')) def load_model_hook(models, input_dir): if model_ema: model_ema_path = os.path.join(input_dir, 'pytorch_model_ema.bin') if os.path.exists(model_ema_path): model_ema_state = torch.load(model_ema_path, map_location='cpu') load_res = model_ema.load_state_dict(model_ema_state) print(f"Loading ema weights {load_res}") accelerator.register_save_state_pre_hook(save_model_hook) accelerator.register_load_state_pre_hook(load_model_hook) # Create the Optimizer optimizer = create_optimizer(args, runner.dit) logger.info(f"optimizer: {optimizer}") # Create the LR scheduler num_training_steps_per_epoch = args.iters_per_epoch args.max_train_steps = args.epochs * num_training_steps_per_epoch warmup_iters = args.warmup_epochs * num_training_steps_per_epoch if args.warmup_steps > 0: warmup_iters = args.warmup_steps logger.info(f"LRScheduler: {args.lr_scheduler}, Warmup steps: {warmup_iters * args.gradient_accumulation_steps}") if args.lr_scheduler == 'cosine': lr_schedule_values = cosine_scheduler( args.lr, args.min_lr, args.epochs, num_training_steps_per_epoch, warmup_epochs=args.warmup_epochs, warmup_steps=args.warmup_steps, ) elif args.lr_scheduler == 'constant_with_warmup': lr_schedule_values = constant_scheduler( args.lr, args.epochs, num_training_steps_per_epoch, warmup_epochs=args.warmup_epochs, warmup_steps=args.warmup_steps, ) else: raise NotImplementedError(f"Not Implemented for scheduler {args.lr_scheduler}") # Wrap the model, optmizer, and scheduler with accelerate logger.info(f'before accelerator.prepare') if fsdp_plugin is not None: logger.info(f'show fsdp configs:') print('accelerator.state.fsdp_plugin.use_orig_params', accelerator.state.fsdp_plugin.use_orig_params) print('accelerator.state.fsdp_plugin.sync_module_states', accelerator.state.fsdp_plugin.sync_module_states) print('accelerator.state.fsdp_plugin.forward_prefetch', accelerator.state.fsdp_plugin.forward_prefetch) print('accelerator.state.fsdp_plugin.mixed_precision_policy', accelerator.state.fsdp_plugin.mixed_precision_policy) print('accelerator.state.fsdp_plugin.backward_prefetch', accelerator.state.fsdp_plugin.backward_prefetch) # Only wrapping the trained dit and huge text encoder runner.dit, optimizer = accelerator.prepare(runner.dit, optimizer) # Load the VAE and EMAmodel to GPU if runner.vae: runner.vae.to(device) if runner.text_encoder: runner.text_encoder.to(device) logger.info(f'after accelerator.prepare') logger.info(f'{runner.dit}') if model_ema and (not args.load_model_ema_to_cpu): model_ema.to(device) if accelerator.is_main_process: accelerator.init_trackers(os.path.basename(args.output_dir), config=vars(args)) # Report the training info total_batch_size = args.batch_size * accelerator.num_processes * args.gradient_accumulation_steps logger.info("***** Running training *****") logger.info("LR = %.8f" % args.lr) logger.info("Min LR = %.8f" % args.min_lr) logger.info("Weigth Decay = %.8f" % args.weight_decay) logger.info("Batch size = %d" % total_batch_size) logger.info("Number of training steps = %d" % (num_training_steps_per_epoch * args.epochs)) logger.info("Number of training examples per epoch = %d" % (total_batch_size * num_training_steps_per_epoch)) # Auto resume the checkpoint initial_global_step = auto_resume(args, accelerator) first_epoch = initial_global_step // num_training_steps_per_epoch # Start Train! start_time = time.time() accelerator.wait_for_everyone() for epoch in range(first_epoch, args.epochs): train_stats = train_one_epoch_with_fsdp( runner, model_ema, accelerator, args.model_dtype, train_dataloader, optimizer, lr_schedule_values, device, epoch, args.clip_grad, start_steps=epoch * num_training_steps_per_epoch, args=args, print_freq=args.print_freq, iters_per_epoch=num_training_steps_per_epoch, ema_decay=args.ema_decay, use_temporal_pyramid=args.use_temporal_pyramid, ) if args.output_dir: if (epoch + 1) % args.save_ckpt_freq == 0 or epoch + 1 == args.epochs: if accelerator.sync_gradients: global_step = num_training_steps_per_epoch * (epoch + 1) save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") accelerator.save_state(save_path, safe_serialization=False) logger.info(f"Saved state to {save_path}") accelerator.wait_for_everyone() log_stats = {**{f'train_{k}': v for k, v in train_stats.items()}, 'epoch': epoch, 'n_parameters': n_learnable_parameters} if args.output_dir and accelerator.is_main_process: with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f: f.write(json.dumps(log_stats) + "\n") total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str)) accelerator.wait_for_everyone() accelerator.end_training() if __name__ == '__main__': os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["FSDP_USE_ORIG_PARAMS"] = "true" opts = get_args() if opts.output_dir: Path(opts.output_dir).mkdir(parents=True, exist_ok=True) main(opts)