From 3d6cac57fc70b148da5485c48d02cf80e88feca1 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Tue, 9 Sep 2025 01:13:20 +0000 Subject: [PATCH] Fix runtime errors during training --- .../SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh | 139 ++++++++++++++++++ fastvideo/models/dits/causal_wanvideo.py | 20 +-- .../scheduling_flow_match_euler_discrete.py | 31 +++- fastvideo/models/utils.py | 27 +++- fastvideo/training/distillation_pipeline.py | 18 ++- .../self_forcing_distillation_pipeline.py | 23 +-- 6 files changed, 219 insertions(+), 39 deletions(-) create mode 100644 examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh diff --git a/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh b/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh new file mode 100644 index 000000000..c81181493 --- /dev/null +++ b/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh @@ -0,0 +1,139 @@ +# Basic Info +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +# different cache dir for different processes +export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=offline +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA +export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' + +# Configs +NUM_GPUS=8 + +# Model paths for DMD distillation: +GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Teacher model +FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model + +DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" +VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/SFWan2.1-I2V/validation_better.json" +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name SFwan_t2v_distill_self_forcing_dmd + --output_dir "/mnt/sharefs/users/hao.zhang/wl/sf_checkpoints/ode0_SFwan_t2v_finetune_sf_${lr}_c${critic_lr}" + --wandb_run_name "DEBUG${lr}_c${critic_lr}" + --max_train_steps 4000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 21 + --num_height 480 + --num_width 832 + --num_frames 81 + --warp_denoising_step + --enable_gradient_checkpointing_type "full" + --log_visualization + --simulate_generator_forward + --num_frame_per_block 3 + --enable_gradient_masking + --gradient_mask_last_n_frames 21 +) + +# Parallel arguments +parallel_args=( + --num_gpus 8 # 64 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 # 64 + --hsdp_shard_dim 8 +) + +# Model arguments +model_args=( + --model_path $GENERATOR_MODEL_PATH # TODO: check if you can remove this in this script + --pretrained_model_name_or_path $GENERATOR_MODEL_PATH + --generator_model_path $GENERATOR_MODEL_PATH + --real_score_model_path $REAL_SCORE_MODEL_PATH + --fake_score_model_path $FAKE_SCORE_MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 4 +) + +# Validation arguments +validation_args=( +# --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 50 + --validation_sampling_steps "3" + --validation_guidance_scale "6.0" # not used for dmd inference +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-5 + --mixed_precision "bf16" + --training_state_checkpointing_steps 50 + --weight_only_checkpointing_steps 50 + --weight_decay 0.01 + --betas '0.0,0.999' + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.0 + --dit_precision "fp32" + --flow_shift 5 + --seed 1000 + --use_ema True + --ema_decay 0.99 + --ema_start_step 100 + --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors" +) + +# Self-forcing DMD arguments +dmd_args=( + --dmd_denoising_steps '1000,750,500,250' + --min_timestep_ratio 0.02 + --max_timestep_ratio 0.98 + --dfake_gen_update_ratio 5 + --real_score_guidance_scale 3.0 + --fake_score_learning_rate 8e-6 + --fake_score_betas '0.0,0.999' +) + +# Self-forcing specific arguments +self_forcing_args=( + --independent_first_frame False # Whether to treat first frame independently + --same_step_across_blocks False # Whether to use same denoising step across all blocks + --last_step_only False # Whether to only use the last denoising step + --context_noise 0 # Amount of noise to add during context caching (0 = no noise) + --validate_cache_structure False # Set to True for debugging KV cache issues +) + +torchrun \ +--nnodes 1 \ +--nproc_per_node $NUM_GPUS \ +--master_port $MASTER_PORT \ + fastvideo/training/wan_self_forcing_distillation_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ + "${dmd_args[@]}" \ + "${self_forcing_args[@]}" \ No newline at end of file diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index b2b0e162a..465ff9c20 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -244,27 +244,27 @@ class CausalWanTransformerBlock(nn.Module): current_start: int = 0, cache_start: int | None = None, ) -> torch.Tensor: - logger.info("temb.shape: %s", temb.shape) + # logger.info("temb.shape: %s", temb.shape) num_frames = temb.shape[1] - logger.info("first hidden_states.shape: %s", hidden_states.shape) - logger.info("num_frames: %s", num_frames) + # logger.info("first hidden_states.shape: %s", hidden_states.shape) + # logger.info("num_frames: %s", num_frames) if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) frame_seqlen = hidden_states.shape[1] // temb.shape[1] - logger.info("frame_seqlen: %s", frame_seqlen) + # logger.info("frame_seqlen: %s", frame_seqlen) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 e = self.scale_shift_table + temb.float() - logger.info("e.shape: %s", e.shape) + # logger.info("e.shape: %s", e.shape) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( 6, dim=2) assert shift_msa.dtype == torch.float32 # 1. Self-attention - logger.info("hidden_states.shape: %s", hidden_states.shape) - logger.info("scale_msa.shape: %s", scale_msa.shape) - logger.info("shift_msa.shape: %s", shift_msa.shape) + # logger.info("hidden_states.shape: %s", hidden_states.shape) + # logger.info("scale_msa.shape: %s", scale_msa.shape) + # logger.info("shift_msa.shape: %s", shift_msa.shape) norm_hidden_states_unflattened = self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) # logger.info("norm_hidden_states_unflattened.shape: %s", norm_hidden_states_unflattened.shape) @@ -316,7 +316,7 @@ class CausalWanTransformerBlock(nn.Module): ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) # logger.info("after mlp_residual norm_hidden_states.shape: %s", norm_hidden_states.shape) - logger.info("after mlp_residual hidden_states.shape: %s", hidden_states.shape) + # logger.info("after mlp_residual hidden_states.shape: %s", hidden_states.shape) hidden_states = hidden_states.to(orig_dtype) return hidden_states @@ -513,7 +513,7 @@ class CausalWanTransformer3DModel(BaseDiT): hidden_states = hidden_states.flatten(2).transpose(1, 2) # logger.info("forward inference flattened and transposed hidden_states.shape: %s", hidden_states.shape) - logger.info("timestep shape: %s", timestep.shape) + # logger.info("timestep shape: %s", timestep.shape) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py index 0815ea72a..b39194d2d 100644 --- a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -635,13 +635,34 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, noise: torch.Tensor, timestep: torch.IntTensor, ) -> torch.Tensor: + """ + Args: + clean_latent: the clean latent with shape [B, C, H, W], + where B is batch_size or batch_size * num_frames + noise: the noise with shape [B, C, H, W] + timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames] + + Returns: + the corrupted latent with shape [B, C, H, W] + """ + # If timestep is [bs, num_frames] + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + assert timestep.numel() == clean_latent.shape[0] + elif timestep.ndim == 1: + # If timestep is [1] + if timestep.shape[0] == 1: + timestep = timestep.expand(clean_latent.shape[0]) + else: + assert timestep.numel() == clean_latent.shape[0] + else: + raise ValueError(f"[add_noise] Invalid timestep shape: {timestep.shape}") + # timestep shape should be [B] + self.sigmas = self.sigmas.to(noise.device) - # timestep = timestep.expand(clean_latent.shape[0]) self.timesteps = self.timesteps.to(noise.device) - logger.info("self.timesteps shape: %s", self.timesteps.shape) - logger.info("timestep shape: %s", timestep.shape) timestep_id = torch.argmin( - (self.timesteps.unsqueeze(0).unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) sample = (1 - sigma) * clean_latent + sigma * noise return sample.type_as(noise) @@ -652,4 +673,4 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, return sample def __len__(self) -> int: - return self.config.num_train_timesteps + return self.config.num_train_timesteps \ No newline at end of file diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index 5eaf11fe4..a7adbb99e 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -148,11 +148,30 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor, scheduler: Any) -> torch.Tensor: """ Convert predicted noise to clean latent. + + Args: + pred_noise: the predicted noise with shape [B, C, H, W] + where B is batch_size or batch_size * num_frames + noise_input_latent: the noisy latent with shape [B, C, H, W], + timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames] + scheduler: the scheduler + + Returns: + the predicted video with shape [B, C, H, W] """ - logger.info(f"timestep: {timestep.shape}") - logger.info(f"noise_input_latent: {noise_input_latent.shape}") - logger.info(f"pred_noise: {pred_noise.shape}") - timestep = timestep.expand(noise_input_latent.shape[0]) + # If timestep is [bs, num_frames] + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + assert timestep.numel() == noise_input_latent.shape[0] + elif timestep.ndim == 1: + # If timestep is [1] + if timestep.shape[0] == 1: + timestep = timestep.expand(noise_input_latent.shape[0]) + else: + assert timestep.numel() == noise_input_latent.shape[0] + else: + raise ValueError(f"[pred_noise_to_pred_video] Invalid timestep shape: {timestep.shape}") + # timestep shape should be [B] dtype = pred_noise.dtype device = pred_noise.device pred_noise = pred_noise.float().to(device) diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 8b4c52dcf..dd802295d 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -12,6 +12,7 @@ from typing import Any import imageio import numpy as np import torch +import torch.distributed as dist import torch.nn.functional as F import torchvision from einops import rearrange @@ -728,16 +729,17 @@ class DistillationPipeline(TrainingPipeline): "encoder_hidden_states": training_batch.encoder_hidden_states, "encoder_attention_mask": training_batch.encoder_attention_mask, } - unconditional_dict = { - "encoder_hidden_states": self.negative_prompt_embeds, - "encoder_attention_mask": self.negative_prompt_attention_mask, - } + if getattr(self, "negative_prompt_embeds", None) is not None: + unconditional_dict = { + "encoder_hidden_states": self.negative_prompt_embeds, + "encoder_attention_mask": self.negative_prompt_attention_mask, + } + training_batch.unconditional_dict = unconditional_dict training_batch.dmd_latent_vis_dict = {} training_batch.fake_score_latent_vis_dict = {} training_batch.conditional_dict = conditional_dict - training_batch.unconditional_dict = unconditional_dict training_batch.raw_latent_shape = training_batch.latents.shape training_batch.latents = training_batch.latents.permute(0, 2, 1, 3, 4) self.video_latent_shape = training_batch.latents.shape @@ -762,8 +764,14 @@ class DistillationPipeline(TrainingPipeline): for _ in range(gradient_accumulation_steps): batch = self._prepare_distillation(training_batch) batch = self._get_next_batch(batch) + if dist.get_rank() == 0: + logger.info("batch.latents shape: %s", batch.latents.shape) batch = self._normalize_dit_input(batch) + if dist.get_rank() == 0: + logger.info("batch.latents shape: %s", batch.latents.shape) batch = self._prepare_dit_inputs(batch) + if dist.get_rank() == 0: + logger.info("batch.latents shape: %s", batch.latents.shape) batch = self._build_attention_metadata(batch) batch.attn_metadata_vsa = copy.deepcopy(batch.attn_metadata) if batch.attn_metadata is not None: diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 31d4bad0f..212a71a92 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -158,9 +158,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): noise = noise[:, self.rank_in_sp_group, :, :, :, :] noisy_latent = self.noise_scheduler.add_noise(latents.flatten(0, 1), noise.flatten(0, 1), - timestep).unflatten( - 0, - (1, latents.shape[1])) + torch.tensor([timestep], device=noise.device)) # Step 4: Build input kwargs with KV cache support training_batch = self._build_distill_input_kwargs( @@ -189,7 +187,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): pred_video = pred_noise_to_pred_video( pred_noise=pred_noise.flatten(0, 1), noise_input_latent=noisy_latent.flatten(0, 1), - timestep=timestep, + timestep=torch.tensor([timestep], device=noisy_latent.device), scheduler=self.noise_scheduler).unflatten(0, pred_noise.shape[:2]) self._reset_simulation_caches(kv_cache, crossattn_cache) @@ -253,6 +251,9 @@ class SelfForcingDistillationPipeline(DistillationPipeline): num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0 num_output_frames = num_frames + num_input_frames + if dist.get_rank() == 0: + logger.info(f"num_frames: {num_frames}") + logger.info(f"num_input_frames: {num_input_frames}") output = torch.zeros([batch_size, num_output_frames, num_channels, height, width], device=noise.device, dtype=noise.dtype) @@ -303,7 +304,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): exit_flag = (index == exit_flags[block_index]) timestep = torch.ones([batch_size, current_num_frames], device=noise.device, dtype=torch.int64) * current_timestep - logger.info("timestep shape at initalization: %s", timestep.shape) if not exit_flag: with torch.no_grad(): @@ -311,7 +311,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): training_batch_temp = self._build_distill_input_kwargs( noisy_input, timestep, training_batch.conditional_dict, training_batch) - logger.info("timestep shape in generator: %s", timestep.shape) pred_flow = self.transformer( hidden_states=training_batch_temp.input_kwargs['hidden_states'], encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'], @@ -325,21 +324,18 @@ class SelfForcingDistillationPipeline(DistillationPipeline): denoised_pred = pred_noise_to_pred_video( pred_noise=pred_flow.flatten(0, 1), noise_input_latent=noisy_input.flatten(0, 1), - timestep=timestep.flatten(), + timestep=timestep, scheduler=self.noise_scheduler).unflatten(0, pred_flow.shape[:2]) next_timestep = self.denoising_step_list[index + 1] - logger.info("denoised_pred shape before: %s", denoised_pred.shape) noisy_input = self.noise_scheduler.add_noise( denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), next_timestep * torch.ones([batch_size * current_num_frames], device=noise.device, dtype=torch.long) ).unflatten(0, denoised_pred.shape[:2]) - logger.info("denoised_pred shape after: %s", denoised_pred.shape) else: # Final prediction with gradient control if current_start_frame < start_gradient_frame_index: - logger.info("timestep shape in generator: %s", timestep.shape) with torch.no_grad(): training_batch_temp = self._build_distill_input_kwargs( noisy_input, timestep, training_batch.conditional_dict, training_batch) @@ -354,7 +350,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): current_start=current_start_frame * self.frame_seq_length ).permute(0, 2, 1, 3, 4) else: - logger.info("timestep shape in generator: %s", timestep.shape) training_batch_temp = self._build_distill_input_kwargs( noisy_input, timestep, training_batch.conditional_dict, training_batch) @@ -371,7 +366,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): denoised_pred = pred_noise_to_pred_video( pred_noise=pred_flow.flatten(0, 1), noise_input_latent=noisy_input.flatten(0, 1), - timestep=timestep.flatten(), + timestep=timestep, scheduler=self.noise_scheduler).unflatten(0, pred_flow.shape[:2]) break @@ -380,13 +375,11 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Step 3.3: rerun with timestep zero to update the cache context_timestep = torch.ones_like(timestep) * self.context_noise - logger.info("denoised_pred shape before: %s", denoised_pred.shape) denoised_pred = self.noise_scheduler.add_noise( denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), - context_timestep * torch.ones([batch_size * current_num_frames], device=noise.device, dtype=torch.long) + context_timestep ).unflatten(0, denoised_pred.shape[:2]) - logger.info("denoised_pred shape after: %s", denoised_pred.shape) with torch.no_grad(): training_batch_temp = self._build_distill_input_kwargs(