From 85639d110151d98238efaaa3fdf19192e15c3f2d Mon Sep 17 00:00:00 2001 From: Zhang Peiyuan Date: Fri, 13 Dec 2024 11:52:57 -0800 Subject: [PATCH] [feat] add hunyuan adv (#79) --- fastvideo/distill.py | 2 +- fastvideo/distill/discriminator.py | 9 +- fastvideo/distill_adv.py | 134 +++++++++++--------- fastvideo/models/hunyuan/modules/models.py | 17 ++- fastvideo/models/mochi_hf/modeling_mochi.py | 12 +- pyproject.toml | 7 +- scripts/distill_hunyuan.sh | 12 +- 7 files changed, 105 insertions(+), 88 deletions(-) diff --git a/fastvideo/distill.py b/fastvideo/distill.py index 727aaabdb..64d46b5a1 100644 --- a/fastvideo/distill.py +++ b/fastvideo/distill.py @@ -30,7 +30,7 @@ from fastvideo.utils.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpoint from diffusers import ( FlowMatchEulerDiscreteScheduler, ) -from fastvideo.utils.load import get_no_split_modules, load_transformer +from fastvideo.utils.load import load_transformer from fastvideo.distill.solver import EulerSolver, extract_into_tensor from copy import deepcopy from diffusers.optimization import get_scheduler diff --git a/fastvideo/distill/discriminator.py b/fastvideo/distill/discriminator.py index d586bfb70..96f4925b8 100644 --- a/fastvideo/distill/discriminator.py +++ b/fastvideo/distill/discriminator.py @@ -62,9 +62,10 @@ class Discriminator(nn.Module): stride=8, num_h_per_head=1, adapter_channel_dims=[3072], + total_layers = 48, ): super().__init__() - adapter_channel_dims = adapter_channel_dims * (48 // stride) + adapter_channel_dims = adapter_channel_dims * (total_layers // stride) self.stride = stride self.num_h_per_head = num_h_per_head self.head_num = len(adapter_channel_dims) @@ -89,9 +90,9 @@ class Discriminator(nn.Module): return custom_forward - assert len(features) // self.stride == len(self.heads) - for i in range(0, len(features), self.stride): - for h in self.heads[i // self.stride]: + assert len(features) == len(self.heads) + for i in range(0, len(features)): + for h in self.heads[i]: # out = torch.utils.checkpoint.checkpoint( # create_custom_forward(h), # features[i], diff --git a/fastvideo/distill_adv.py b/fastvideo/distill_adv.py index dfcf092dc..f0112a82c 100644 --- a/fastvideo/distill_adv.py +++ b/fastvideo/distill_adv.py @@ -22,6 +22,7 @@ from torch.distributed.fsdp import ( StateDictType, FullStateDictConfig, ) +from fastvideo.utils.load import load_transformer from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule import json @@ -76,6 +77,7 @@ def gan_d_loss( encoder_hidden_states, encoder_attention_mask, weight, + discriminator_head_stride ): loss = 0.0 # collate sample_fake and sample_real @@ -85,7 +87,8 @@ def gan_d_loss( encoder_hidden_states, timestep, encoder_attention_mask, - output_attn=True, + output_features=True, + output_features_stride=discriminator_head_stride, return_dict=False, )[1] real_features = teacher_transformer( @@ -93,7 +96,8 @@ def gan_d_loss( encoder_hidden_states, timestep, encoder_attention_mask, - output_attn=True, + output_features=True, + output_features_stride=discriminator_head_stride, return_dict=False, )[1] @@ -115,6 +119,7 @@ def gan_g_loss( encoder_hidden_states, encoder_attention_mask, weight, + discriminator_head_stride ): loss = 0.0 features = teacher_transformer( @@ -122,7 +127,8 @@ def gan_g_loss( encoder_hidden_states, timestep, encoder_attention_mask, - output_attn=True, + output_features=True, + output_features_stride=discriminator_head_stride, return_dict=False, )[1] fake_outputs = discriminator( @@ -135,20 +141,19 @@ def gan_g_loss( return loss -def train_one_step_mochi( +def distill_one_step_adv( transformer, + model_type, teacher_transformer, optimizer, discriminator, discriminator_optimizer, - global_step, lr_scheduler, loader, noise_scheduler, solver, noise_random_generator, sp_size, - precondition_outputs, max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, @@ -157,6 +162,7 @@ def train_one_step_mochi( not_apply_cfg_solver, distill_cfg, adv_weight, + discriminator_head_stride ): optimizer.zero_grad() discriminator_optimizer.zero_grad() @@ -167,7 +173,7 @@ def train_one_step_mochi( latents_attention_mask, encoder_attention_mask, ) = next(loader) - model_input = normalize_mochi_dit_input(latents) + model_input = normalize_dit_input(model_type, latents) noise = torch.randn_like(model_input) bsz = model_input.shape[0] index = torch.randint( @@ -284,6 +290,7 @@ def train_one_step_mochi( encoder_hidden_states.float(), encoder_attention_mask, 1.0, + discriminator_head_stride ) g_loss += g_gan_loss g_loss.backward() @@ -308,6 +315,7 @@ def train_one_step_mochi( encoder_hidden_states, encoder_attention_mask, 1.0, + discriminator_head_stride, ) d_loss.backward() @@ -347,21 +355,14 @@ def main(args): main_print(f"--> loading model from {args.pretrained_model_name_or_path}") # keep the master weight to float32 - if args.dit_model_name_or_path: - transformer = transformer = MochiTransformer3DModel.from_pretrained( - args.dit_model_name_or_path, - torch_dtype=torch.float32, - # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, - ) - else: - transformer = MochiTransformer3DModel.from_pretrained( - args.pretrained_model_name_or_path, - subfolder="transformer", - torch_dtype=torch.float32, - # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, - ) + transformer = load_transformer( + args.model_type, + args.dit_model_name_or_path, + args.pretrained_model_name_or_path, + torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16, + ) teacher_transformer = deepcopy(transformer) - discriminator = Discriminator(args.discriminator_head_stride) + discriminator = Discriminator(args.discriminator_head_stride, total_layers = 48 if args.model_type =="mochi" else 40) if args.use_lora: transformer.requires_grad_(False) @@ -383,15 +384,20 @@ def main(args): 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 + fsdp_kwargs, no_split_modules = get_dit_fsdp_kwargs( + transformer, + args.fsdp_sharding_startegy, + args.use_lora, + args.use_cpu_offload, + args.master_weight_type, ) discriminator_fsdp_kwargs = get_discriminator_fsdp_kwargs(args.master_weight_type) if args.use_lora: + assert args.model_type == "mochi", "LoRA is only supported for Mochi model." transformer.config.lora_rank = args.lora_rank transformer.config.lora_alpha = args.lora_alpha transformer.config.lora_target_modules = ["to_k", "to_q", "to_v", "to_out.0"] - transformer._no_split_modules = ["MochiTransformerBlock"] + transformer._no_split_modules = no_split_modules fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer) transformer = FSDP( @@ -409,8 +415,12 @@ def main(args): main_print(f"--> model loaded") if args.gradient_checkpointing: - apply_fsdp_checkpointing(transformer, args.selective_checkpointing) - apply_fsdp_checkpointing(teacher_transformer, args.selective_checkpointing) + apply_fsdp_checkpointing( + transformer, no_split_modules, args.selective_checkpointing + ) + apply_fsdp_checkpointing( + teacher_transformer, no_split_modules, args.selective_checkpointing + ) # Set model as trainable. transformer.train() teacher_transformer.requires_grad_(False) @@ -562,39 +572,49 @@ def main(args): step_times = deque(maxlen=100) # log_validation(args, transformer, device, - # torch.bfloat16, init_steps, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold, ema=False) - + # torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False) + def get_num_phases(multi_phased_distill_schedule, step): + # step-phase,step-phase + multi_phases = multi_phased_distill_schedule.split(",") + phase = multi_phases[-1].split("-")[-1] + for step_phases in multi_phases: + phase_step, phase = step_phases.split("-") + if step <= int(phase_step): + return int(phase) + return phase for i in range(init_steps): _ = next(loader) for step in range(init_steps + 1, args.max_train_steps + 1): + assert args.multi_phased_distill_schedule is not None + num_phases = get_num_phases(args.multi_phased_distill_schedule, step) start_time = time.time() ( generator_loss, generator_grad_norm, discriminator_loss, discriminator_grad_norm, - ) = train_one_step_mochi( + ) = distill_one_step_adv( transformer, + args.model_type, teacher_transformer, optimizer, discriminator, discriminator_optimizer, - step, lr_scheduler, loader, noise_scheduler, solver, noise_random_generator, args.sp_size, - args.precondition_outputs, args.max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, args.num_euler_timesteps, - args.validation_sampling_steps, + num_phases, args.not_apply_cfg_solver, args.distill_cfg, args.adv_weight, + args.discriminator_head_stride ) step_time = time.time() - start_time @@ -633,15 +653,17 @@ def main(args): ) else: # Your existing checkpoint saving code - save_checkpoint_generator_discriminator( - transformer, - optimizer, - discriminator, - discriminator_optimizer, - rank, - args.output_dir, - step, - ) + # TODO + # save_checkpoint_generator_discriminator( + # transformer, + # optimizer, + # discriminator, + # discriminator_optimizer, + # rank, + # args.output_dir, + # step, + # ) + save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps) main_print(f"--> checkpoint saved at step {step}") dist.barrier() if args.log_validation and step % args.validation_steps == 0: @@ -655,25 +677,17 @@ def main(args): shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold, + linear_range=args.linear_range, ema=False, ) + if args.use_lora: save_lora_checkpoint( transformer, optimizer, rank, args.output_dir, args.max_train_steps ) else: - save_checkpoint( - transformer, optimizer, rank, args.output_dir, args.max_train_steps - ) - save_checkpoint( - discriminator, - discriminator_optimizer, - rank, - args.output_dir, - step, - discriminator=True, - ) + save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps) if get_sequence_parallel_state(): destroy_sequence_parallel_group() @@ -682,6 +696,9 @@ def main(args): if __name__ == "__main__": parser = argparse.ArgumentParser() + parser.add_argument( + "--model_type", type=str, default="mochi", help="The type of model to train." + ) # dataset & dataloader parser.add_argument("--data_json_path", type=str, required=True) parser.add_argument("--num_frames", type=int, default=163) @@ -712,16 +729,9 @@ if __name__ == "__main__": 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( - "--precondition_outputs", - action="store_true", - help="Whether to precondition the outputs of the model.", - ) - # validation & logs - parser.add_argument("--validation_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_sampling_steps", type=str, default="64") + parser.add_argument("--validation_guidance_scale", type=str, default="4.5") parser.add_argument("--validation_steps", type=float, default=64) parser.add_argument("--log_validation", action="store_true") parser.add_argument("--tracker_project_name", type=str, default=None) @@ -750,6 +760,7 @@ if __name__ == "__main__": " training using `--resume_from_checkpoint`." ), ) + parser.add_argument("--validation_prompt_dir", type=str) parser.add_argument("--shift", type=float, default=1.0) parser.add_argument( "--resume_from_checkpoint", @@ -866,6 +877,7 @@ if __name__ == "__main__": "--lora_rank", type=int, default=128, help="LoRA rank parameter. " ) parser.add_argument("--fsdp_sharding_startegy", default="full") + parser.add_argument("--multi_phased_distill_schedule", type=str, default=None) parser.add_argument( "--gradient_accumulation_steps", type=int, diff --git a/fastvideo/models/hunyuan/modules/models.py b/fastvideo/models/hunyuan/modules/models.py index d056aadcf..9a0cbe7a9 100644 --- a/fastvideo/models/hunyuan/modules/models.py +++ b/fastvideo/models/hunyuan/modules/models.py @@ -601,7 +601,8 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): encoder_hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_attention_mask: torch.Tensor, - output_attn=False, + output_features=False, + output_features_stride = 8, attention_kwargs: Optional[Dict[str, Any]] = None, return_dict: bool = False, guidance=None, @@ -663,6 +664,8 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): # Merge txt and img to pass through single stream blocks. x = torch.cat((img, txt), 1) + if output_features: + features_list = [] if len(self.single_blocks) > 0: for _, block in enumerate(self.single_blocks): single_block_args = [ @@ -674,6 +677,8 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): ] x = block(*single_block_args) + if output_features and _ % output_features_stride == 0: + features_list.append(x[:, :img_seq_len, ...]) img = x[:, :img_seq_len, ...] @@ -681,10 +686,12 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) img = self.unpatchify(img, tt, th, tw) - if return_dict: - out["x"] = img - return out - return (img,) + assert return_dict == False, "return_dict is not supported." + if output_features: + features_list = torch.stack(features_list, dim=0) + else: + features_list = None + return (img, features_list) def unpatchify(self, x, t, h, w): """ diff --git a/fastvideo/models/mochi_hf/modeling_mochi.py b/fastvideo/models/mochi_hf/modeling_mochi.py index 3876577f1..3a2f7f2a6 100644 --- a/fastvideo/models/mochi_hf/modeling_mochi.py +++ b/fastvideo/models/mochi_hf/modeling_mochi.py @@ -621,7 +621,8 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): encoder_hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_attention_mask: torch.Tensor, - output_attn=False, + output_features=False, + output_features_stride = 8, attention_kwargs: Optional[Dict[str, Any]] = None, return_dict: bool = False, ) -> torch.Tensor: @@ -697,7 +698,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): encoder_attention_mask, temb, image_rotary_emb, - output_attn, + output_features, **ckpt_kwargs, ) else: @@ -707,9 +708,10 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): encoder_attention_mask=encoder_attention_mask, temb=temb, image_rotary_emb=image_rotary_emb, - output_attn=output_attn, + output_attn=output_features, ) - attn_outputs_list.append(attn_outputs) + if i % output_features_stride == 0: + attn_outputs_list.append(attn_outputs) hidden_states = self.norm_out(hidden_states, temb) hidden_states = self.proj_out(hidden_states) @@ -724,7 +726,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): # remove `lora_scale` from each PEFT layer unscale_lora_layers(self, lora_scale) - if not output_attn: + if not output_features: attn_outputs_list = None else: attn_outputs_list = torch.stack(attn_outputs_list, dim=0) diff --git a/pyproject.toml b/pyproject.toml index 675aadba9..99c679b6f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,14 +21,9 @@ dependencies = [ "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", -"watch", "gpustat", "peft==0.13.2", "liger_kernel==0.4.1", "einops==0.8.0", "wheel==0.44.0"] +"watch", "gpustat", "peft==0.13.2", "liger_kernel==0.4.1", "einops==0.8.0", "wheel==0.44.0", "loguru"] -[project.optional-dependencies] -hunyuan = [ - "loguru" -] - [tool.setuptools.packages.find] exclude = ["assets*", "docker*", "docs", "scripts*"] diff --git a/scripts/distill_hunyuan.sh b/scripts/distill_hunyuan.sh index 6c90152e3..9b8516550 100644 --- a/scripts/distill_hunyuan.sh +++ b/scripts/distill_hunyuan.sh @@ -3,19 +3,19 @@ export WANDB_MODE=online export WANDB_API_KEY=4f6de3765d6464f43e0506ec7d785641af645e73 -torchrun --nnodes 1 --nproc_per_node 8\ - fastvideo/distill.py\ +torchrun --nnodes 1 --nproc_per_node 4\ + fastvideo/distill_adv.py\ --seed 42\ --pretrained_model_name_or_path data/hunyuan\ --dit_model_name_or_path data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt\ --model_type "hunyuan" \ --cache_dir "data/.cache"\ - --data_json_path "data/Hunyuan-Mixkit-Data/videos2caption.json"\ - --validation_prompt_dir "data/Hunyuan-Mixkit-Data/validation"\ + --data_json_path "data/Hunyuan-Distill-Data/videos2caption.json"\ + --validation_prompt_dir "data/Hunyuan-Distill-Data/validation"\ --gradient_checkpointing\ --train_batch_size=1\ - --num_latent_t 24\ - --sp_size 4\ + --num_latent_t 1\ + --sp_size 1\ --train_sp_batch_size 1\ --dataloader_num_workers 4\ --gradient_accumulation_steps=1\