From 897afb8b947c5bbd94a59e71d5202a7e262c4e86 Mon Sep 17 00:00:00 2001 From: Zhang Peiyuan Date: Fri, 15 Nov 2024 19:28:41 -0800 Subject: [PATCH] Add LADD (#45) Co-authored-by: Peiyuan Zhang --- README.md | 19 +- fastvideo/distill.py | 509 +++++++++++ fastvideo/fsdp_util.py | 51 +- fastvideo/gan/discriminator.py | 123 +++ fastvideo/model/mochi_latents_stat.py | 33 - fastvideo/model/mochi_latents_utils.py | 39 + fastvideo/model/modeling_mochi.py | 132 ++- fastvideo/model/norm.py | 168 ++++ fastvideo/sample/sample_t2v_mochi.py | 16 +- fastvideo/train.py | 92 +- fastvideo/utils/encode_vae_t5.py | 867 ------------------- inference_mochi.sh | 79 +- pyproject.toml | 6 +- scripts/data_preprocess/finetune_data_gen.sh | 6 +- scripts/finetune_T_J.sh | 117 +++ t2v_debug_multi.sh | 71 +- 16 files changed, 1268 insertions(+), 1060 deletions(-) create mode 100644 fastvideo/distill.py create mode 100644 fastvideo/gan/discriminator.py delete mode 100644 fastvideo/model/mochi_latents_stat.py create mode 100644 fastvideo/model/mochi_latents_utils.py create mode 100644 fastvideo/model/norm.py delete mode 100644 fastvideo/utils/encode_vae_t5.py create mode 100644 scripts/finetune_T_J.sh diff --git a/README.md b/README.md index c2d834900..fff6289fa 100644 --- a/README.md +++ b/README.md @@ -8,9 +8,7 @@ conda create -n fastvideo python=3.10.12 conda activate fastvideo pip3 install torch==2.5.0 torchvision --index-url https://download.pytorch.org/whl/cu121 pip3 install -U xformers==0.0.28.post2 --index-url https://download.pytorch.org/whl/cu121 -pip3 install ray -cd .. && git clone https://github.com/huggingface/diffusers -cd diffusers && pip install -e . && cd ../FastVideo-OSP +pip install git+https://github.com/huggingface/diffusers.git@76b7d86a9a5c0c2186efa09c4a67b5f5666ac9e3 ``` ``` @@ -26,7 +24,7 @@ python scripts/download_hf.py --repo_id=Stealths-Video/dummyVid --local_dir=data python scripts/download_hf.py --repo_id=Stealths-Video/mochi_diffuser --local_dir=data/mochi --repo_type=model python scripts/download_hf.py --repo_id=Stealths-Video/Mochi-Synthetic-Data --local_dir=data/Mochi-Synthetic-Data --repo_type=dataset python scripts/download_hf.py --repo_id=Stealths-Video/Encoder_Overfit_Data --local_dir=data/Encoder_Overfit_Data --repo_type=dataset -python scripts/download_hf.py --repo_id=Stealths-Video/BW-Finetune-Dataset --local_dir=data/BW-Finetune-Dataset --repo_type=dataset +python scripts/download_hf.py --repo_id=Stealths-Video/Black-Myth-Recaption --local_dir=Stealths-Video/BLACK-MYTH-Finetune-Dataset --repo_type=dataset ``` ## How to overfit @@ -35,10 +33,11 @@ bash scripts/overfit.sh ``` Make sure to edit data/Mochi-Synthetic-Data/videos2caption.json such that this is only one video in the dataset (you can copy multiple annotations of the same video). Also make sure to edit the prompt in scripts/overfit.shto match the prompt in the training data. I observe the overfitting after 50 steps. -## TODO -- [X] Delete all npu related stuff. -- [ ] Remove inpaint. -- [ ] Create dummy debug data. -- [ ] Add Mochi -- [ ] Add Mochi VAE +## Feature for Branding: +- [] LoRA +- [] Low mem. Only load DiT during training. +- [] selective gradient checkpointing +- [] Fused kernel +- [] correct attention mask issue in HF's implementation +- [] \ No newline at end of file diff --git a/fastvideo/distill.py b/fastvideo/distill.py new file mode 100644 index 000000000..ac5f7a9d2 --- /dev/null +++ b/fastvideo/distill.py @@ -0,0 +1,509 @@ +import argparse +import math +import os +from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \ + destroy_sequence_parallel_group, get_sequence_parallel_state, nccl_info +from fastvideo.utils.communications import sp_parallel_dataloader_wrapper, broadcast +from diffusers.training_utils import cast_training_params, compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3 + +from fastvideo.model.mochi_latents_utils import normalize_mochi_dit_input +from fastvideo.utils.validation import log_validation +import time +from torch.utils.data import DataLoader +import torch +from torch.distributed.fsdp import ( + FullyShardedDataParallel as FSDP, + StateDictType, + FullStateDictConfig, +) +import json +from torch.utils.data.distributed import DistributedSampler +import wandb +from accelerate.utils import set_seed +from tqdm.auto import tqdm +from fastvideo.fsdp_util import get_dit_fsdp_kwargs, get_discriminator_fsdp_kwargs, apply_fsdp_checkpointing +from diffusers import ( + FlowMatchEulerDiscreteScheduler, +) +import torch.nn.functional as F + +import numpy as np +from fastvideo.model.modeling_mochi import MochiTransformer3DModel +from diffusers.utils import check_min_version +from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_function +from fastvideo.gan.discriminator import LADDDiscriminator +import torch.distributed as dist +from safetensors.torch import save_file +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.31.0") + + +import sys +import pdb + +class ForkedPdb(pdb.Pdb): + """A Pdb subclass that may be used + from a forked multiprocessing child + + """ + def interaction(self, *args, **kwargs): + _stdin = sys.stdin + try: + sys.stdin = open('/dev/stdin') + pdb.Pdb.interaction(self, *args, **kwargs) + finally: + sys.stdin = _stdin + + + +def main_print(content): + if int(os.environ['LOCAL_RANK']) <= 0: + print(content) + +def save_checkpoint(transformer: MochiTransformer3DModel, rank, output_dir, step): + main_print(f"--> saving checkpoint at step {step}") + with FSDP.state_dict_type( + transformer, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + ): + cpu_state = transformer.state_dict() + + 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") + save_file(cpu_state, weight_path) + config_dict = dict(transformer.config) + 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 get_sigmas(noise_scheduler, device, timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype) + schedule_timesteps = noise_scheduler.timesteps.to(device) + timesteps = timesteps.to(device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + +def student_forward(transformer, latents, encoder_hidden_states, encoder_attention_mask, sp_size): + batch_size = latents.shape[0] + noise = torch.randn_like(latents) + # uniformly sample from [1.0, 0.75, 0.5, 0.25] + p = [0.25, 0.25, 0.25, 0.25] + sigmas = np.random.choice([1.0, 0.75, 0.5, 0.25], batch_size, p=p) + sigmas = torch.tensor(sigmas, device=latents.device, dtype=latents.dtype) + student_timesteps = (sigmas * 1000).long() + if sp_size > 1: + # Make sure that the timesteps are the same across all sp processes. + broadcast(student_timesteps) + + while len(sigmas.shape) < latents.ndim: + sigmas = sigmas.unsqueeze(-1) + noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise + + model_pred = transformer( + noisy_model_input, + encoder_hidden_states, + 1000 - student_timesteps, + encoder_attention_mask, # B, L + return_dict= False, + )[0] + # calculate X1 + model_pred = model_pred * sigmas + noisy_model_input + return model_pred + + +def teacher_forward(transformer, latents, encoder_hidden_states, encoder_attention_mask, sp_size, noise_scheduler, weighting_scheme, logit_mean, logit_std, mode_scale): + batch_size = latents.shape[0] + noise = torch.randn_like(latents) + u = compute_density_for_timestep_sampling( + weighting_scheme=weighting_scheme, + batch_size=batch_size, + logit_mean=logit_mean, + logit_std=logit_std, + mode_scale=mode_scale, + ) + indices = (u * noise_scheduler.config.num_train_timesteps).long() + timesteps = noise_scheduler.timesteps[indices].to(device=latents.device) + if sp_size > 1: + # Make sure that the timesteps are the same across all sp processes. + broadcast(timesteps) + sigmas = get_sigmas(noise_scheduler, latents.device, timesteps, n_dim=latents.ndim, dtype=latents.dtype) + noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise + teacher_features = transformer( + noisy_model_input, + encoder_hidden_states, + noise_scheduler.config.num_train_timesteps - timesteps, + encoder_attention_mask, # B, L + return_dict= False, + output_attn=True + )[1] + return teacher_features + +def train_one_step_generator(student_transformer, teacher_transformer, discriminator, student_optimizer, loader,noise_scheduler, gradient_accumulation_steps, sp_size, max_grad_norm, weighting_scheme, logit_mean, logit_std, mode_scale): + total_loss = 0.0 + teacher_transformer.requires_grad_(False) + discriminator.requires_grad_(False) + student_transformer.requires_grad_(True) + student_optimizer.zero_grad() + for _ in range(gradient_accumulation_steps): + latents, encoder_hidden_states, latents_attention_mask, encoder_attention_mask = next(loader) + latents = normalize_mochi_dit_input(latents) + + student_pred = student_forward(student_transformer, latents, encoder_hidden_states, encoder_attention_mask, sp_size) + + teacher_features = teacher_forward(teacher_transformer, student_pred, encoder_hidden_states, encoder_attention_mask, sp_size, noise_scheduler, weighting_scheme, logit_mean, logit_std, mode_scale) + + logits = discriminator(teacher_features, encoder_hidden_states) + + loss = (F.relu(torch.ones_like(logits) - logits)).mean() / gradient_accumulation_steps + loss.backward() + + avg_loss = loss.detach().clone() + dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG) + total_loss += avg_loss.item() + + + grad_norm = student_transformer.clip_grad_norm_(max_grad_norm) + student_optimizer.step() + return total_loss, grad_norm.item() + +def train_one_step_discriminator(student_transformer, teacher_transformer, discriminator, discriminator_optimizer, loader,noise_scheduler, gradient_accumulation_steps, sp_size, max_grad_norm, weighting_scheme, logit_mean, logit_std, mode_scale): + total_loss = 0.0 + teacher_transformer.requires_grad_(False) + student_transformer.requires_grad_(False) + discriminator.requires_grad_(True) + discriminator_optimizer.zero_grad() + for _ in range(gradient_accumulation_steps): + latents, encoder_hidden_states, latents_attention_mask, encoder_attention_mask = next(loader) + latents = normalize_mochi_dit_input(latents) + with torch.no_grad(): + student_pred = student_forward(student_transformer, latents, encoder_hidden_states, encoder_attention_mask, sp_size) + teacher_features_fake = teacher_forward(teacher_transformer, student_pred, encoder_hidden_states, encoder_attention_mask, sp_size, noise_scheduler, weighting_scheme, logit_mean, logit_std, mode_scale) + teacher_features_real = teacher_forward(teacher_transformer, latents, encoder_hidden_states, encoder_attention_mask, sp_size, noise_scheduler, weighting_scheme, logit_mean, logit_std, mode_scale) + fake_logits = discriminator(teacher_features_fake, encoder_hidden_states) + real_logits = discriminator(teacher_features_real, encoder_hidden_states) + fake_loss = (F.relu(torch.ones_like(fake_logits) + fake_logits)).mean() + real_loss = (F.relu(torch.ones_like(real_logits) - real_logits)).mean() + + loss = (fake_loss + real_loss)/ 2 / gradient_accumulation_steps + + loss.backward() + + avg_loss = loss.detach().clone() + dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG) + total_loss += avg_loss.item() + + + grad_norm = discriminator.clip_grad_norm_(max_grad_norm) + discriminator_optimizer.step() + return total_loss, grad_norm.item() + + +def main(args): + # use LayerNorm, GeLu, SiLu always as fp32 mode + # TODO: + if args.enable_stable_fp32: + raise NotImplementedError("enable_stable_fp32 is not supported now.") + torch.backends.cuda.matmul.allow_tf32 = True + + local_rank = int(os.environ['LOCAL_RANK']) + rank = int(os.environ['RANK']) + world_size = int(os.environ['WORLD_SIZE']) + dist.init_process_group("nccl") + torch.cuda.set_device(local_rank) + device = torch.cuda.current_device() + initialize_sequence_parallel_state(args.sp_size) + + + + + # If passed along, set the training seed now. On GPU... + if args.seed is not None: + # TODO: t within the same seq parallel group should be the same. Noise should be different. + set_seed(args.seed + rank) + + # Handle the repository creation + if rank <=0 and args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weigths to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + + + # Create model: + + main_print(f"--> loading model...") + student_transformer = MochiTransformer3DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + ) + teacher_transformer = MochiTransformer3DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + load_dtype=torch.bfloat16, # TODO: FP8 or INT8 weight type + ) + # feature_dim, cnn_channel_size, text_cond_dim, num_layers, layers_stride): + discriminator = LADDDiscriminator(feature_dim=3072, text_cond_dim=4096, num_layers=48, layers_stride=4) + # get num of discriminator parameters in billion + main_print(f"--> discriminator parameters: {sum(p.numel() for p in discriminator.parameters()) / 1e9} B") + + teacher_transformer.requires_grad_(False) + + + + main_print(f"--> Initializing FSDP...") + dit_fsdp_kwargs = get_dit_fsdp_kwargs("full") + student_transformer = FSDP( + student_transformer, + **dit_fsdp_kwargs, + ) + teacher_transformer = FSDP( + teacher_transformer, + **dit_fsdp_kwargs, + ) + discriminator_fsdp_kwargs = get_discriminator_fsdp_kwargs("full") + discriminator = FSDP( + discriminator, + **discriminator_fsdp_kwargs, + ) + + + main_print(f"--> model loaded") + if args.gradient_checkpointing: + apply_fsdp_checkpointing(teacher_transformer) + apply_fsdp_checkpointing(student_transformer) + + + # Set model as trainable. + student_transformer.train() + + noise_scheduler = FlowMatchEulerDiscreteScheduler() + + + student_optimizer = torch.optim.AdamW( + student_transformer.parameters(), + lr=args.student_learning_rate, + betas=(0.9,0.999), + weight_decay=0.01, + eps=1e-8, + ) + discriminator_optimizer = torch.optim.AdamW( + discriminator.parameters(), + lr=args.discriminator_learning_rate, + betas=(0.9,0.999), + weight_decay=0.01, + eps=1e-8, + ) + + train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg) + sampler = DistributedSampler(train_dataset, rank=rank, num_replicas=world_size, shuffle=True) + train_dataloader = DataLoader( + train_dataset, + sampler=sampler, + collate_fn=latent_collate_function, + pin_memory=True, + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + drop_last=True, + ) + + + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size) + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if rank <= 0: + project = args.tracker_project_name or "fastvideo_distill" + wandb.init(project=project, config=args) + + # Train! + total_batch_size = args.train_batch_size * world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size + main_print("***** Running training *****") + main_print(f" Num examples = {len(train_dataset)}") + main_print(f" Dataloader size = {len(train_dataloader)}") + main_print(f" Num Epochs = {args.num_train_epochs}") + main_print(f" Instantaneous batch size per device = {args.train_batch_size}") + main_print(f" Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}") + main_print(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + main_print(f" Total optimization steps = {args.max_train_steps}") + main_print(f" Total student training parameters per FSDP shard = {sum(p.numel() for p in student_transformer.parameters() if p.requires_grad) / 1e9} B") + main_print(f" Total discriminator training parameters per FSDP shard = {sum(p.numel() for p in discriminator.parameters() if p.requires_grad) / 1e9} B") + + main_print(f" Master weight dtype: {student_transformer.parameters().__next__().dtype}") + + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + assert NotImplementedError("resume_from_checkpoint is not supported now.") + # TODO + + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=0, + desc="Steps", + # Only show the progress bar once on each machine. + disable= local_rank > 0, + ) + + + + loader = sp_parallel_dataloader_wrapper(train_dataloader, device, args.train_batch_size, args.sp_size, args.train_sp_batch_size) + + for step in range(1, args.discriminator_warmup_steps+1): + discriminator_loss, discriminator_grad_norm = train_one_step_discriminator(student_transformer, teacher_transformer, discriminator, \ + discriminator_optimizer, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size, \ + args.max_grad_norm, args.weighting_scheme, args.logit_mean, args.logit_std, args.mode_scale) + progress_bar.set_postfix({"d_loss": discriminator_loss, "d_grad_n": discriminator_grad_norm}) + progress_bar.update(1) + if rank <= 0: + wandb.log({"discriminator_loss": discriminator_loss, "discriminator_grad_norm": discriminator_grad_norm}, step=step) + + for step in range(args.discriminator_warmup_steps+1, args.max_train_steps+1): + stage = (step - args.discriminator_warmup_steps) // args.gan_switch_n_steps + if stage % 2 == 1: + # Train generator + generator_loss, student_grad_norm= train_one_step_generator(student_transformer, teacher_transformer, discriminator, \ + student_optimizer, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size, \ + args.max_grad_norm, args.weighting_scheme, args.logit_mean, args.logit_std, args.mode_scale) + progress_bar.set_postfix({"g_loss": generator_loss, "g_grad_n": student_grad_norm}) + if rank <= 0: + wandb.log({"generator_loss": generator_loss, "student_grad_norm": student_grad_norm}, step=step) + + else: + # Train discriminator + discriminator_loss, discriminator_grad_norm = train_one_step_discriminator(student_transformer, teacher_transformer, discriminator, \ + discriminator_optimizer, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size, \ + args.max_grad_norm, args.weighting_scheme, args.logit_mean, args.logit_std, args.mode_scale) + progress_bar.set_postfix({"d_loss": discriminator_loss, "d_grad_n": discriminator_grad_norm}) + if rank <= 0: + wandb.log({"discriminator_loss": discriminator_loss, "discriminator_grad_norm": discriminator_grad_norm}, step=step) + + progress_bar.update(1) + + if args.log_validation and step % args.validation_steps == 0: + log_validation(args, student_transformer, device, + torch.bfloat16, step) + + if step % args.checkpointing_steps == 0: + save_checkpoint(student_transformer, rank, args.output_dir, step) + dist.barrier() + + if get_sequence_parallel_state(): + destroy_sequence_parallel_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + # dataset & dataloader + parser.add_argument("--data_json_path", type=str, required=True) + parser.add_argument("--num_frames", type=int, default=163) + parser.add_argument("--dataloader_num_workers", type=int, default=10, help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.") + parser.add_argument("--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader.") + parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.") + parser.add_argument("--group_frame", action="store_true") # TODO + parser.add_argument("--group_resolution", action="store_true") # TODO + + # text encoder & vae & diffusion model + parser.add_argument("--pretrained_model_name_or_path", type=str) + parser.add_argument("--cache_dir", type=str, default='./cache_dir') + parser.add_argument('--enable_stable_fp32', action='store_true') # TODO + + # diffusion setting + parser.add_argument("--ema_decay", type=float, default=0.999) + parser.add_argument("--ema_start_step", type=int, default=0) + parser.add_argument('--cfg', type=float, default=0.1) + parser.add_argument( + "--weighting_scheme", + type=str, + default="uniform", + choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"], + ) + parser.add_argument( + "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--mode_scale", + type=float, + default=1.29, + help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", + ) + # validation & logs + parser.add_argument("--validation_prompt_dir", type=str) + parser.add_argument("--uncond_prompt_dir", type=str) + parser.add_argument("--validation_sampling_steps", type=int, default=64) + parser.add_argument('--validation_guidance_scale', type=float, default=4.5) + parser.add_argument('--validation_steps', type=float, default=4.5) + parser.add_argument("--log_validation", action="store_true") + parser.add_argument("--tracker_project_name", type=str, default=None) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument("--output_dir", type=str, default=None, help="The output directory where the model predictions and checkpoints will be written.") + parser.add_argument("--checkpoints_total_limit", type=int, default=None, help=("Max number of checkpoints to store.")) + parser.add_argument("--checkpointing_steps", type=int, default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument("--resume_from_checkpoint", type=str, default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument("--logging_dir", type=str, default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + + # optimizer & scheduler & Training + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform. If provided, overrides num_train_epochs.") + parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.") + parser.add_argument("--student_learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.") + parser.add_argument("--discriminator_learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.") + parser.add_argument("--discriminator_warmup_steps", type=int, default=0) + parser.add_argument("--gan_switch_n_steps", type=int, default=1) + parser.add_argument("--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.") + parser.add_argument("--lr_warmup_steps", type=int, default=10, help="Number of steps for the warmup in the lr scheduler.") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.") + parser.add_argument("--allow_tf32", action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + + parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel") + parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training") + + args = parser.parse_args() + main(args) diff --git a/fastvideo/fsdp_util.py b/fastvideo/fsdp_util.py index 8a0d82977..5b71f48cf 100644 --- a/fastvideo/fsdp_util.py +++ b/fastvideo/fsdp_util.py @@ -35,14 +35,33 @@ non_reentrant_wrapper = partial( check_fn = lambda submodule: isinstance(submodule, MochiTransformerBlock) -def apply_fsdp_checkpointing(model): +def apply_fsdp_checkpointing(model, p=1): + # https://github.com/foundation-model-stack/fms-fsdp/blob/408c7516d69ea9b6bcd4c0f5efab26c0f64b3c2d/fms_fsdp/policies/ac_handler.py#L16 """apply activation checkpointing to model returns None as model is updated directly """ print(f"--> applying fdsp activation checkpointing...") + block_idx = 0 + cut_off = 1 / 2 + # when passing p as a fraction number (e.g. 1/3), it will be interpreted + # as a string in argv, thus we need eval("1/3") here for fractions. + p = eval(p) if isinstance(p, str) else p + + def selective_checkpointing(submodule): + nonlocal block_idx + nonlocal cut_off + + if isinstance(submodule, MochiTransformerBlock): + block_idx += 1 + if block_idx * p >= cut_off: + cut_off += 1 + return True + return False + + apply_activation_checkpointing( - model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn + model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=selective_checkpointing ) bf16_mix = MixedPrecision( @@ -57,7 +76,7 @@ bf16_mix = MixedPrecision( -def get_fsdp_kwargs(sharding_strategy, use_lora=False, cpu_offload=False): +def get_dit_fsdp_kwargs(sharding_strategy, use_lora=False, cpu_offload=False): if use_lora: auto_wrap_policy = fsdp_auto_wrap_policy else: @@ -74,6 +93,8 @@ def get_fsdp_kwargs(sharding_strategy, use_lora=False, cpu_offload=False): if sharding_strategy == "full": sharding_strategy = ShardingStrategy.FULL_SHARD + elif sharding_strategy == "hybrid_full": + sharding_strategy = ShardingStrategy.HYBRID_SHARD elif sharding_strategy == "none": sharding_strategy = ShardingStrategy.NO_SHARD auto_wrap_policy = None @@ -101,4 +122,28 @@ def get_fsdp_kwargs(sharding_strategy, use_lora=False, cpu_offload=False): return fsdp_kwargs + + + +def get_discriminator_fsdp_kwargs(sharding_strategy): + + auto_wrap_policy = None + + + # Use existing mixed precision settings + + mixed_precision = bf16_mix + sharding_strategy = ShardingStrategy.NO_SHARD + device_id = torch.cuda.current_device() + fsdp_kwargs = { + "auto_wrap_policy": auto_wrap_policy, + "mixed_precision": mixed_precision, + "sharding_strategy": sharding_strategy, + "device_id": device_id, + "limit_all_gathers": True, + } + + return fsdp_kwargs + + \ No newline at end of file diff --git a/fastvideo/gan/discriminator.py b/fastvideo/gan/discriminator.py new file mode 100644 index 000000000..e8bdcd7cc --- /dev/null +++ b/fastvideo/gan/discriminator.py @@ -0,0 +1,123 @@ + +import torch +from torch import nn +import numpy as np +from torch.nn.utils.parametrizations import spectral_norm +import os +class DummyDiscriminator(nn.Module): + def __init__(self, dim_in, num_layers): + super().__init__() + self.layers = nn.ModuleList() + for _ in range(num_layers): + self.layers.append(nn.Linear(dim_in, 1)) + + def forward(self, features): + logits = [] + for layer, feature in zip(self.layers, features): + mean = feature.mean(dim=1) + logits.append(layer(mean)) + return torch.cat(logits, dim=1) + + +class ResidualBlock(nn.Module): + def __init__(self, fn): + super().__init__() + self.fn = fn + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (self.fn(x) + x) / np.sqrt(2) + + +class SpectralConv1d(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.conv = spectral_norm(nn.Conv1d(*args, **kwargs)) + def forward(self, x): + return self.conv(x) + +class BatchNormLocal(nn.Module): + def __init__(self, num_features: int, affine: bool = True, virtual_bs: int = 8, eps: float = 1e-5): + super().__init__() + self.virtual_bs = virtual_bs + self.eps = eps + self.affine = affine + + if self.affine: + self.weight = nn.Parameter(torch.ones(num_features)) + self.bias = nn.Parameter(torch.zeros(num_features)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shape = x.size() + + # Calculate stats. + mean = x.mean([0, 2], keepdim=True) + var = x.var([0, 2], keepdim=True, unbiased=False) + x = (x - mean) / (torch.sqrt(var + self.eps)) + + if self.affine: + x = x * self.weight[None, :, None] + self.bias[None, :, None] + + return x.view(shape) + +def make_block(channels: int, kernel_size: int) -> nn.Module: + return nn.Sequential( + SpectralConv1d( + channels, + channels, + kernel_size = kernel_size, + padding = kernel_size//2, + padding_mode = 'circular', + ), + BatchNormLocal(channels), + nn.LeakyReLU(0.2, True), + ) + +class DiscHead(nn.Module): + def __init__(self, feature_dim: int, text_c_dim: int, cmap_dim: int = 64, cnn_dim=512): + super().__init__() + self.channels = feature_dim + self.text_c_dim = text_c_dim + self.cmap_dim = cmap_dim + self.down_proj = SpectralConv1d(feature_dim, cnn_dim, kernel_size=1, padding=0) + self.main = nn.Sequential( + make_block(cnn_dim, kernel_size=1), + ResidualBlock(make_block(cnn_dim, kernel_size=9)) + ) + + self.cmapper = nn.Linear(self.text_c_dim, cmap_dim) + self.cls = SpectralConv1d(cnn_dim, cmap_dim, kernel_size=1, padding=0) + + def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: + h = self.down_proj(x) + h = self.main(h) + out = self.cls(h) + + cmap = self.cmapper(c).unsqueeze(-1) + out = (out * cmap).sum(1, keepdim=True) * (1 / np.sqrt(self.cmap_dim)) + + return out + + +class LADDDiscriminator(nn.Module): + def __init__(self, feature_dim, text_cond_dim, num_layers, layers_stride): + super().__init__() + heads = [] + for i in range(0, num_layers, layers_stride): + heads.append(DiscHead(feature_dim, text_cond_dim)) + self.heads = nn.ModuleList(heads) + self.layers_stride = layers_stride + self.num_layers = num_layers + + def forward(self, features, text_conditions) -> torch.Tensor: + text_conditions = text_conditions.mean(1) + # layer, B, L, C -> layer, B, C, L + features = features.transpose(2, 3) + logits = [] + for i in range(0, self.num_layers, self.layers_stride): + head = self.heads[i//self.layers_stride] + feat = features[i] + logits.append(head(feat, text_conditions).view(feat.size(0), -1)) + logits = torch.cat(logits, dim=1) + + + return logits diff --git a/fastvideo/model/mochi_latents_stat.py b/fastvideo/model/mochi_latents_stat.py deleted file mode 100644 index 915b57cc6..000000000 --- a/fastvideo/model/mochi_latents_stat.py +++ /dev/null @@ -1,33 +0,0 @@ -class MochiStat: - latents_mean = [ - -0.06730895953510081, - -0.038011381506090416, - -0.07477820912866141, - -0.05565264470995561, - 0.012767231469026969, - -0.04703542746246419, - 0.043896967884726704, - -0.09346305707025976, - -0.09918314763016893, - -0.008729793427399178, - -0.011931556316503654, - -0.0321993391887285 - ] - latents_std = [ - 0.9263795028493863, - 0.9248894543193766, - 0.9393059390890617, - 0.959253732819592, - 0.8244560132752793, - 0.917259975397747, - 0.9294154431013696, - 1.3720942357788521, - 0.881393668867029, - 0.9168315692124348, - 0.9185249279345552, - 0.9274757570805041 - ] - scaling_factor = 1.0 - -mochi_stat = MochiStat() - \ No newline at end of file diff --git a/fastvideo/model/mochi_latents_utils.py b/fastvideo/model/mochi_latents_utils.py new file mode 100644 index 000000000..0919e4b1b --- /dev/null +++ b/fastvideo/model/mochi_latents_utils.py @@ -0,0 +1,39 @@ +import torch +mochi_latents_mean = torch.tensor([ + -0.06730895953510081, + -0.038011381506090416, + -0.07477820912866141, + -0.05565264470995561, + 0.012767231469026969, + -0.04703542746246419, + 0.043896967884726704, + -0.09346305707025976, + -0.09918314763016893, + -0.008729793427399178, + -0.011931556316503654, + -0.0321993391887285 + ]).view(1, 12, 1, 1, 1) +mochi_latents_std = torch.tensor([ + 0.9263795028493863, + 0.9248894543193766, + 0.9393059390890617, + 0.959253732819592, + 0.8244560132752793, + 0.917259975397747, + 0.9294154431013696, + 1.3720942357788521, + 0.881393668867029, + 0.9168315692124348, + 0.9185249279345552, + 0.9274757570805041 + ]).view(1, 12, 1, 1, 1) +mochi_scaling_factor = 1.0 + + +def normalize_mochi_dit_input(latents): + latents_mean = mochi_latents_mean.to(latents.device, latents.dtype) + latents_std = mochi_latents_std.to(latents.device, latents.dtype) + latents = (latents - latents_mean) / latents_std + return latents + + \ No newline at end of file diff --git a/fastvideo/model/modeling_mochi.py b/fastvideo/model/modeling_mochi.py index 3c74b47e0..57c0e22c4 100644 --- a/fastvideo/model/modeling_mochi.py +++ b/fastvideo/model/modeling_mochi.py @@ -16,21 +16,87 @@ from typing import Any, Dict, Optional, Tuple import torch import torch.nn as nn - +import diffusers from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.utils import is_torch_version, logging from diffusers.utils.torch_utils import maybe_allow_in_graph -from diffusers.models.attention import FeedForward +from diffusers.models.attention import FeedForward as HF_FeedForward from diffusers.models.attention_processor import Attention from diffusers.models.embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.modeling_utils import ModelMixin from diffusers.models.normalization import AdaLayerNormContinuous, LuminaLayerNormContinuous, MochiRMSNormZero, RMSNorm + + from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info from fastvideo.utils.communications import all_gather, all_to_all_4D import torch.nn.functional as F from diffusers.utils.torch_utils import is_torch_version, maybe_allow_in_graph +from einops import rearrange + +import numbers +from flash_attn import flash_attn_varlen_qkvpacked_func +from flash_attn.bert_padding import pad_input, unpad_input + +from liger_kernel.ops.swiglu import LigerSiLUMulFunction +class FeedForward(HF_FeedForward): + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + dropout: float = 0.0, + activation_fn: str = "geglu", + final_dropout: bool = False, + inner_dim=None, + bias: bool = True, + ): + super().__init__(dim, dim_out, mult, dropout, activation_fn, final_dropout, inner_dim, bias) + assert activation_fn == "swiglu" + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.net[0].proj(hidden_states) + hidden_states, gate = hidden_states.chunk(2, dim=-1) + + return self.net[2]( + LigerSiLUMulFunction.apply(gate, hidden_states) + ) + +def flash_attn_no_pad(qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None): + # adapted from https://github.com/Dao-AILab/flash-attention/blob/13403e81157ba37ca525890f2f0f2137edf75311/flash_attn/flash_attention.py#L27 + batch_size = qkv.shape[0] + seqlen = qkv.shape[1] + nheads = qkv.shape[-2] + x = rearrange(qkv, 'b s three h d -> b s (three h d)') + x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(x, key_padding_mask) + + + x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads) + output_unpad = flash_attn_varlen_qkvpacked_func( + x_unpad, cu_seqlens, max_s, dropout_p, + softmax_scale=softmax_scale, causal=causal + ) + output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), + indices, batch_size, seqlen), + 'b s (h d) -> b s h d', h=nheads) + return output + + logger = logging.get_logger(__name__) # pylint: disable=invalid-name +import pdb, sys +class ForkedPdb(pdb.Pdb): + """A Pdb subclass that may be used + from a forked multiprocessing child + + """ + def interaction(self, *args, **kwargs): + _stdin = sys.stdin + try: + sys.stdin = open('/dev/stdin') + pdb.Pdb.interaction(self, *args, **kwargs) + finally: + sys.stdin = _stdin + class MochiAttnProcessor2_0: @@ -45,6 +111,7 @@ class MochiAttnProcessor2_0: attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, image_rotary_emb: Optional[torch.Tensor] = None, ) -> torch.Tensor: @@ -113,38 +180,47 @@ class MochiAttnProcessor2_0: query = apply_rotary_emb(query, freqs_cos, freqs_sin) key = apply_rotary_emb(key, freqs_cos, freqs_sin) - query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) - encoder_query, encoder_key, encoder_value = ( - encoder_query.transpose(1, 2), - encoder_key.transpose(1, 2), - encoder_value.transpose(1, 2), - ) - # [b, h, s, d] - sequence_length = query.size(2) - encoder_sequence_length = encoder_query.size(2) + # query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) + # encoder_query, encoder_key, encoder_value = ( + # encoder_query.transpose(1, 2), + # encoder_key.transpose(1, 2), + # encoder_value.transpose(1, 2), + # ) + # [b, s, h, d] + sequence_length = query.size(1) + encoder_sequence_length = encoder_query.size(1) - # Hint: please check encoder_query.shape - query = torch.cat([query, encoder_query], dim=2) - key = torch.cat([key, encoder_key], dim=2) - value = torch.cat([value, encoder_value], dim=2) - - - - hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + # H + query = torch.cat([query, encoder_query], dim=1).unsqueeze(2) + key = torch.cat([key, encoder_key], dim=1).unsqueeze(2) + value = torch.cat([value, encoder_value], dim=1).unsqueeze(2) + # B, S, 3, H, D + qkv = torch.cat([query, key, value], dim=2) + attn_mask = encoder_attention_mask[:, :].bool() + attn_mask = F.pad(attn_mask, (sequence_length, 0), value=True) + # hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) + + hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask = None, dropout_p=0.0, is_causal=False) + + # valid_lengths = encoder_attention_mask.sum(dim=1) + sequence_length + # def no_padding_mask(score, b, h, q_idx, kv_idx): + # return torch.where(kv_idx < valid_lengths[b],score, -float("inf")) + + # hidden_states = flex_attention(query, key, value, score_mod=no_padding_mask) if get_sequence_parallel_state(): hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( - (sequence_length, encoder_sequence_length), dim=2 + (sequence_length, encoder_sequence_length), dim=1 ) - # B, H, S, D - hidden_states = all_to_all_4D(hidden_states, scatter_dim=2, gather_dim=1) - encoder_hidden_states = all_gather(encoder_hidden_states, dim=1).contiguous() - hidden_states = hidden_states.transpose(1,2).flatten(2, 3) + # B, S, H, D + hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2) + encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous() + hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) - encoder_hidden_states = encoder_hidden_states.transpose(1, 2).flatten(2, 3) + encoder_hidden_states = encoder_hidden_states.flatten(2, 3) encoder_hidden_states = encoder_hidden_states.to(query.dtype) else: - hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) + hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( @@ -257,6 +333,7 @@ class MochiTransformerBlock(nn.Module): self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, temb: torch.Tensor, image_rotary_emb: Optional[torch.Tensor] = None, output_attn = False, @@ -274,6 +351,7 @@ class MochiTransformerBlock(nn.Module): hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, image_rotary_emb=image_rotary_emb, + encoder_attention_mask=encoder_attention_mask ) hidden_states = hidden_states + self.norm2(attn_hidden_states) * torch.tanh(gate_msa).unsqueeze(1) @@ -498,6 +576,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin): create_custom_forward(block), hidden_states, encoder_hidden_states, + encoder_attention_mask, temb, image_rotary_emb, output_attn, @@ -507,6 +586,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin): hidden_states, encoder_hidden_states, attn_outputs = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, temb=temb, image_rotary_emb=image_rotary_emb, output_attn = output_attn, diff --git a/fastvideo/model/norm.py b/fastvideo/model/norm.py new file mode 100644 index 000000000..c8df72ca9 --- /dev/null +++ b/fastvideo/model/norm.py @@ -0,0 +1,168 @@ +import numbers +from typing import Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from liger_kernel.ops.rms_norm import LigerRMSNormFunction + +from liger_kernel.ops.layer_norm import LigerLayerNormFunction + +class RMSNorm(nn.Module): + # mbedding_dim, eps=eps, elementwise_affine=elementwise_affine) + def __init__(self, dim, eps: float, elementwise_affine: bool = True): + super().__init__() + + self.eps = eps + + if isinstance(dim, numbers.Integral): + dim = (dim,) + + self.dim = torch.Size(dim) + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim)) + else: + weight = torch.ones(dim) + self.register_buffer("weight", weight, persistent=False) + + def forward(self, hidden_states): + + return LigerRMSNormFunction.apply( + hidden_states, + self.weight, + self.eps, + 0.0, + "gemma", + True, + ) + +class LayerNorm(nn.Module): + def __init__(self, dim, eps: float = 1e-5, elementwise_affine: bool = True, bias: bool = True): + super().__init__() + + self.eps = eps + + if isinstance(dim, numbers.Integral): + dim = (dim,) + + self.dim = torch.Size(dim) + + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim)) + self.bias = nn.Parameter(torch.zeros(dim)) if bias else None + else: + weight = torch.ones(dim) + bias = torch.zeros(dim) + self.register_buffer("weight", weight, persistent=False) + self.register_buffer("bias", bias, persistent=False) + def forward(self, hidden_states): + return LigerLayerNormFunction.apply( + hidden_states, self.weight, self.bias, self.eps + ) + +class MochiRMSNormZero(nn.Module): + r""" + Adaptive RMS Norm used in Mochi. + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + """ + + def __init__( + self, embedding_dim: int, hidden_dim: int, eps: float = 1e-5, elementwise_affine: bool = False + ) -> None: + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, hidden_dim) + # TODO: There is a bug in HF. Remove hardcode + self.norm = RMSNorm(hidden_dim//4, eps=eps, elementwise_affine=elementwise_affine) + + def forward( + self, hidden_states: torch.Tensor, emb: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb = self.linear(self.silu(emb)) + scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1) + hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + + return hidden_states, gate_msa, scale_mlp, gate_mlp + + +class LuminaLayerNormContinuous(nn.Module): + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters + # because the output is immediately scaled and shifted by the projected conditioning embeddings. + # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters. + # However, this is how it was implemented in the original code, and it's rather likely you should + # set `elementwise_affine` to False. + elementwise_affine=True, + eps=1e-5, + bias=True, + norm_type="layer_norm", + out_dim: Optional[int] = None, + ): + super().__init__() + + # AdaLN + self.silu = nn.SiLU() + self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias) + + if norm_type == "layer_norm": + self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias) + elif norm_type == "rms_norm": + self.norm = RMSNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + self.linear_2 = None + if out_dim is not None: + self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias) + def forward( + self, + x: torch.Tensor, + conditioning_embedding: torch.Tensor, + ) -> torch.Tensor: + # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) + emb = self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) + scale = emb + x = self.norm(x) * (1 + scale)[:, None, :] + + if self.linear_2 is not None: + x = self.linear_2(x) + + return x + +class AdaLayerNormContinuous(nn.Module): + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters + # because the output is immediately scaled and shifted by the projected conditioning embeddings. + # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters. + # However, this is how it was implemented in the original code, and it's rather likely you should + # set `elementwise_affine` to False. + elementwise_affine=True, + eps=1e-5, + bias=True, + norm_type="layer_norm", + ): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias) + if norm_type == "layer_norm": + self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias) + elif norm_type == "rms_norm": + self.norm = RMSNorm(embedding_dim, eps, elementwise_affine) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: + # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) + emb = self.linear(self.silu(conditioning_embedding).to(x.dtype)) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x diff --git a/fastvideo/sample/sample_t2v_mochi.py b/fastvideo/sample/sample_t2v_mochi.py index 372325525..3fed389b3 100644 --- a/fastvideo/sample/sample_t2v_mochi.py +++ b/fastvideo/sample/sample_t2v_mochi.py @@ -17,19 +17,6 @@ import pdb import copy from typing import Dict -class ForkedPdb(pdb.Pdb): - """A Pdb subclass that may be used - from a forked multiprocessing child - - """ - def interaction(self, *args, **kwargs): - _stdin = sys.stdin - try: - sys.stdin = open('/dev/stdin') - pdb.Pdb.interaction(self, *args, **kwargs) - finally: - sys.stdin = _stdin - def initialize_distributed(): local_rank = int(os.getenv('RANK', 0)) world_size = int(os.getenv('WORLD_SIZE', 1)) @@ -121,7 +108,6 @@ def main(args): if args.transformer_path is not None: transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path, torch_dtype=torch.bfloat16) - else: transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder = 'transformer/', torch_dtype=torch.bfloat16) if args.lora_checkpoint_dir is not None: @@ -163,7 +149,7 @@ def main(args): if prompts is not None: for video, prompt in zip(videos, prompts): suffix = prompt.split(".")[0] - export_to_video(video, args.output_path + f"_{suffix}.mp4", fps=30) + export_to_video(video, args.output_path + f"{suffix}.mp4", fps=30) else: export_to_video(videos[0], args.output_path + ".mp4", fps=30) diff --git a/fastvideo/train.py b/fastvideo/train.py index b392cae5f..aadc283ca 100644 --- a/fastvideo/train.py +++ b/fastvideo/train.py @@ -9,7 +9,7 @@ from diffusers.training_utils import cast_training_params, compute_density_for_t from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \ destroy_sequence_parallel_group, get_sequence_parallel_state, nccl_info from fastvideo.utils.communications import sp_parallel_dataloader_wrapper, broadcast -from fastvideo.model.mochi_latents_stat import mochi_stat +from fastvideo.model.mochi_latents_utils import normalize_mochi_dit_input from fastvideo.utils.validation import log_validation import time from torch.utils.data import DataLoader @@ -24,7 +24,7 @@ from torch.utils.data.distributed import DistributedSampler import wandb from accelerate.utils import set_seed from tqdm.auto import tqdm -from fastvideo.fsdp_util import get_fsdp_kwargs, apply_fsdp_checkpointing +from fastvideo.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpointing import diffusers from diffusers import ( FlowMatchEulerDiscreteScheduler, @@ -102,25 +102,22 @@ def get_sigmas(noise_scheduler, device, timesteps, n_dim=4, dtype=torch.float32) return sigma -def train_one_step_mochi(transformer, optimizer, lr_scheduler, loader,noise_scheduler, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm): +def train_one_step_mochi(transformer, optimizer, lr_scheduler, loader,noise_scheduler, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm, weighting_scheme, logit_mean, logit_std, mode_scale): total_loss = 0.0 optimizer.zero_grad() for _ in range(gradient_accumulation_steps): latents, encoder_hidden_states, latents_attention_mask, encoder_attention_mask = next(loader) - - latents = latents.to(dtype=torch.bfloat16) - encoder_hidden_states = encoder_hidden_states.to(dtype=torch.bfloat16) - latents_mean = ( - torch.tensor(mochi_stat.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(mochi_stat.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype) - ) - latents = (latents - latents_mean) / latents_std + latents = normalize_mochi_dit_input(latents) batch_size = latents.shape[0] noise = torch.randn_like(latents) - u = torch.rand(size=(batch_size,), device="cpu") + u = compute_density_for_timestep_sampling( + weighting_scheme=weighting_scheme, + batch_size=batch_size, + logit_mean=logit_mean, + logit_std=logit_std, + mode_scale=mode_scale, + ) indices = (u * noise_scheduler.config.num_train_timesteps).long() timesteps = noise_scheduler.timesteps[indices].to(device=latents.device) if sp_size > 1: @@ -145,21 +142,19 @@ def train_one_step_mochi(transformer, optimizer, lr_scheduler, loader,noise_sche else: target = latents - noise - loss = torch.mean((model_pred.float() - target.float()) ** 2) - + loss = torch.mean((model_pred.float() - target.float()) ** 2) / gradient_accumulation_steps + loss.backward() avg_loss = loss.detach().clone() dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG) - total_loss += avg_loss.item() / gradient_accumulation_steps + total_loss += avg_loss.item() - transformer.clip_grad_norm_(max_grad_norm) - + grad_norm = transformer.clip_grad_norm_(max_grad_norm) optimizer.step() lr_scheduler.step() - - return total_loss + return total_loss, grad_norm.item() def get_lora_model(transformer, lora_config): transformer.requires_grad_(False) @@ -286,8 +281,9 @@ def main(args): transformer = get_lora_model(transformer, lora_config) main_print(f" Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M") - main_print(f"--> Initializing FSDP with sharding strategy: full") - fsdp_kwargs = get_fsdp_kwargs("full", args.use_lora, args.use_cpu_offload) + main_print(f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}") + fsdp_kwargs = get_dit_fsdp_kwargs(args.fsdp_sharding_startegy, args.use_lora, args.use_cpu_offload) + if args.use_lora: transformer.config.lora_rank = args.lora_rank @@ -304,7 +300,7 @@ def main(args): main_print(f"--> model loaded") if args.gradient_checkpointing: - apply_fsdp_checkpointing(transformer) + apply_fsdp_checkpointing(transformer, args.selective_checkpointing) # Set model as trainable. transformer.train() @@ -312,8 +308,8 @@ def main(args): noise_scheduler = FlowMatchEulerDiscreteScheduler() params_to_optimize = transformer.parameters() - if args.use_lora: - params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize)) + params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize)) + optimizer = torch.optim.AdamW( params_to_optimize, lr=args.learning_rate, @@ -352,13 +348,10 @@ def main(args): drop_last=True, ) - # We need to recalculate our total training steps as the size of the training dataloader may have changed. num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size) - # Afterwards we recalculate our number of training epochs args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) - # We need to initialize the trackers we use, and also store our configuration. - # The trackers initializes automatically on the main process. + if rank <= 0: project = args.tracker_project_name or "fastvideo" wandb.init(project=project, config=args) @@ -402,7 +395,7 @@ def main(args): next(loader) for step in range(init_steps + 1, args.max_train_steps+1): start_time = time.time() - loss = train_one_step_mochi(transformer, optimizer, lr_scheduler, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, args.max_grad_norm) + loss, grad_norm= train_one_step_mochi(transformer, optimizer, lr_scheduler, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, args.max_grad_norm, args.weighting_scheme, args.logit_mean, args.logit_std, args.mode_scale) step_time = time.time() - start_time step_times.append(step_time) @@ -410,7 +403,8 @@ def main(args): progress_bar.set_postfix({ "loss": f"{loss:.4f}", - "step_time": f"{step_time:.2f}s" + "step_time": f"{step_time:.2f}s", + "grad_norm": grad_norm }) progress_bar.update(1) if rank <= 0: @@ -418,20 +412,25 @@ def main(args): "train_loss": loss, "learning_rate": lr_scheduler.get_last_lr()[0], "step_time": step_time, - "avg_step_time": avg_step_time + "avg_step_time": avg_step_time, + "grad_norm": grad_norm }, step=step) if step % args.checkpointing_steps == 0: - save_checkpoint(transformer, rank, args.output_dir, step) if args.use_lora: # Save LoRA weights save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step) - + else: + # Your existing checkpoint saving code + save_checkpoint(transformer, rank, args.output_dir, step) + dist.barrier() if args.log_validation and step % args.validation_steps == 0: log_validation(args, transformer, device, torch.bfloat16, step) if args.use_lora: save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps) + else: + save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps) if get_sequence_parallel_state(): destroy_sequence_parallel_group() @@ -506,6 +505,7 @@ if __name__ == "__main__": parser.add_argument("--lr_warmup_steps", type=int, default=10, help="Number of steps for the warmup in the lr scheduler.") parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.") + parser.add_argument("--selective_checkpointing", type=float, default=1.0) parser.add_argument("--allow_tf32", action="store_true", help=( "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" @@ -527,8 +527,28 @@ if __name__ == "__main__": parser.add_argument("--use_lora", action="store_true", default=False, help="Whether to use LoRA for finetuning.") parser.add_argument("--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA.") parser.add_argument("--lora_rank", type=int, default=128, help="LoRA rank parameter. ") + parser.add_argument("--fsdp_sharding_startegy", default="full") + + parser.add_argument( + "--weighting_scheme", + type=str, + default="uniform", + choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"], + ) + parser.add_argument( + "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--mode_scale", + type=float, + default=1.29, + help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", + ) # lr_scheduler - parser.add_argument("--lr_scheduler", type=str, default="cosine_with_restarts", + parser.add_argument("--lr_scheduler", type=str, default="constant", help=( 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' ' "constant", "constant_with_warmup"]' diff --git a/fastvideo/utils/encode_vae_t5.py b/fastvideo/utils/encode_vae_t5.py deleted file mode 100644 index ce908ca69..000000000 --- a/fastvideo/utils/encode_vae_t5.py +++ /dev/null @@ -1,867 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. - -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -""" -A minimal training script for DiT using PyTorch DDP. -""" -import argparse -import logging -import math -import os -import shutil -from pathlib import Path -import gc -import wandb -from einops import rearrange -from tqdm import tqdm - -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \ - destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state -from fastvideo.utils.communications import prepare_parallel_data, broadcast -import time -from torch.utils.data import DataLoader -from copy import deepcopy -import accelerate -import torch -from torch.nn import functional as F -import transformers -from accelerate import Accelerator -from accelerate.logging import get_logger -from accelerate.utils import DistributedType, ProjectConfiguration, set_seed -from packaging import version -from tqdm.auto import tqdm - -import diffusers -from diffusers import ( - AutoencoderKL, - FlowMatchEulerDiscreteScheduler, -) -from fastvideo.model.modeling_mochi import MochiTransformer3DModel -from fastvideo.model.pipeline_mochi import MochiPipeline -from diffusers.optimization import get_scheduler -from diffusers.training_utils import compute_snr -from diffusers.utils import check_min_version, is_wandb_available - -from fastvideo.utils.ema import EMAModel -from fastvideo.dataset import getdataset -from fastvideo.utils.dataset_utils import Collate, LengthGroupedSampler - - -# Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.24.0") -logger = get_logger(__name__) - - -@torch.inference_mode() -def log_validation(args, model, vae, text_encoder, tokenizer, accelerator, weight_dtype, global_step, ema=False): - logger.info(f"Running validation....\n") - model = accelerator.unwrap_model(model) - scheduler = FlowMatchEulerDiscreteScheduler() - mochi_pipeline = MochiPipeline(vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - scheduler=scheduler, - transformer=model).to(device=accelerator.device) - videos = [] - for prompt in args.validaiton_prompt: - logger.info('Processing the ({}) prompt'.format(prompt)) - video = mochi_pipeline( - args.prompt, - num_frames=args.num_frames, - height=args.max_height, - width=args.max_width, - num_inference_steps=args.validation_sampling_steps, - guidance_scale=args.guidance_scale, - ).images - videos.append(video[0]) - # import ipdb;ipdb.set_trace() - gc.collect() - torch.cuda.empty_cache() - videos = torch.stack(videos).numpy() - videos = rearrange(videos, 'b t h w c -> b t c h w') - for tracker in accelerator.trackers: - if videos.shape[1] == 1: - images = rearrange(videos, 'b 1 c h w -> (b 1) h w c') - logs = { - f"{'ema_' if ema else ''}validation": [ - wandb.Image(image, caption=f"{i}: {prompt}") - for i, (image, prompt) in enumerate(zip(images, args.validaiton_prompt)) - ] - } - else: - logs = { - f"{'ema_' if ema else ''}validation": [ - wandb.Video(video, caption=f"{i}: {prompt}", fps=24) - for i, (video, prompt) in enumerate(zip(videos, args.validaiton_prompt)) - ] - } - tracker.log(logs, step=global_step) - - del mochi_pipeline - gc.collect() - torch.cuda.empty_cache() - - -class ProgressInfo: - def __init__(self, global_step, train_loss=0.0): - self.global_step = global_step - self.train_loss = train_loss - - -################################################################################# -# Training Loop # -################################################################################# - -def main(args): - logging_dir = Path(args.output_dir, args.logging_dir) - - # use LayerNorm, GeLu, SiLu always as fp32 mode - # TODO: - if args.enable_stable_fp32: - replace_with_fp32_forwards() - - accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) - - accelerator = Accelerator( - gradient_accumulation_steps=args.gradient_accumulation_steps, - mixed_precision=args.mixed_precision, - log_with=args.report_to, - project_config=accelerator_project_config, - log_with="wandb" - ) - - if args.num_frames != 1 and args.use_image_num == 0: - initialize_sequence_parallel_state(args.sp_size) - - 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.") - - # Make one log on every process with the configuration for debugging. - 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 passed along, set the training seed now. On GPU... - if args.seed is not None: - set_seed(args.seed) - - # Handle the repository creation - if accelerator.is_main_process: - if args.output_dir is not None: - os.makedirs(args.output_dir, exist_ok=True) - - # For mixed precision training we cast all non-trainable weigths 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 - - # Create model: - kwargs = {} - ae = CausalVAEModelWrapper(args.ae_path, cache_dir=args.cache_dir, **kwargs).eval() - if args.enable_tiling: - ae.vae.enable_tiling() - ae.vae.tile_sample_stride = args.tile_sample_stride - - kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True} - text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval() - - ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae] - ae.vae_scale_factor = (ae_stride_t, ae_stride_h, ae_stride_w) - assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})" - args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w - args.ae_stride = args.ae_stride_h - patch_size = args.model[-3:] - patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) - args.patch_size = patch_size_h - args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w - assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})" - # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})." - assert args.max_height % ae_stride_h == 0, f"Height must be divisible by ae_stride_h, but found Height ({args.max_height}), ae_stride_h ({ae_stride_h})." - assert args.max_width % ae_stride_h == 0, f"Width size must be divisible by ae_stride_h, but found Width ({args.max_width}), ae_stride_h ({ae_stride_h})." - - args.stride_t = ae_stride_t * patch_size_t - args.stride = ae_stride_h * patch_size_h - latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w) - ae.latent_size = latent_size - - if args.num_frames % 2 == 1: - args.latent_size_t = latent_size_t = (args.num_frames - 1) // ae_stride_t + 1 - else: - latent_size_t = args.num_frames // ae_stride_t - model = Diffusion_models[args.model]( - in_channels=ae_channel_config[args.ae], - out_channels=ae_channel_config[args.ae], - # caption_channels=4096, - # cross_attention_dim=1152, - attention_bias=True, - sample_size=latent_size, - sample_size_t=latent_size_t, - num_vector_embeds=None, - activation_fn="gelu-approximate", - num_embeds_ada_norm=1000, - use_linear_projection=False, - only_cross_attention=False, - double_self_attention=False, - upcast_attention=False, - # norm_type="ada_norm_single", - norm_elementwise_affine=False, - norm_eps=1e-6, - attention_type='default', - attention_mode=args.attention_mode, - interpolation_scale_h=args.interpolation_scale_h, - interpolation_scale_w=args.interpolation_scale_w, - interpolation_scale_t=args.interpolation_scale_t, - downsampler=args.downsampler, - use_rope=args.use_rope, - # text_max_length=args.text_max_length, - use_stable_fp32=args.enable_stable_fp32, - ) - model.gradient_checkpointing = args.gradient_checkpointing - - # # use pretrained model? - if args.pretrained: - model_state_dict = model.state_dict() - if 'safetensors' in args.pretrained: # pixart series - from safetensors.torch import load_file as safe_load - # import ipdb;ipdb.set_trace() - pretrained_checkpoint = safe_load(args.pretrained, device="cpu") - pretrained_keys = set(list(pretrained_checkpoint.keys())) - model_keys = set(list(model_state_dict.keys())) - common_keys = list(pretrained_keys & model_keys) - checkpoint = {k: pretrained_checkpoint[k] for k in common_keys if model_state_dict[k].numel() == pretrained_checkpoint[k].numel()} - # if checkpoint['pos_embed.proj.weight'].shape != model.pos_embed.proj.weight.shape and checkpoint['pos_embed.proj.weight'].ndim == 4: - # logger.info(f"Resize pos_embed, {checkpoint['pos_embed.proj.weight'].shape} -> {model.pos_embed.proj.weight.shape}") - # repeat = model.pos_embed.proj.weight.shape[2] - # checkpoint['pos_embed.proj.weight'] = checkpoint['pos_embed.proj.weight'].unsqueeze(2).repeat(1, 1, repeat, 1, 1) / float(repeat) - # del checkpoint['proj_out.weight'], checkpoint['proj_out.bias'] - else: # latest stage training weight - checkpoint = torch.load(args.pretrained, map_location='cpu') - if 'model' in checkpoint: - checkpoint = checkpoint['model'] - missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False) - logger.info(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}') - logger.info(f'Successfully load {len(model_state_dict) - len(missing_keys)}/{len(model_state_dict)} keys from {args.pretrained}!') - - # Freeze vae and text encoders. - ae.vae.requires_grad_(False) - text_enc.requires_grad_(False) - # Set model as trainable. - model.train() - - noise_scheduler = DDPMScheduler() - # Move unet, vae and text_encoder to device and cast to weight_dtype - # The VAE is in float32 to avoid NaN losses. - ae.vae.to(accelerator.device, dtype=torch.float32) - # ae.vae.to(accelerator.device, dtype=weight_dtype) - text_enc.to(accelerator.device, dtype=weight_dtype) - - # Create EMA for the unet. - if args.use_ema: - ema_model = deepcopy(model) - ema_model = EMAModel(ema_model.parameters(), decay=args.ema_decay, update_after_step=args.ema_start_step, - model_cls=Diffusion_models_class[args.model], model_config=ema_model.config) - - # `accelerate` 0.16.0 will have better support for customized saving - 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 args.use_ema: - ema_model.save_pretrained(os.path.join(output_dir, "model_ema")) - - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "model")) - if weights: # Don't pop if empty - # make sure to pop weight so that corresponding model is not saved again - weights.pop() - - def load_model_hook(models, input_dir): - if args.use_ema: - load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), Diffusion_models_class[args.model]) - ema_model.load_state_dict(load_model.state_dict()) - ema_model.to(accelerator.device) - del load_model - - for i in range(len(models)): - # pop models so that they are not loaded again - model = models.pop() - - # load diffusers style into model - load_model = Diffusion_models_class[args.model].from_pretrained(input_dir, subfolder="model") - model.register_to_config(**load_model.config) - - model.load_state_dict(load_model.state_dict()) - del load_model - - accelerator.register_save_state_pre_hook(save_model_hook) - accelerator.register_load_state_pre_hook(load_model_hook) - - # Enable TF32 for faster training on Ampere GPUs, - # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices - if args.allow_tf32: - torch.backends.cuda.matmul.allow_tf32 = True - - if args.scale_lr: - args.learning_rate = ( - args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes - ) - - params_to_optimize = model.parameters() - # Optimizer creation - if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"): - logger.warning( - f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]." - "Defaulting to adamW" - ) - args.optimizer = "adamw" - - if args.use_8bit_adam and not args.optimizer.lower() == "adamw": - logger.warning( - f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was " - f"set to {args.optimizer.lower()}" - ) - - if args.optimizer.lower() == "adamw": - if args.use_8bit_adam: - try: - import bitsandbytes as bnb - except ImportError: - raise ImportError( - "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." - ) - - optimizer_class = bnb.optim.AdamW8bit - else: - optimizer_class = torch.optim.AdamW - - optimizer = optimizer_class( - params_to_optimize, - lr=args.learning_rate, - betas=(args.adam_beta1, args.adam_beta2), - weight_decay=args.adam_weight_decay, - eps=args.adam_epsilon, - ) - - if args.optimizer.lower() == "prodigy": - try: - import prodigyopt - except ImportError: - raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") - - optimizer_class = prodigyopt.Prodigy - - if args.learning_rate <= 0.1: - logger.warning( - "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" - ) - - optimizer = optimizer_class( - params_to_optimize, - lr=args.learning_rate, - betas=(args.adam_beta1, args.adam_beta2), - beta3=args.prodigy_beta3, - weight_decay=args.adam_weight_decay, - eps=args.adam_epsilon, - decouple=args.prodigy_decouple, - use_bias_correction=args.prodigy_use_bias_correction, - safeguard_warmup=args.prodigy_safeguard_warmup, - ) - logger.info(f"optimizer: {optimizer}") - - # Setup data: - train_dataset = getdataset(args) - sampler = LengthGroupedSampler( - args.train_batch_size, - world_size=accelerator.num_processes, - lengths=train_dataset.lengths, - group_frame=args.group_frame, - group_resolution=args.group_resolution, - ) if (args.group_frame or args.group_resolution) else None - train_dataloader = DataLoader( - train_dataset, - shuffle=sampler is None, - # pin_memory=True, - collate_fn=Collate(args), - batch_size=args.train_batch_size, - num_workers=args.dataloader_num_workers, - sampler=sampler if (args.group_frame or args.group_resolution) else None, - drop_last=True, - # prefetch_factor=4 - ) - logger.info(f'after train_dataloader') - - # Scheduler and math around the number of training steps. - overrode_max_train_steps = False - num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) - if args.max_train_steps is None: - args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch - overrode_max_train_steps = True - - lr_scheduler = get_scheduler( - args.lr_scheduler, - optimizer=optimizer, - num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, - num_training_steps=args.max_train_steps * args.gradient_accumulation_steps, - ) - - # Prepare everything with our `accelerator`. - # model.requires_grad_(False) - # model.pos_embed.requires_grad_(True) - logger.info(f'before accelerator.prepare') - model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( - model, optimizer, train_dataloader, lr_scheduler - ) - logger.info(f'after accelerator.prepare') - if args.use_ema: - ema_model.to(accelerator.device) - - # We need to recalculate our total training steps as the size of the training dataloader may have changed. - num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) - if overrode_max_train_steps: - args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch - # Afterwards we recalculate our number of training epochs - args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) - - # We need to initialize the trackers we use, and also store our configuration. - # The trackers initializes automatically on the main process. - if accelerator.is_main_process: - accelerator.init_trackers(os.path.basename(args.output_dir), config=vars(args)) - - # Train! - total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps - total_batch_size = total_batch_size // args.sp_size * args.train_sp_batch_size - logger.info("***** Running training *****") - logger.info(f" Num examples = {len(train_dataset)}") - logger.info(f" Num Epochs = {args.num_train_epochs}") - logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") - logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") - logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") - logger.info(f" Total optimization steps = {args.max_train_steps}") - logger.info(f" Total training parameters = {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9} B") - global_step = 0 - first_epoch = 0 - - # Potentially load in the weights and states from a previous save - if args.resume_from_checkpoint: - if args.resume_from_checkpoint != "latest": - path = os.path.basename(args.resume_from_checkpoint) - 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 '{args.resume_from_checkpoint}' does not exist. Starting a new training run." - ) - args.resume_from_checkpoint = None - 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 - first_epoch = global_step // num_update_steps_per_epoch - - - - else: - initial_global_step = 0 - - progress_bar = tqdm( - range(0, args.max_train_steps), - initial=initial_global_step, - desc="Steps", - # Only show the progress bar once on each machine. - disable=not accelerator.is_local_main_process, - ) - progress_info = ProgressInfo(global_step, train_loss=0.0) - - def sync_gradients_info(loss): - # Checks if the accelerator has performed an optimization step behind the scenes - if args.use_ema: - ema_model.step(model.parameters()) - progress_bar.update(1) - progress_info.global_step += 1 - end_time = time.time() - one_step_duration = end_time - start_time - accelerator.log({"train_loss": progress_info.train_loss}, step=progress_info.global_step) - progress_info.train_loss = 0.0 - - # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues. - if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process: - if progress_info.global_step % args.checkpointing_steps == 0: - # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` - if accelerator.is_main_process and args.checkpoints_total_limit is not None: - checkpoints = os.listdir(args.output_dir) - checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] - checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) - - # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints - if len(checkpoints) >= args.checkpoints_total_limit: - num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 - removing_checkpoints = checkpoints[0:num_to_remove] - - logger.info( - f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" - ) - logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") - - for removing_checkpoint in removing_checkpoints: - removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) - shutil.rmtree(removing_checkpoint) - - save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}") - accelerator.save_state(save_path) - logger.info(f"Saved state to {save_path}") - - logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} - progress_bar.set_postfix(**logs) - - def run(model_input, model_kwargs, prof): - global start_time - start_time = time.time() - - noise = torch.randn_like(model_input) - if args.noise_offset: - # https://www.crosslabs.org//blog/diffusion-with-offset-noise - noise += args.noise_offset * torch.randn((model_input.shape[0], model_input.shape[1], 1, 1, 1), - device=model_input.device) - - bsz = model_input.shape[0] - current_step_frame = model_input.shape[2] - # Sample a random timestep for each image without bias. - timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device) - if current_step_frame != 1 and get_sequence_parallel_state(): # image do not need sp - broadcast(timesteps) - - # Add noise to the model input according to the noise magnitude at each timestep - # (this is the forward diffusion process) - - noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps) - - model_pred = model( - noisy_model_input, - timesteps, - **model_kwargs - )[0] - # Get the target for loss depending on the prediction type - if args.prediction_type is not None: - # set prediction_type of scheduler if defined - noise_scheduler.register_to_config(prediction_type=args.prediction_type) - - if noise_scheduler.config.prediction_type == "epsilon": - target = noise - elif noise_scheduler.config.prediction_type == "v_prediction": - target = noise_scheduler.get_velocity(model_input, noise, timesteps) - elif noise_scheduler.config.prediction_type == "sample": - # We set the target to latents here, but the model_pred will return the noise sample prediction. - target = model_input - # We will have to subtract the noise residual from the prediction to get the target sample. - model_pred = model_pred - noise - else: - raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") - - mask = model_kwargs.get('attention_mask', None) - if torch.all(mask.bool()): - mask = None - if get_sequence_parallel_state(): - assert mask is None - b, c, _, _, _ = model_pred.shape - if mask is not None: - mask = mask.unsqueeze(1).repeat(1, c, 1, 1, 1).float() # b t h w -> b c t h w - mask = mask.reshape(b, -1) - if args.snr_gamma is None: - # model_pred: b c t h w, attention_mask: b t h w - loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") - loss = loss.reshape(b, -1) - if mask is not None: - loss = (loss * mask).sum() / mask.sum() # mean loss on unpad patches - else: - loss = loss.mean() - else: - # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556. - # Since we predict the noise instead of x_0, the original formulation is slightly changed. - # This is discussed in Section 4.2 of the same paper. - snr = compute_snr(noise_scheduler, timesteps) - mse_loss_weights = torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min( - dim=1 - )[0] - if noise_scheduler.config.prediction_type == "epsilon": - mse_loss_weights = mse_loss_weights / snr - elif noise_scheduler.config.prediction_type == "v_prediction": - mse_loss_weights = mse_loss_weights / (snr + 1) - loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") - loss = loss.reshape(b, -1) - mse_loss_weights = mse_loss_weights.reshape(b, 1) - if mask is not None: - loss = (loss * mask * mse_loss_weights).sum() / mask.sum() # mean loss on unpad patches - else: - loss = (loss * mse_loss_weights).mean() - - # Gather the losses across all processes for logging (if we use distributed training). - avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() - progress_info.train_loss += avg_loss.detach().item() / args.gradient_accumulation_steps - - # Backpropagate - accelerator.backward(loss) - if accelerator.sync_gradients: - params_to_clip = model.parameters() - accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - if accelerator.sync_gradients: - sync_gradients_info(loss) - - if accelerator.is_main_process: - - if progress_info.global_step % args.checkpointing_steps == 0: - - if args.enable_tracker: - log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, - weight_dtype, progress_info.global_step) - - if args.use_ema: - # Store the UNet parameters temporarily and load the EMA parameters to perform inference. - ema_model.store(model.parameters()) - ema_model.copy_to(model.parameters()) - log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator, - weight_dtype, progress_info.global_step, ema=True) - # Switch back to the original UNet parameters. - ema_model.restore(model.parameters()) - - if prof is not None: - prof.step() - - - return loss - - def train_one_step(step_, data_item_, prof_=None): - train_loss = 0.0 - x, attn_mask, input_ids, cond_mask = data_item_ - if args.group_frame or args.group_resolution: - if not args.group_frame: - each_latent_frame = torch.any(attn_mask.flatten(-2), dim=-1).int().sum(-1).tolist() - # logger.info(f'rank: {accelerator.process_index}, step {step_}, special batch has attention_mask ' - # f'each_latent_frame: {each_latent_frame}') - print(f'rank: {accelerator.process_index}, step {step_}, special batch has attention_mask ' - f'each_latent_frame: {each_latent_frame}') - assert not torch.any(torch.isnan(x)), 'torch.any(torch.isnan(x))' - x = x.to(accelerator.device, dtype=ae.vae.dtype) # B C T+num_images H W, 16 + 4 - - attn_mask = attn_mask.to(accelerator.device) # B T+num_images H W - input_ids = input_ids.to(accelerator.device) # B 1+num_images L - cond_mask = cond_mask.to(accelerator.device) # B 1+num_images L - # if accelerator.process_index == 0: - # logger.info(f'rank: {accelerator.process_index}, x: {x.shape}, attn_mask: {attn_mask.shape}') - - with torch.no_grad(): - B, N, L = input_ids.shape # B 1+num_images L - input_ids_ = input_ids.reshape(-1, L) - cond_mask_ = cond_mask.reshape(-1, L) - cond = text_enc(input_ids_, cond_mask_) # B 1+num_images L D - cond = cond.reshape(B, N, L, -1) - x = ae.encode(x) # B C T H W - current_step_frame = x.shape[2] - current_step_sp_state = get_sequence_parallel_state() - if args.sp_size != 1: # enable sp - if current_step_frame == 1: # but image do not need sp - set_sequence_parallel_state(False) - else: - set_sequence_parallel_state(True) - if get_sequence_parallel_state(): - x, cond, attn_mask, cond_mask, use_image_num = prepare_parallel_data(x, cond, attn_mask, cond_mask, - args.use_image_num) - for iter in range(args.train_batch_size * args.sp_size // args.train_sp_batch_size): - with accelerator.accumulate(model): - st_idx = iter * args.train_sp_batch_size - ed_idx = (iter + 1) * args.train_sp_batch_size - model_kwargs = dict(encoder_hidden_states=cond[st_idx: ed_idx], - attention_mask=attn_mask[st_idx: ed_idx], - encoder_attention_mask=cond_mask[st_idx: ed_idx], use_image_num=use_image_num) - run(x[st_idx: ed_idx], model_kwargs, prof_) - - else: - with accelerator.accumulate(model): - assert not torch.any(torch.isnan(x)), 'after vae' - x = x.to(weight_dtype) - model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask, - encoder_attention_mask=cond_mask, use_image_num=args.use_image_num) - run(x, model_kwargs, prof_) - - set_sequence_parallel_state(current_step_sp_state) # in case the next step use sp, which need broadcast(timesteps) - - if progress_info.global_step >= args.max_train_steps: - return True - - return False - - def train_all_epoch(prof_=None): - for epoch in range(first_epoch, args.num_train_epochs): - progress_info.train_loss = 0.0 - if progress_info.global_step >= args.max_train_steps: - return True - - for step, data_item in enumerate(train_dataloader): - if train_one_step(step, data_item, prof_): - break - - - train_all_epoch() - accelerator.wait_for_everyone() - accelerator.end_training() - if get_sequence_parallel_state(): - destroy_sequence_parallel_group() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - # dataset & dataloader - parser.add_argument("--dataset", type=str, required=True) - parser.add_argument("--data", type=str, required='') - parser.add_argument("--cache_vae_latent", action="store_true") - parser.add_argument("--train_fps", type=int, default=24) - parser.add_argument("--drop_short_ratio", type=float, default=1.0) - parser.add_argument("--speed_factor", type=float, default=1.0) - parser.add_argument("--num_frames", type=int, default=65) - parser.add_argument("--video_length_tolerance_range", type=int, default=2.0) - parser.add_argument("--max_height", type=int, default=320) - parser.add_argument("--max_width", type=int, default=240) - parser.add_argument("--use_image_num", type=int, default=0) - parser.add_argument("--text_max_length", type=int, default=512) - parser.add_argument('--cfg', type=float, default=0.1) - parser.add_argument("--dataloader_num_workers", type=int, default=10, help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.") - parser.add_argument("--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader.") - parser.add_argument("--group_frame", action="store_true") - parser.add_argument("--group_resolution", action="store_true") - - # text encoder & vae & diffusion model - parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="Latte-XL/122") - parser.add_argument('--enable_8bit_t5', action='store_true') - parser.add_argument('--tile_sample_stride', type=float, default=0.125) - parser.add_argument('--enable_tiling', action='store_true') - parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="xformers") - parser.add_argument('--use_rope', action='store_true') - parser.add_argument('--interpolation_scale_h', type=float, default=1.0) - parser.add_argument('--interpolation_scale_w', type=float, default=1.0) - parser.add_argument('--interpolation_scale_t', type=float, default=1.0) - parser.add_argument("--downsampler", type=str, default=None) - parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse") - parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse") - parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl') - parser.add_argument("--cache_dir", type=str, default='./cache_dir') - parser.add_argument("--pretrained", type=str, default=None) - parser.add_argument('--enable_stable_fp32', action='store_true') - parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.") - - # diffusion setting - parser.add_argument("--snr_gamma", type=float, default=None, help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. More details here: https://arxiv.org/abs/2303.09556.") - parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.") - parser.add_argument("--ema_decay", type=float, default=0.999) - parser.add_argument("--ema_start_step", type=int, default=0) - parser.add_argument("--noise_offset", type=float, default=0.02, help="The scale of noise offset.") - parser.add_argument("--prediction_type", type=str, default=None, help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.") - - # validation & logs - parser.add_argument("--validaiton_prompt", nargs='+', help="List of prompts to use for validation.") - parser.add_argument("--validation_sampling_steps", type=int, default=64) - parser.add_argument('--guidance_scale', type=float, default=4.5) - parser.add_argument("--enable_tracker", action="store_true") - parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") - parser.add_argument("--output_dir", type=str, default=None, help="The output directory where the model predictions and checkpoints will be written.") - parser.add_argument("--checkpoints_total_limit", type=int, default=None, help=("Max number of checkpoints to store.")) - parser.add_argument("--checkpointing_steps", type=int, default=500, - help=( - "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" - " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" - " training using `--resume_from_checkpoint`." - ), - ) - parser.add_argument("--resume_from_checkpoint", type=str, default=None, - help=( - "Whether training should be resumed from a previous checkpoint. Use a path saved by" - ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' - ), - ) - parser.add_argument("--logging_dir", type=str, default="logs", - help=( - "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" - " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." - ), - ) - 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.' - ), - ) - # optimizer & scheduler - parser.add_argument("--num_train_epochs", type=int, default=100) - parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform. If provided, overrides num_train_epochs.") - parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.") - parser.add_argument("--optimizer", type=str, default="adamW", help='The optimizer type to use. Choose between ["AdamW", "prodigy"]') - parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.") - parser.add_argument("--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.") - parser.add_argument("--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler.") - parser.add_argument("--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW") - parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers.") - parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers.") - parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") - parser.add_argument("--adam_weight_decay", type=float, default=1e-02, help="Weight decay to use for unet params") - parser.add_argument("--adam_weight_decay_text_encoder", type=float, default=None, help="Weight decay to use for text_encoder") - parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer and Prodigy optimizers.") - parser.add_argument("--prodigy_use_bias_correction", type=bool, default=True, help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW") - parser.add_argument("--prodigy_safeguard_warmup", type=bool, default=True, help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. Ignored if optimizer is adamW") - parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") - parser.add_argument("--prodigy_beta3", type=float, default=None, - help="coefficients for computing the Prodidy stepsize using running averages. If set to None, " - "uses the value of square root of beta2. Ignored if optimizer is adamW", - ) - parser.add_argument("--lr_scheduler", type=str, default="constant", - help=( - 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' - ' "constant", "constant_with_warmup"]' - ), - ) - parser.add_argument("--allow_tf32", action="store_true", - help=( - "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" - " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" - ), - ) - parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"], - help=( - "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" - " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" - " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." - ), - ) - - parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") - parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel") - parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training") - - args = parser.parse_args() - main(args) diff --git a/inference_mochi.sh b/inference_mochi.sh index 591c387ac..83d7e1ec3 100644 --- a/inference_mochi.sh +++ b/inference_mochi.sh @@ -1,24 +1,87 @@ -#!/bin/bash +num_gpus=2 +prompts=( + "In this animated scene, Tom the cat is peeking out from behind a tall stack of books in a cozy library. His eyes are wide with curiosity as he spots Jerry, the tiny brown mouse, sitting on a book cover at the top of the pile. Jerry is holding a small, crumpled piece of cheese in his paws, looking down at Tom with a mischievous grin. The style of the animation brings warmth to the scene with rich, textured colors that highlight the humorous interaction between the characters." + "In this playful scene, Tom, the gray-blue cat, is dressed as a chef, complete with a white hat and apron. He’s standing over a steaming pot on a stove, stirring with a large spoon. Meanwhile, Jerry, the brown mouse, is perched on the edge of the counter, tossing in tiny sprigs of herbs. Their expressions suggest a lighthearted teamwork moment. The animation style emphasizes vibrant colors and exaggerated movements that enhance the lively interaction between the characters." +) + +for prompt in "${prompts[@]}"; do + torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ + fastvideo/sample/sample_t2v_mochi.py \ + --model_path data/mochi \ + --prompts "$prompt" \ + --num_frames 79 \ + --height 480 \ + --width 848 \ + --num_inference_steps 64 \ + --guidance_scale 4.5 \ + --output_path "outputs_video/T_J_original" \ + --seed 12345 +done + + + + + + + + +num_gpus=2 +prompts=( + "In this animated scene, Tom the cat is peeking out from behind a tall stack of books in a cozy library. His eyes are wide with curiosity as he spots Jerry, the tiny brown mouse, sitting on a book cover at the top of the pile. Jerry is holding a small, crumpled piece of cheese in his paws, looking down at Tom with a mischievous grin. The style of the animation brings warmth to the scene with rich, textured colors that highlight the humorous interaction between the characters." + "In this playful scene, Tom, the gray-blue cat, is dressed as a chef, complete with a white hat and apron. He’s standing over a steaming pot on a stove, stirring with a large spoon. Meanwhile, Jerry, the brown mouse, is perched on the edge of the counter, tossing in tiny sprigs of herbs. Their expressions suggest a lighthearted teamwork moment. The animation style emphasizes vibrant colors and exaggerated movements that enhance the lively interaction between the characters." +) + +for prompt in "${prompts[@]}"; do + torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ + fastvideo/sample/sample_t2v_mochi.py \ + --model_path data/mochi \ + --prompts "$prompt" \ + --num_frames 79 \ + --height 480 \ + --width 848 \ + --num_inference_steps 64 \ + --guidance_scale 4.5 \ + --output_path "outputs_video/T_J_FT_normal_79_1000" \ + --transformer_path data/outputs/T_J_FT/checkpoint-1000 \ + --seed 12345 +done + + + + + + + +black myth: a character with a pig-like face, dressed in a green robe and holding a string of black beads. the character is riding a white horse outdoors, with a mountainous landscape visible in the background. The horse is a strong, muscular white steed with a slightly wild, untamed mane flowing in the wind, suggesting speed and vigor. Its eyes are alert, with a slight intensity that complements the adventurous aura of the character. +black myth: a character dressed in elaborate, ornate armor that is richly decorated with intricate designs and patterns. the armor appears to be of a high-quality material, possibly metal, and is adorned with various embellishments such as jewels and metallic accents. the character is riding a white horse outdoors, with a mountainous landscape visible in the background. the lighting suggests it might be either dawn or dusk, casting a soft glow on the scene. the character's pose is dynamic, suggesting movement or action. + + + num_gpus=4 - -torchrun --nnodes=1 --nproc_per_node=4 --master_port 29503 \ +torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ fastvideo/sample/sample_t2v_mochi.py \ --model_path data/mochi \ - --prompts "A cute fluffy panda eating Chinese food in a restaurant" \ - --num_frames 163 \ + --prompts "A hyper-realistic sci-fi scene featuring a brave monkey astronaut, floating in zero gravity aboard a futuristic spacecraft. The monkey wears a sleek, advanced spacesuit with glowing blue accents, gazing through the spacecraft window at a mesmerizing nebula and distant planets. The spacecraft interior is filled with detailed control panels and holographic displays, creating an immersive, high-tech atmosphere." \ + --num_frames 91 \ --height 480 \ --width 848 \ --num_inference_steps 64 \ --guidance_scale 4.5 \ - --output_path outputs_video/debug \ + --output_path outputs_video/\ --seed 12345 -num_gpus=4 -torchrun --nnodes=1 --nproc_per_node=4 --master_port 29503 \ + + + + + + +num_gpus=2 +torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ fastvideo/sample/sample_t2v_mochi.py \ --model_path data/mochi/ \ --prompt_embed_path data/Encoder_Overfit_Data/prompt_embed/0.pt \ diff --git a/pyproject.toml b/pyproject.toml index c270a8db4..b1e6fa7b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,15 +13,15 @@ classifiers = [ "License :: OSI Approved :: Apache Software License", ] dependencies = [ -"transformers==4.46.1", "accelerate==1.0.1", "tokenizers==0.20.1", "diffusers==0.31.0", "albumentations==1.4.20", "av==13.1.0", +"transformers==4.46.1", "accelerate==1.0.1", "tokenizers==0.20.1", "albumentations==1.4.20", "av==13.1.0", "decord==0.6.0", "einops==0.8.0", "fastapi==0.115.3", "gdown==5.2.0", "h5py==3.12.1", "idna==3.6", "imageio==2.36.0", "matplotlib==3.9.2", "numpy==1.26.3", "omegaconf==2.3.0", "opencv-python==4.10.0.84", "opencv-python-headless==4.10.0.84", "pandas==2.2.3", "pillow==10.2.0", "pydub==0.25.1", "pytorch-lightning==2.4.0", "pytorchvideo==0.1.5", "PyYAML==6.0.1", "regex==2024.9.11", "requests==2.31.0", "scikit-learn==1.5.2", "scipy==1.14.1", "six==1.16.0", "test-tube==0.7.5", "timm==1.0.11", "torchdiffeq==0.2.4", "torchmetrics==1.5.1", "tqdm==4.66.5", "urllib3==2.2.0", "uvicorn==0.32.0", "scikit-video==1.1.11", "imageio-ffmpeg==0.5.1", "sentencepiece==0.2.0", "beautifulsoup4==4.12.3", "ftfy==6.3.0", -"moviepy==1.0.3", "wandb==0.18.5", "tensorboard==2.18.0", "pydantic==2.9.2", "gradio==5.3.0", "huggingface_hub==0.26.1", "protobuf==5.28.3", "peft==0.13.2" -"watch", "gpustat" +"moviepy==1.0.3", "wandb==0.18.5", "tensorboard==2.18.0", "pydantic==2.9.2", "gradio==5.3.0", "huggingface_hub==0.26.1", "protobuf==5.28.3", +"watch", "gpustat", "peft==0.13.2", "liger_kernel==0.4.1" ] [project.optional-dependencies] diff --git a/scripts/data_preprocess/finetune_data_gen.sh b/scripts/data_preprocess/finetune_data_gen.sh index eef958ae5..05777d457 100644 --- a/scripts/data_preprocess/finetune_data_gen.sh +++ b/scripts/data_preprocess/finetune_data_gen.sh @@ -4,7 +4,7 @@ MODEL_PATH="/ephemeral/hao.zhang/outputfolder/ckptfolder/mochi_diffuser" DATA_MERGE_PATH="/ephemeral/hao.zhang/resourcefolder/Mochi-Synthetic-Data-BW-Finetune/merge.txt" OUTPUT_DIR="./data/BW-Finetune-Synthetic-Data_test" -torchrun --nproc_per_node=$GPU_NUM \ +rchrun --nproc_per_node=$GPU_NUM \ ./fastvideo/utils/data_preprocess/finetune_data_VAE.py \ --model_path $MODEL_PATH \ --data_merge_path $DATA_MERGE_PATH \ @@ -15,6 +15,10 @@ torchrun --nproc_per_node=$GPU_NUM \ --dataloader_num_workers 1 \ --output_dir=$OUTPUT_DIR +to + + + torchrun --nproc_per_node=$GPU_NUM \ ./fastvideo/utils/data_preprocess/finetune_data_T5.py \ --model_path $MODEL_PATH \ diff --git a/scripts/finetune_T_J.sh b/scripts/finetune_T_J.sh new file mode 100644 index 000000000..0f37472c3 --- /dev/null +++ b/scripts/finetune_T_J.sh @@ -0,0 +1,117 @@ +torchrun --nnodes 4 --nproc_per_node 4 \ + --node_rank=3 \ + --rdzv_id=456 \ + --rdzv_backend=c10d \ + --rdzv_endpoint=172.23.30.15:29500 \ + fastvideo/train.py \ + --seed 42 \ + --pretrained_model_name_or_path data/mochi \ + --cache_dir "data/.cache" \ + --data_json_path "data/BLACK-MYTH-Finetune-Dataset/videos2caption.json" \ + --validation_prompt_dir "data/BLACK-MYTH-Finetune-Dataset/validation_prompt_embed_mask" \ + --uncond_prompt_dir "data/BLACK-MYTH-Finetune-Dataset/uncond_prompt_embed_mask" \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --num_latent_t 16 \ + --sp_size 4 \ + --train_sp_batch_size 4 \ + --dataloader_num_workers 4 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=4000 \ + --learning_rate=2e-5 \ + --mixed_precision="bf16" \ + --checkpointing_steps=500 \ + --validation_steps=100 \ + --validation_sampling_steps 64 \ + --checkpoints_total_limit 3 \ + --allow_tf32 \ + --ema_start_step 0 \ + --cfg 0.0 \ + --ema_decay 0.999 \ + --log_validation \ + --output_dir="data/outputs/black_myth_correct_video_mask" \ + --weighting_scheme "uniform" \ + --num_frames 91 \ + --selective_checkpointing 1.0 + + +torchrun --nnodes 4 --nproc_per_node 4 \ + --node_rank=3 \ + --rdzv_id=456 \ + --rdzv_backend=c10d \ + --rdzv_endpoint=172.23.30.15:29500 \ + fastvideo/train.py \ + --seed 42 \ + --pretrained_model_name_or_path data/mochi \ + --cache_dir "data/.cache" \ + --data_json_path "data/BLACK-MYTH-DREAM-BOOTH/videos2caption.json" \ + --validation_prompt_dir "data/BLACK-MYTH-DREAM-BOOTH/validation_prompt_embed_mask" \ + --uncond_prompt_dir "data/BLACK-MYTH-DREAM-BOOTH/uncond_prompt_embed_mask" \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --num_latent_t 16 \ + --sp_size 4 \ + --train_sp_batch_size 4 \ + --dataloader_num_workers 4 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=4000 \ + --learning_rate=2e-5 \ + --mixed_precision="bf16" \ + --checkpointing_steps=100 \ + --validation_steps=50 \ + --validation_sampling_steps 64 \ + --checkpoints_total_limit 3 \ + --allow_tf32 \ + --ema_start_step 0 \ + --cfg 0.1 \ + --ema_decay 0.999 \ + --log_validation \ + --output_dir="data/outputs/black_myth_correct_video_mask" \ + --weighting_scheme "uniform" \ + --num_frames 91 \ + --selective_checkpointing 1.0 + + + + +torchrun --nnodes 1 --nproc_per_node 4 \ + fastvideo/train.py \ + --seed 42 \ + --pretrained_model_name_or_path data/mochi \ + --cache_dir "data/.cache" \ + --data_json_path "data/BLACK-MYTH-Finetune-Dataset/videos2caption.json" \ + --validation_prompt_dir "data/BLACK-MYTH-Finetune-Dataset/validation_prompt_embed_mask" \ + --uncond_prompt_dir "data/BLACK-MYTH-Finetune-Dataset/uncond_prompt_embed_mask" \ + --gradient_checkpointing \ + --train_batch_size=1 \ + --num_latent_t 16 \ + --sp_size 4 \ + --train_sp_batch_size 4 \ + --dataloader_num_workers 4 \ + --gradient_accumulation_steps=1 \ + --max_train_steps=2000 \ + --learning_rate=2e-5 \ + --mixed_precision="bf16" \ + --checkpointing_steps=100 \ + --validation_steps=100 \ + --validation_sampling_steps 64 \ + --checkpoints_total_limit 3 \ + --allow_tf32 \ + --ema_start_step 0 \ + --cfg 0.0 \ + --ema_decay 0.999 \ + --log_validation \ + --output_dir="data/outputs/black_myth_dream" \ + --weighting_scheme "uniform" \ + --num_frames 91 \ + + +| 4/2000 [01:23<10:59:33, 19.83s/it, loss=0.74, grad_norm=2.6] + + +| 2/2000 [00:44<11:59:27, 21.61s/it, loss=0.873, grad_norm=1.57] +| 3/2000 [01:03<11:18:51, 20.40s/it, loss=0.813, grad_norm=1.48] + + +# FSDP Hybrid: + diff --git a/t2v_debug_multi.sh b/t2v_debug_multi.sh index 1c9045036..39a88377b 100644 --- a/t2v_debug_multi.sh +++ b/t2v_debug_multi.sh @@ -1,3 +1,4 @@ +export WANDB_MODE=online torchrun --nnodes 1 --nproc_per_node 4 \ fastvideo/train.py \ --seed 42 \ @@ -5,72 +6,26 @@ torchrun --nnodes 1 --nproc_per_node 4 \ --cache_dir "data/.cache" \ --data_json_path "data/Encoder_Overfit_Data/videos2caption.json" \ --validation_prompt_dir "data/Encoder_Overfit_Data/validation_prompt_embed_mask" \ - --uncond_prompt_dir "data/Encoder_Overfit_Data/uncond_prompt_embed_mask" \ + --uncond_prompt_dir "data/Encoder_Overfit_Dataa/uncond_prompt_embed_mask" \ --gradient_checkpointing \ --train_batch_size=1 \ - --num_latent_t 2 \ - --sp_size 2 \ - --train_sp_batch_size 1 \ - --dataloader_num_workers 1 \ - --gradient_accumulation_steps 4 \ - --max_train_steps 200 \ - --learning_rate 1e-4 \ - --mixed_precision="bf16" \ - --checkpointing_steps 6 \ - --validation_steps 40 \ - --validation_sampling_steps 2 \ - --checkpoints_total_limit 3 \ - --allow_tf32 \ - --ema_start_step 0 \ - --cfg 0.1 \ - --ema_decay 0.999 \ - --log_validation \ - --output_dir="data/outputs/BW_Testrun_Test" \ - --use_lora \ - --resume_from_lora_checkpoint data/outputs/BW_Testrun/lora-checkpoint-5 \ - -torchrun --nnodes 1 --nproc_per_node 4 \ - fastvideo/train.py \ - --seed 42 \ - --pretrained_model_name_or_path data/mochi \ - --cache_dir data/.cache \ - --data_json_path data/Encoder_Overfit_Data/videos2caption.json \ - --validation_prompt_dir data/Encoder_Overfit_Data/validation_prompt_embed_mask \ - --uncond_prompt_dir data/Encoder_Overfit_Data/uncond_prompt_embed_mask \ - --gradient_checkpointing \ - --train_batch_size 1 \ - --num_latent_t 12 \ + --num_latent_t 14 \ --sp_size 2 \ --train_sp_batch_size 2 \ - --dataloader_num_workers 1 \ - --gradient_accumulation_steps 1 \ - --max_train_steps 400 \ - --learning_rate 1e-4 \ - --mixed_precision bf16 \ - --checkpointing_steps 200 \ - --validation_steps 50 \ + --dataloader_num_workers 4 \ + --gradient_accumulation_steps=2 \ + --max_train_steps=2000 \ + --learning_rate=1e-5 \ + --mixed_precision="bf16" \ + --checkpointing_steps=250 \ + --validation_steps 250 \ --validation_sampling_steps 64 \ --checkpoints_total_limit 3 \ --allow_tf32 \ --ema_start_step 0 \ - --cfg 0.1 \ + --cfg 0.0 \ --ema_decay 0.999 \ --log_validation \ - --output_dir data/outputs/BW_Testrun_Test/ \ - --lora_rank 128 \ - --lora_alpha 256 \ - --use_lora \ - --use_cpu_offload \ - --resume_from_lora_checkpoint data/outputs/BW_Testrun_2/lora-checkpoint-100 \ - - - - -# SP single node: 13 -# SP multi node: 20 -# SP multi node with larger batch size: 15 -# No SP Single node: 10.5 -# No SP Multe node: 11 - - + --output_dir="data/outputs/T_J_FT" \ + --weighting_scheme "uniform"