From aabf1fe135014fb2757b6dee7af66f713c34ee00 Mon Sep 17 00:00:00 2001 From: RandNMR73 Date: Tue, 9 Sep 2025 05:04:16 +0000 Subject: [PATCH] inference works after changes added new branch Stop backprop through kv_cache Enable timestep warping & using SelfForcing scheduler checkpoint text preprocessing ready fix rotary embedding Change Wan DiT to have 0 numerical diff with SF's Wan wan2.2 training doesn't crash training hangs --- examples/distill/SFWan2.2-A14B/distill_dmd.sh | 160 ++++++++++ examples/inference/basic/basic.py | 14 +- fastvideo/configs/pipelines/wan.py | 6 + fastvideo/fastvideo_args.py | 4 +- fastvideo/layers/layernorm.py | 10 +- fastvideo/models/dits/causal_wanvideo.py | 96 +++--- fastvideo/models/dits/wanvideo.py | 220 ++++++++----- .../pipelines/stages/causal_denoising.py | 14 +- .../transformers/test_ori_causal_wanvideo.py | 292 ++++++++++++++++++ .../tests/transformers/test_ori_wanvideo.py | 133 ++++++++ .../test_train_ori_causal_wanvideo.py | 144 +++++++++ fastvideo/training/distillation_pipeline.py | 192 ++++++++++-- .../self_forcing_distillation_pipeline.py | 87 +++++- fastvideo/training/training_pipeline.py | 201 +++++++++--- fastvideo/training/training_utils.py | 4 +- wan2.2_train.sh | 57 ++++ 16 files changed, 1414 insertions(+), 220 deletions(-) create mode 100644 examples/distill/SFWan2.2-A14B/distill_dmd.sh create mode 100644 fastvideo/tests/transformers/test_ori_causal_wanvideo.py create mode 100644 fastvideo/tests/transformers/test_ori_wanvideo.py create mode 100644 fastvideo/tests/transformers/test_train_ori_causal_wanvideo.py create mode 100644 wan2.2_train.sh diff --git a/examples/distill/SFWan2.2-A14B/distill_dmd.sh b/examples/distill/SFWan2.2-A14B/distill_dmd.sh new file mode 100644 index 000000000..b4b9bc99f --- /dev/null +++ b/examples/distill/SFWan2.2-A14B/distill_dmd.sh @@ -0,0 +1,160 @@ +#!/bin/bash +#SBATCH --job-name=t2v +#SBATCH --partition=main +#SBATCH --nodes=4 +#SBATCH --ntasks=4 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=dmd_t2v_output/t2v_%j.out +#SBATCH --error=dmd_t2v_output/t2v_%j.err +#SBATCH --exclusive + +# Basic Info +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=1 +export NCCL_DEBUG=INFO +export NCCL_DEBUG_SUBSYS=INIT,NET +CUDA_LAUNCH_BLOCKING=1 +# different cache dir for different processes +export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export TOKENIZERS_PARALLELISM=false +export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f" +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN + +# Configs +NUM_GPUS=8 + +# Model paths for Self-Forcing DMD distillation with Wan2.2: +GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Updated to Wan2.2 +REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model +FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model + +DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/" +# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset +# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/" +VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json" +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd # Updated for Wan2.2 + --output_dir "/mnt/sharefs/users/hao.zhang/SFwan2.2_t2v_finetune" + --max_train_steps 4000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 21 + --num_height 448 # Updated to match Wan2.2 config + --num_width 832 # Updated to match Wan2.2 config + --num_frames 81 # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓) + --enable_gradient_checkpointing_type "full" + --log_visualization + --simulate_generator_forward + --num_frame_per_block 3 # Frame generation block size for self-forcing + --enable_gradient_masking + --gradient_mask_last_n_frames 21 +) + +# Parallel arguments +parallel_args=( + --num_gpus 32 # 64 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 # 64 + --hsdp_shard_dim 32 +) + +# 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 5 + --validation_sampling_steps "4" + --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 500 + --weight_only_checkpointing_steps 500 + --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 +) + +# 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' + --warp_denoising_step +) + +# Self-forcing specific arguments +self_forcing_args=( + --independent_first_frame False # Whether to treat first frame independently + --same_step_across_blocks True # 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 +) + +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node $NUM_GPUS \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$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/examples/inference/basic/basic.py b/examples/inference/basic/basic.py index f12a4969e..316621eaf 100644 --- a/examples/inference/basic/basic.py +++ b/examples/inference/basic/basic.py @@ -9,9 +9,9 @@ def main(): # If a local path is provided, FastVideo will make a best effort # attempt to identify the optimal arguments. generator = VideoGenerator.from_pretrained( - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "Wan-AI/Wan2.1-T2V-14B-Diffusers", # FastVideo will automatically handle distributed setup - num_gpus=1, + num_gpus=4, use_fsdp_inference=True, dit_cpu_offload=False, vae_cpu_offload=False, @@ -25,9 +25,7 @@ def main(): # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" # Generate videos with the same simple API, regardless of GPU count prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." + "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open." ) video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") @@ -35,11 +33,7 @@ def main(): # Generate another video with a different prompt, without reloading the # model! prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") + "The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object") video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index 4a45c4df7..e782fe681 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -69,6 +69,7 @@ class WanT2V720PConfig(WanT2V480PConfig): # Denoising stage flow_shift: float | None = 5.0 + warp_denoising_step: bool = True @dataclass @@ -133,6 +134,11 @@ class Wan2_2_T2V_A14B_Config(WanT2V480PConfig): flow_shift: float | None = 12.0 boundary_ratio: float | None = 0.875 + # self-forcing params + dmd_denoising_steps: list[int] | None = field( + default_factory=lambda: [1000, 750, 500, 250]) + warp_denoising_step: bool = True + def __post_init__(self) -> None: self.dit_config.boundary_ratio = self.boundary_ratio diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 59a384459..db770d383 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -164,7 +164,7 @@ class FastVideoArgs: # dmd_denoising_steps: List[int] | None = field(default=None) # MoE parameters used by Wan2.2 - boundary_ratio: float | None = None + boundary_ratio: float | None = 0.875 @property def training_mode(self) -> bool: @@ -1151,4 +1151,4 @@ class TrainingArgs(FastVideoArgs): def parse_int_list(value: str) -> list[int]: if not value: return [] - return [int(x.strip()) for x in value.split(",")] + return [int(x.strip()) for x in value.split(",")] \ No newline at end of file diff --git a/fastvideo/layers/layernorm.py b/fastvideo/layers/layernorm.py index d93e1d966..1f9405f39 100644 --- a/fastvideo/layers/layernorm.py +++ b/fastvideo/layers/layernorm.py @@ -212,9 +212,9 @@ class ScaleResidualLayerNormScaleShift(nn.Module): frame_seqlen = normalized.shape[1] // num_frames modulated = ( normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * - (1.0 + scale) + shift).flatten(1, 2) + (1 + scale) + shift).flatten(1, 2) else: - modulated = normalized * (1.0 + scale) + shift + modulated = normalized * (1 + scale) + shift return modulated, residual_output @@ -267,13 +267,13 @@ class LayerNormScaleShift(nn.Module): frame_seqlen = normalized.shape[1] // num_frames output = ( normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * - (1.0 + scale) + shift).flatten(1, 2) + (1 + scale) + shift).flatten(1, 2) else: # scale.shape: [batch_size, 1, inner_dim] # shift.shape: [batch_size, 1, inner_dim] - output = normalized * (1.0 + scale) + shift + output = normalized * (1 + scale) + shift if self.compute_dtype == torch.float32: output = output.to(x.dtype) - return output + return output \ No newline at end of file diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index a02b68875..bbaa6fe9d 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -179,7 +179,7 @@ class CausalWanTransformerBlock(nn.Module): super().__init__() # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False) self.to_q = ReplicatedLinear(dim, dim, bias=True) self.to_k = ReplicatedLinear(dim, dim, bias=True) self.to_v = ReplicatedLinear(dim, dim, bias=True) @@ -212,8 +212,7 @@ class CausalWanTransformerBlock(nn.Module): norm_type="layer", eps=eps, elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) # 2. Cross-attention # Only T2V for now @@ -226,8 +225,7 @@ class CausalWanTransformerBlock(nn.Module): norm_type="layer", eps=eps, elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") @@ -252,29 +250,34 @@ class CausalWanTransformerBlock(nn.Module): if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) num_frames = temb.shape[1] - frame_seqlen = hidden_states.shape[1] // num_frames + frame_seqlen = hidden_states.shape[1] // num_frames bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 - e = self.scale_shift_table + temb.float() + e = self.scale_shift_table + temb # e.shape: [batch_size, num_frames, 6, inner_dim] assert e.shape == (bs, num_frames, 6, self.hidden_dim) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( 6, dim=2) # *_msa.shape: [batch_size, num_frames, 1, inner_dim] - assert shift_msa.dtype == torch.float32 + # assert shift_msa.dtype == torch.float32 + + # logger.info("temb sum: %s, dtype: %s", temb.float().sum().item(), temb.dtype) + # logger.info("scale_msa sum: %s, dtype: %s", scale_msa.float().sum().item(), scale_msa.dtype) + # logger.info("shift_msa sum: %s, dtype: %s", shift_msa.float().sum().item(), shift_msa.dtype) # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * - (1 + scale_msa) + shift_msa).flatten(1, 2).to(orig_dtype) + norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1 + scale_msa) + shift_msa).flatten(1, 2) + # logger.info("norm_hidden_states sum: %s, shape: %s", norm_hidden_states.float().sum().item(), norm_hidden_states.shape) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) if self.norm_q is not None: - query = self.norm_q(query) + query = self.norm_q.forward_native(query) if self.norm_k is not None: - key = self.norm_k(key) + key = self.norm_k.forward_native(key) query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) @@ -288,8 +291,6 @@ class CausalWanTransformerBlock(nn.Module): null_shift = null_scale = torch.tensor([0], device=hidden_states.device) norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) # 2. Cross-attention attn_output = self.attn2(norm_hidden_states, @@ -298,13 +299,10 @@ class CausalWanTransformerBlock(nn.Module): crossattn_cache=crossattn_cache) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) - hidden_states = hidden_states.to(orig_dtype) return hidden_states @@ -367,8 +365,7 @@ class CausalWanTransformer3DModel(BaseDiT): norm_type="layer", eps=config.eps, elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) self.proj_out = nn.Linear( inner_dim, config.out_channels * math.prod(config.patch_size)) self.scale_shift_table = nn.Parameter( @@ -491,12 +488,16 @@ class CausalWanTransformer3DModel(BaseDiT): ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos.float(), - freqs_sin.float()) if freqs_cos is not None else None + freqs_cis = (freqs_cos, + freqs_sin) if freqs_cos is not None else None hidden_states = self.patch_embedding(hidden_states) + grid_sizes = torch.stack( + [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) hidden_states = hidden_states.flatten(2).transpose(1, 2) + encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1) + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) @@ -543,14 +544,9 @@ class CausalWanTransformer3DModel(BaseDiT): hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, - post_patch_height, - post_patch_width, p_t, p_h, p_w, - -1) - hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) - output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + output = self.unpatchify(hidden_states, grid_sizes) - return output + return torch.stack(output) def _forward_train(self, hidden_states: torch.Tensor, @@ -591,8 +587,8 @@ class CausalWanTransformer3DModel(BaseDiT): ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos.float(), - freqs_sin.float()) if freqs_cos is not None else None + freqs_cis = (freqs_cos, + freqs_sin) if freqs_cos is not None else None # Construct blockwise causal attn mask if self.block_mask is None: @@ -605,8 +601,12 @@ class CausalWanTransformer3DModel(BaseDiT): ) hidden_states = self.patch_embedding(hidden_states) + grid_sizes = torch.stack( + [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) hidden_states = hidden_states.flatten(2).transpose(1, 2) + encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1) + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) @@ -641,14 +641,9 @@ class CausalWanTransformer3DModel(BaseDiT): hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, - post_patch_height, - post_patch_width, p_t, p_h, p_w, - -1) - hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) - output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + output = self.unpatchify(hidden_states, grid_sizes) - return output + return torch.stack(output) def forward( self, @@ -659,3 +654,30 @@ class CausalWanTransformer3DModel(BaseDiT): return self._forward_inference(*args, **kwargs) else: return self._forward_train(*args, **kwargs) + + + def unpatchify(self, x, grid_sizes): + r""" + + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + + + Returns: + Tensor: + Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] + """ + + c = self.out_channels + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = u.permute(6, 0, 3, 1, 4, 2, 5) + # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous()) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out \ No newline at end of file diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index ac79e62e8..9931ed215 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -1,3 +1,5 @@ +import torch +import torch.nn as nn # SPDX-License-Identifier: Apache-2.0 import math @@ -37,16 +39,14 @@ class WanImageEmbedding(torch.nn.Module): def __init__(self, in_features: int, out_features: int): super().__init__() - self.norm1 = FP32LayerNorm(in_features) + self.norm1 = nn.LayerNorm(in_features) self.ff = MLP(in_features, in_features, out_features, act_type="gelu") - self.norm2 = FP32LayerNorm(out_features) + self.norm2 = nn.LayerNorm(out_features) - def forward(self, - encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: - dtype = encoder_hidden_states_image.dtype + def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: hidden_states = self.norm1(encoder_hidden_states_image) hidden_states = self.ff(hidden_states) - hidden_states = self.norm2(hidden_states).to(dtype) + hidden_states = self.norm2(hidden_states) return hidden_states @@ -62,7 +62,7 @@ class WanTimeTextImageEmbedding(nn.Module): super().__init__() self.time_embedder = TimestepEmbedder( - dim, frequency_embedding_size=time_freq_dim, act_layer="silu") + dim, frequency_embedding_size=time_freq_dim, act_layer="silu", freq_dtype=torch.float64) self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") @@ -156,12 +156,12 @@ class WanT2VCrossAttention(WanSelfAttention): b, n, d = x.size(0), self.num_heads, self.head_dim # compute query, key, value - q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) + q = self.norm_q.forward_native(self.to_q(x)[0]).view(b, -1, n, d) if crossattn_cache is not None: if not crossattn_cache["is_init"]: crossattn_cache["is_init"] = True - k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) crossattn_cache["k"] = k crossattn_cache["v"] = v @@ -169,7 +169,7 @@ class WanT2VCrossAttention(WanSelfAttention): k = crossattn_cache["k"] v = crossattn_cache["v"] else: - k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) # compute attention @@ -213,10 +213,10 @@ class WanI2VCrossAttention(WanSelfAttention): b, n, d = x.size(0), self.num_heads, self.head_dim # compute query, key, value - q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) - k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + q = self.norm_q.forward_native(self.to_q(x)[0]).view(b, -1, n, d) + k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) - k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view( + k_img = self.norm_added_k.forward_native(self.add_k_proj(context_img)[0]).view( b, -1, n, d) v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) img_x = self.attn(q, k_img, v_img) @@ -247,7 +247,7 @@ class WanTransformerBlock(nn.Module): super().__init__() # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False) self.to_q = ReplicatedLinear(dim, dim, bias=True) self.to_k = ReplicatedLinear(dim, dim, bias=True) self.to_v = ReplicatedLinear(dim, dim, bias=True) @@ -278,29 +278,29 @@ class WanTransformerBlock(nn.Module): norm_type="layer", eps=eps, elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) # 2. Cross-attention if added_kv_proj_dim is not None: # I2V - self.attn2 = WanI2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, + self.attn2 = WanI2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, eps=eps) + else: # T2V - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, + self.attn2 = WanT2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, eps=eps) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") @@ -319,12 +319,11 @@ class WanTransformerBlock(nn.Module): hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype - # assert orig_dtype != torch.float32 if temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table.unsqueeze(0) + temb.float() + self.scale_shift_table.unsqueeze(0) + temb ).chunk(6, dim=2) # batch_size, seq_len, 1, inner_dim shift_msa = shift_msa.squeeze(2) @@ -335,22 +334,38 @@ class WanTransformerBlock(nn.Module): c_gate_msa = c_gate_msa.squeeze(2) else: # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) - e = self.scale_shift_table + temb.float() + e = self.scale_shift_table + temb shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( 6, dim=1) - assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * - (1 + scale_msa) + shift_msa).to(orig_dtype) + norm_hidden_states_base = self.norm1(hidden_states) + + # Handle broadcasting for ti2v case where we have per-frame embeddings + if temb.dim() == 4 and scale_msa.shape[1] != norm_hidden_states_base.shape[1]: + # For ti2v: scale_msa has shape [batch, num_frames, dim] but we need [batch, seq_len, dim] + # We need to repeat each frame's embedding across its corresponding sequence positions + seq_len = norm_hidden_states_base.shape[1] + num_frames = scale_msa.shape[1] + frames_per_seq = seq_len // num_frames + + # Repeat each frame's embedding across its sequence positions for all embeddings + scale_msa = scale_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_len, -1) + shift_msa = shift_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_len, -1) + gate_msa = gate_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_len, -1) + c_shift_msa = c_shift_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_len, -1) + c_scale_msa = c_scale_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_len, -1) + c_gate_msa = c_gate_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_len, -1) + + norm_hidden_states = norm_hidden_states_base * (1 + scale_msa) + shift_msa query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) if self.norm_q is not None: - query = self.norm_q(query) + query = self.norm_q.forward_native(query) if self.norm_k is not None: - key = self.norm_k(key) + key = self.norm_k.forward_native(key) query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) @@ -370,26 +385,20 @@ class WanTransformerBlock(nn.Module): null_shift = null_scale = torch.tensor([0], device=hidden_states.device) norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) # 2. Cross-attention - attn_output = self.attn2(norm_hidden_states, - context=encoder_hidden_states, + attn_output = self.attn2(norm_hidden_states, + context=encoder_hidden_states, context_lens=None) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) - hidden_states = hidden_states.to(orig_dtype) return hidden_states - class WanTransformerBlock_VSA(nn.Module): def __init__(self, @@ -406,7 +415,7 @@ class WanTransformerBlock_VSA(nn.Module): super().__init__() # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False) self.to_q = ReplicatedLinear(dim, dim, bias=True) self.to_k = ReplicatedLinear(dim, dim, bias=True) self.to_v = ReplicatedLinear(dim, dim, bias=True) @@ -438,8 +447,7 @@ class WanTransformerBlock_VSA(nn.Module): norm_type="layer", eps=eps, elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) # 2. Cross-attention if added_kv_proj_dim is not None: @@ -459,8 +467,7 @@ class WanTransformerBlock_VSA(nn.Module): norm_type="layer", eps=eps, elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") @@ -480,23 +487,55 @@ class WanTransformerBlock_VSA(nn.Module): bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 - e = self.scale_shift_table + temb.float() - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( - 6, dim=1) - assert shift_msa.dtype == torch.float32 + + print(f"self.norm1(hidden_states).shape: {self.norm1(hidden_states).shape}") + print(f"scale_msa.shape: {scale_msa.shape}") + print(f"shift_msa.shape: {shift_msa.shape}") + + if temb.dim() == 4: + # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table.unsqueeze(0) + temb + ).chunk(6, dim=2) + # batch_size, seq_len, 1, inner_dim + shift_msa = shift_msa.squeeze(2) + scale_msa = scale_msa.squeeze(2) + gate_msa = gate_msa.squeeze(2) + c_shift_msa = c_shift_msa.squeeze(2) + c_scale_msa = c_scale_msa.squeeze(2) + c_gate_msa = c_gate_msa.squeeze(2) + + # Handle broadcasting for ti2v case where we have per-frame embeddings + if scale_msa.shape[1] != seq_length: + # For ti2v: scale_msa has shape [batch, num_frames, dim] but we need [batch, seq_len, dim] + num_frames = scale_msa.shape[1] + frames_per_seq = seq_length // num_frames + + # Repeat each frame's embedding across its sequence positions for all embeddings + scale_msa = scale_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_length, -1) + shift_msa = shift_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_length, -1) + gate_msa = gate_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_length, -1) + c_shift_msa = c_shift_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_length, -1) + c_scale_msa = c_scale_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_length, -1) + c_gate_msa = c_gate_msa.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(bs, seq_length, -1) + else: + # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) + e = self.scale_shift_table + temb + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( + 6, dim=1) # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * - (1 + scale_msa) + shift_msa).to(orig_dtype) + norm_hidden_states = (self.norm1(hidden_states) * + (1 + scale_msa) + shift_msa) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) gate_compress, _ = self.to_gate_compress(norm_hidden_states) if self.norm_q is not None: - query = self.norm_q(query) + query = self.norm_q.forward_native(query) if self.norm_k is not None: - key = self.norm_k(key) + key = self.norm_k.forward_native(key) query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) @@ -521,8 +560,6 @@ class WanTransformerBlock_VSA(nn.Module): null_shift = null_scale = torch.tensor([0], device=hidden_states.device) norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) # 2. Cross-attention attn_output = self.attn2(norm_hidden_states, @@ -530,17 +567,15 @@ class WanTransformerBlock_VSA(nn.Module): context_lens=None) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) - hidden_states = hidden_states.to(orig_dtype) return hidden_states + class WanTransformer3DModel(CachableDiT): _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions _compile_conditions = WanVideoConfig()._compile_conditions @@ -562,6 +597,7 @@ class WanTransformer3DModel(CachableDiT): self.num_channels_latents = config.num_channels_latents self.patch_size = config.patch_size self.text_len = config.text_len + self.attention_head_dim = self.hidden_size // self.num_attention_heads # 1. Patch & position embedding self.patch_embedding = PatchEmbed(in_chans=config.in_channels, @@ -598,8 +634,7 @@ class WanTransformer3DModel(CachableDiT): norm_type="layer", eps=config.eps, elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + dtype=torch.float32) self.proj_out = nn.Linear( inner_dim, config.out_channels * math.prod(config.patch_size)) self.scale_shift_table = nn.Parameter( @@ -659,10 +694,12 @@ class WanTransformer3DModel(CachableDiT): rope_theta=10000) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos.float(), - freqs_sin.float()) if freqs_cos is not None else None + freqs_cis = (freqs_cos, + freqs_sin) if freqs_cos is not None else None hidden_states = self.patch_embedding(hidden_states) + grid_sizes = torch.stack( + [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) hidden_states = hidden_states.flatten(2).transpose(1, 2) # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v) @@ -672,6 +709,8 @@ class WanTransformer3DModel(CachableDiT): else: ts_seq_len = None + encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1) + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len) if ts_seq_len is not None: @@ -721,6 +760,17 @@ class WanTransformer3DModel(CachableDiT): shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2) shift = shift.squeeze(2) scale = scale.squeeze(2) + + # Handle broadcasting for ti2v case where we have per-frame embeddings + if shift.shape[1] != hidden_states.shape[1]: + # For ti2v: shift has shape [batch, num_frames, dim] but we need [batch, seq_len, dim] + seq_len = hidden_states.shape[1] + num_frames = shift.shape[1] + frames_per_seq = seq_len // num_frames + + # Repeat each frame's embedding across its sequence positions + shift = shift.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(batch_size, seq_len, -1) + scale = scale.unsqueeze(2).repeat(1, 1, frames_per_seq, 1).view(batch_size, seq_len, -1) else: # batch_size, inner_dim shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) @@ -728,14 +778,35 @@ class WanTransformer3DModel(CachableDiT): hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, - post_patch_height, - post_patch_width, p_t, p_h, p_w, - -1) - hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) - output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + output = self.unpatchify(hidden_states, grid_sizes) - return output + return torch.stack(output) + + def unpatchify(self, x, grid_sizes): + r""" + + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + + + Returns: + Tensor: + Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] + """ + + c = self.out_channels + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = u.permute(6, 0, 3, 1, 4, 2, 5) + # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous()) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out def maybe_cache_states(self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor) -> None: @@ -827,5 +898,4 @@ class WanTransformer3DModel(CachableDiT): if self.is_even: return hidden_states + self.previous_residual_even else: - return hidden_states + self.previous_residual_odd - \ No newline at end of file + return hidden_states + self.previous_residual_odd \ No newline at end of file diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 925c570f4..d12b56c70 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -57,21 +57,27 @@ class CausalDMDDenosingStage(DenoisingStage): fastvideo_args: FastVideoArgs, ) -> ForwardBatch: target_dtype = torch.bfloat16 + logger.info("Using target dtype: %s", target_dtype, local_main_process_only=False) autocast_enabled = (target_dtype != torch.float32 ) and not fastvideo_args.disable_autocast + logger.info("Using autocast enabled: %s", autocast_enabled, local_main_process_only=False) + latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] + logger.info("Using latent seq length: %s", latent_seq_length, local_main_process_only=False) patch_ratio = self.transformer.config.arch_config.patch_size[ -1] * self.transformer.config.arch_config.patch_size[-2] + logger.info("Using patch ratio: %s", patch_ratio, local_main_process_only=False) self.frame_seq_length = latent_seq_length // patch_ratio + logger.info("Using frame seq length: %s", self.frame_seq_length, local_main_process_only=False) # TODO(will): make this a parameter once we add i2v support - independent_first_frame = self.transformer.independent_first_frame - + independent_first_frame = self.transformer.independent_first_frame if hasattr(self.transformer, 'independent_first_frame') else False + logger.info("Using independent first frame: %s", independent_first_frame, local_main_process_only=False) # Timesteps for DMD timesteps = torch.tensor( fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu() - + logger.info("Using timesteps: %s", timesteps, local_main_process_only=False) if fastvideo_args.pipeline_config.warp_denoising_step: logger.info("Warping timesteps...") scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), @@ -79,7 +85,7 @@ class CausalDMDDenosingStage(DenoisingStage): dtype=torch.float32))) timesteps = scheduler_timesteps[1000 - timesteps] timesteps = timesteps.to(get_local_torch_device()) - logger.info("Using timesteps: %s", timesteps) + logger.info("Using timesteps: %s", timesteps, local_main_process_only=False) # Image kwargs (kept empty unless caller provides compatible args) image_kwargs: dict = {} diff --git a/fastvideo/tests/transformers/test_ori_causal_wanvideo.py b/fastvideo/tests/transformers/test_ori_causal_wanvideo.py new file mode 100644 index 000000000..7be5fbbc8 --- /dev/null +++ b/fastvideo/tests/transformers/test_ori_causal_wanvideo.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +import math + +import numpy as np +import pytest +import torch +from fastvideo.wan.modules.causal_model import CausalWanModel + +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.forward_context import set_forward_context +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TransformerLoader +from fastvideo.models.dits.causal_wanvideo import CausalWanTransformer3DModel +from fastvideo.utils import maybe_download_model +from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +BASE_MODEL_PATH = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + 'data', BASE_MODEL_PATH)) +TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") + + +@pytest.mark.usefixtures("distributed_setup") +def test_ori_causal_wan_transformer(): + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.bfloat16 + precision_str = "bf16" + args = FastVideoArgs(model_path=TRANSFORMER_PATH, + dit_cpu_offload=True, + pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str)) + args.device = device + + loader = TransformerLoader() + model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision) + + model1 = CausalWanModel.from_pretrained( + "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device, + torch_dtype=precision).to(device, dtype=precision).requires_grad_(False) + causal_state_dict = torch.load("/mnt/weka/home/hao.zhang/wei/Self-Forcing/checkpoints/self_forcing_dmd.pt")["generator_ema"] + new_state_dict = {} + for k, v in causal_state_dict.items(): + if k.startswith("model."): + new_state_dict[k.replace("model.", "")] = v + causal_state_dict = new_state_dict + model1.load_state_dict(causal_state_dict) + + total_params = sum(p.numel() for p in model1.parameters()) + # Calculate weight sum for model1 (converting to float64 to avoid overflow) + weight_sum_model1 = sum( + p.to(torch.float64).sum().item() for p in model1.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model1 = weight_sum_model1 / total_params + logger.info("Model 1 weight sum: %s", weight_sum_model1) + logger.info("Model 1 weight mean: %s", weight_mean_model1) + + # Calculate weight sum for model2 (converting to float64 to avoid overflow) + total_params_model2 = sum(p.numel() for p in model2.parameters()) + weight_sum_model2 = sum( + p.to(torch.float64).sum().item() for p in model2.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model2 = weight_sum_model2 / total_params_model2 + logger.info("Model 2 weight sum: %s", weight_sum_model2) + logger.info("Model 2 weight mean: %s", weight_mean_model2) + + weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2) + logger.info("Weight sum difference: %s", weight_sum_diff) + weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2) + logger.info("Weight mean difference: %s", weight_mean_diff) + + # Set both models to eval mode + model1 = model1.eval() + model2 = model2.eval() + + # Create identical inputs for both models + batch_size = 1 + text_seq_len = 30 + + # Video latents [B, C, T, H, W] + hidden_states = torch.randn(batch_size, + 16, + 12, + 160, + 90, + device=device, + dtype=precision) + block_sizes = [3 for _ in range(4)] + timesteps = [1000, 750, 500, 250] + + # Text embeddings [B, L, D] (including global token) + encoder_hidden_states = torch.randn(batch_size, + text_seq_len + 1, + 4096, + device=device, + dtype=precision) + + output1 = _causal_inference(model1, hidden_states.clone(), encoder_hidden_states.clone(), block_sizes, timesteps, precision) + logger.info("Finish inference for model1") + output2 = _causal_inference(model2, hidden_states.clone(), encoder_hidden_states.clone(), block_sizes, timesteps, precision) + + # Check if outputs have the same shape + assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" + assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}" + + logger.info("Output 1 Sum: %s", output1.float().sum().item()) + logger.info("Output 2 Sum: %s", output2.float().sum().item()) + + # Check if outputs are similar (allowing for small numerical differences) + max_diff = torch.max(torch.abs(output1 - output2)) + mean_diff = torch.mean(torch.abs(output1 - output2)) + logger.info("Max Diff: %s", max_diff.item()) + logger.info("Mean Diff: %s", mean_diff.item()) + assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}" + # mean diff + assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}" + +def _causal_inference(transformer, latents, prompt_embeds, block_sizes, timesteps, target_dtype): + forward_batch = ForwardBatch( + data_type="dummy", + ) + start_index = 0 + pos_start_base = 0 + frame_seq_length = latents.shape[-1] * latents.shape[-2] // (WanVideoConfig().arch_config.patch_size[-1] * WanVideoConfig().arch_config.patch_size[-2]) + seq_len = frame_seq_length * latents.shape[2] + kv_cache1 = _initialize_kv_cache(transformer, batch_size=latents.shape[0], + kv_cache_size=frame_seq_length * latents.shape[2], + dtype=target_dtype, + device=latents.device) + crossattn_cache = _initialize_crossattn_cache( + transformer, + batch_size=latents.shape[0], + max_text_len=WanVideoConfig().arch_config.text_len, + dtype=target_dtype, + device=latents.device) + for current_num_frames, t_cur in zip(block_sizes, timesteps): + # logger.info(f"Current frame idx: {start_index}, Current timestep: {t_cur}") + # logger.info(f"k cache sum: {sum(kv_cache['k'].float().sum().item() for kv_cache in kv_cache1)}, v cache sum: {sum(kv_cache['v'].float().sum().item() for kv_cache in kv_cache1)}") + # logger.info(f"latents sum: {latents.float().sum().item()}, encoder_hidden_states sum: {prompt_embeds.float().sum().item()}") + current_latents = latents[:, :, start_index:start_index + + current_num_frames, :, :] + + attn_metadata = None + + with set_forward_context(current_timestep=0, + attn_metadata=attn_metadata, + forward_batch=forward_batch): + # Run transformer; follow DMD stage pattern + t_expanded_noise = t_cur * torch.ones( + (current_latents.shape[0], 1), + device=current_latents.device, + dtype=torch.long) + if isinstance(transformer, CausalWanModel): + pred_noise_btchw = transformer( + x=current_latents, + context=prompt_embeds, + t=t_expanded_noise, + seq_len=seq_len, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + frame_seq_length + ) + elif isinstance(transformer, CausalWanTransformer3DModel): + pred_noise_btchw = transformer( + current_latents, + prompt_embeds, + t_expanded_noise, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + frame_seq_length, + start_frame=start_index + ) + + # Write back and advance + latents[:, :, start_index:start_index + + current_num_frames, :, :] = pred_noise_btchw.clone() + + # Re-run with context timestep to update KV cache using clean context + context_noise = 0 + t_context = torch.ones([latents.shape[0]], + device=latents.device, + dtype=torch.long) * int(context_noise) + context_bcthw = pred_noise_btchw.to(target_dtype) + with set_forward_context(current_timestep=0, + attn_metadata=attn_metadata, + forward_batch=forward_batch): + t_expanded_context = t_context.unsqueeze(1) + if isinstance(transformer, CausalWanModel): + _ = transformer( + x=context_bcthw, + context=prompt_embeds, + t=t_expanded_context, + seq_len=seq_len, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + frame_seq_length + ) + elif isinstance(transformer, CausalWanTransformer3DModel): + _ = transformer( + context_bcthw, + prompt_embeds, + t_expanded_context, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + frame_seq_length, + start_frame=start_index + ) + start_index += current_num_frames + + return latents + +def _initialize_kv_cache(transformer, batch_size, kv_cache_size, dtype, device) -> None: + """ + Initialize a Per-GPU KV cache aligned with the Wan model assumptions. + """ + kv_cache1 = [] + if isinstance(transformer, CausalWanModel): + num_attention_heads = transformer.num_heads + attention_head_dim = transformer.dim // transformer.num_heads + elif isinstance(transformer, CausalWanTransformer3DModel): + num_attention_heads = transformer.num_attention_heads + attention_head_dim = transformer.attention_head_dim + + for _ in range(len(transformer.blocks)): + kv_cache1.append({ + "k": + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device), + }) + + return kv_cache1 + +def _initialize_crossattn_cache(transformer, batch_size, max_text_len, dtype, + device) -> None: + """ + Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions. + """ + crossattn_cache = [] + if isinstance(transformer, CausalWanModel): + num_attention_heads = transformer.num_heads + attention_head_dim = transformer.dim // transformer.num_heads + elif isinstance(transformer, CausalWanTransformer3DModel): + num_attention_heads = transformer.num_attention_heads + attention_head_dim = transformer.attention_head_dim + + for _ in range(len(transformer.blocks)): + crossattn_cache.append({ + "k": + torch.zeros([ + batch_size, max_text_len, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size, max_text_len, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "is_init": + False, + }) + return crossattn_cache \ No newline at end of file diff --git a/fastvideo/tests/transformers/test_ori_wanvideo.py b/fastvideo/tests/transformers/test_ori_wanvideo.py new file mode 100644 index 000000000..11b5a3460 --- /dev/null +++ b/fastvideo/tests/transformers/test_ori_wanvideo.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +import math + +import numpy as np +import pytest +import torch +from fastvideo.wan.modules.model import WanModel + +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.forward_context import set_forward_context +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TransformerLoader +from fastvideo.utils import maybe_download_model +from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + 'data', BASE_MODEL_PATH)) +TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") + + +@pytest.mark.usefixtures("distributed_setup") +def test_ori_wan_transformer(): + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.bfloat16 + precision_str = "bf16" + args = FastVideoArgs(model_path=TRANSFORMER_PATH, + dit_cpu_offload=True, + pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str)) + args.device = device + + loader = TransformerLoader() + model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision) + + model1 = WanModel.from_pretrained( + "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device, + torch_dtype=precision).to(device, dtype=precision).requires_grad_(False) + + total_params = sum(p.numel() for p in model1.parameters()) + # Calculate weight sum for model1 (converting to float64 to avoid overflow) + weight_sum_model1 = sum( + p.to(torch.float64).sum().item() for p in model1.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model1 = weight_sum_model1 / total_params + logger.info("Model 1 weight sum: %s", weight_sum_model1) + logger.info("Model 1 weight mean: %s", weight_mean_model1) + + # Calculate weight sum for model2 (converting to float64 to avoid overflow) + total_params_model2 = sum(p.numel() for p in model2.parameters()) + weight_sum_model2 = sum( + p.to(torch.float64).sum().item() for p in model2.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model2 = weight_sum_model2 / total_params_model2 + logger.info("Model 2 weight sum: %s", weight_sum_model2) + logger.info("Model 2 weight mean: %s", weight_mean_model2) + + weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2) + logger.info("Weight sum difference: %s", weight_sum_diff) + weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2) + logger.info("Weight mean difference: %s", weight_mean_diff) + + # Set both models to eval mode + model1 = model1.eval() + model2 = model2.eval() + + # Create identical inputs for both models + batch_size = 1 + text_seq_len = 30 + seq_len = math.ceil((160 * 90) / + (2 * 2) * + 21) + + # Video latents [B, C, T, H, W] + hidden_states = torch.randn(batch_size, + 16, + 21, + 160, + 90, + device=device, + dtype=precision) + + # Text embeddings [B, L, D] (including global token) + encoder_hidden_states = torch.randn(batch_size, + text_seq_len + 1, + 4096, + device=device, + dtype=precision) + + # Timestep + timestep = torch.tensor([500], device=device, dtype=precision) + + forward_batch = ForwardBatch( + data_type="dummy", + ) + + # with torch.amp.autocast('cuda', dtype=precision): + output1 = model1( + x=hidden_states, + context=encoder_hidden_states, + t=timestep, + seq_len=seq_len, + ) + with set_forward_context( + current_timestep=0, + attn_metadata=None, + forward_batch=forward_batch, + ): + output2 = model2(hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep) + + # Check if outputs have the same shape + assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" + assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}" + + # Check if outputs are similar (allowing for small numerical differences) + max_diff = torch.max(torch.abs(output1 - output2)) + mean_diff = torch.mean(torch.abs(output1 - output2)) + logger.info("Max Diff: %s", max_diff.item()) + logger.info("Mean Diff: %s", mean_diff.item()) + assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}" + # mean diff + assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}" \ No newline at end of file diff --git a/fastvideo/tests/transformers/test_train_ori_causal_wanvideo.py b/fastvideo/tests/transformers/test_train_ori_causal_wanvideo.py new file mode 100644 index 000000000..13a2c7680 --- /dev/null +++ b/fastvideo/tests/transformers/test_train_ori_causal_wanvideo.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +import math + +import numpy as np +import pytest +import torch +from fastvideo.wan.modules.causal_model import CausalWanModel + +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.forward_context import set_forward_context +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TransformerLoader +from fastvideo.utils import maybe_download_model +from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +BASE_MODEL_PATH = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + 'data', BASE_MODEL_PATH)) +TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") + + +@pytest.mark.usefixtures("distributed_setup") +def test_train_ori_causal_wan_transformer(): + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.bfloat16 + precision_str = "bf16" + args = FastVideoArgs(model_path=TRANSFORMER_PATH, + dit_cpu_offload=True, + pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str)) + args.device = device + + loader = TransformerLoader() + model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision) + + model1 = CausalWanModel.from_pretrained( + "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device, + torch_dtype=precision).to(device, dtype=precision).requires_grad_(False) + causal_state_dict = torch.load("/mnt/weka/home/hao.zhang/wei/Self-Forcing/checkpoints/self_forcing_dmd.pt")["generator_ema"] + new_state_dict = {} + for k, v in causal_state_dict.items(): + if k.startswith("model."): + new_state_dict[k.replace("model.", "")] = v + causal_state_dict = new_state_dict + model1.load_state_dict(causal_state_dict) + + model1.num_frame_per_block = 3 + model2.num_frame_per_block = 3 + + total_params = sum(p.numel() for p in model1.parameters()) + # Calculate weight sum for model1 (converting to float64 to avoid overflow) + weight_sum_model1 = sum( + p.to(torch.float64).sum().item() for p in model1.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model1 = weight_sum_model1 / total_params + logger.info("Model 1 weight sum: %s", weight_sum_model1) + logger.info("Model 1 weight mean: %s", weight_mean_model1) + + # Calculate weight sum for model2 (converting to float64 to avoid overflow) + total_params_model2 = sum(p.numel() for p in model2.parameters()) + weight_sum_model2 = sum( + p.to(torch.float64).sum().item() for p in model2.parameters()) + # Also calculate mean for more stable comparison + weight_mean_model2 = weight_sum_model2 / total_params_model2 + logger.info("Model 2 weight sum: %s", weight_sum_model2) + logger.info("Model 2 weight mean: %s", weight_mean_model2) + + weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2) + logger.info("Weight sum difference: %s", weight_sum_diff) + weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2) + logger.info("Weight mean difference: %s", weight_mean_diff) + + # Set both models to eval mode + model1 = model1.eval() + model2 = model2.eval() + + # Create identical inputs for both models + batch_size = 1 + text_seq_len = 30 + seq_len = math.ceil((160 * 90) / + (2 * 2) * + 21) + + # Video latents [B, C, T, H, W] + hidden_states = torch.randn(batch_size, + 16, + 21, + 160, + 90, + device=device, + dtype=precision) + + # Text embeddings [B, L, D] (including global token) + encoder_hidden_states = torch.randn(batch_size, + text_seq_len + 1, + 4096, + device=device, + dtype=precision) + + # Timestep + timestep = torch.randint(0, 1000, (batch_size, 21), device=device, dtype=torch.long) + logger.info("timestep: %s", timestep) + + forward_batch = ForwardBatch( + data_type="dummy", + ) + + # with torch.amp.autocast('cuda', dtype=precision): + output1 = model1( + x=hidden_states, + context=encoder_hidden_states, + t=timestep, + seq_len=seq_len, + ) + with set_forward_context( + current_timestep=0, + attn_metadata=None, + forward_batch=forward_batch, + ): + output2 = model2(hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep) + + # Check if outputs have the same shape + assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" + assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}" + + # Check if outputs are similar (allowing for small numerical differences) + max_diff = torch.max(torch.abs(output1 - output2)) + mean_diff = torch.mean(torch.abs(output1 - output2)) + logger.info("Max Diff: %s", max_diff.item()) + logger.info("Mean Diff: %s", mean_diff.item()) + assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}" + # mean diff + assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}" \ No newline at end of file diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index e75c1286b..ba97d1d08 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 @@ -39,7 +40,7 @@ from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.training.training_utils import ( EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint, - shift_timestep) + shift_timestep, compute_density_for_timestep_sampling, get_sigmas) from fastvideo.utils import (is_vsa_available, maybe_download_model, set_random_seed, verify_model_config_and_directory) @@ -71,6 +72,7 @@ class DistillationPipeline(TrainingPipeline): current_trainstep: int video_latent_shape: tuple[int, ...] video_latent_shape_sp: tuple[int, ...] + train_fake_score_transformer_2: bool = False def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): raise RuntimeError( @@ -90,40 +92,82 @@ class DistillationPipeline(TrainingPipeline): self.noise_scheduler = FlowMatchEulerDiscreteScheduler( shift=self.timestep_shift) + if self.training_args.boundary_ratio is not None: + self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps + else: + self.boundary_timestep = None + if training_args.real_score_model_path: - logger.info("Loading real score transformer from: %s", - training_args.real_score_model_path) + logger.info( + f"Loading real score transformer from: {training_args.real_score_model_path}" + ) self.real_score_transformer = self.load_module_from_path( training_args.real_score_model_path, "transformer", training_args) + try: + self.real_score_transformer_2 = self.load_module_from_path( + training_args.real_score_model_path, "transformer_2", + training_args) + logger.info("Loaded real score transformer_2 for MoE support") + except: + logger.info("real score transformer_2 not found, using single transformer") + self.real_score_transformer_2 = None else: - self.real_score_transformer = self.get_module( - "real_score_transformer") + self.real_score_transformer = self.get_module("real_score_transformer") + self.real_score_transformer_2 = None if training_args.fake_score_model_path: - logger.info("Loading fake score transformer from: %s", - training_args.fake_score_model_path) + logger.info( + f"Loading fake score transformer from: {training_args.fake_score_model_path}" + ) self.fake_score_transformer = self.load_module_from_path( training_args.fake_score_model_path, "transformer", training_args) + try: + self.fake_score_transformer_2 = self.load_module_from_path( + training_args.fake_score_model_path, "transformer_2", + training_args) + logger.info("Loaded fake score transformer_2 for MoE support") + except: + logger.info("fake score transformer_2 not found, using single transformer") + self.fake_score_transformer_2 = None else: - self.fake_score_transformer = self.get_module( - "fake_score_transformer") + self.fake_score_transformer = self.get_module("fake_score_transformer") + self.fake_score_transformer_2 = None self.real_score_transformer.requires_grad_(False) self.real_score_transformer.eval() + if self.real_score_transformer_2 is not None: + self.real_score_transformer_2.requires_grad_(False) + self.real_score_transformer_2.eval() + + # Set training modes for fake score transformers (trainable) self.fake_score_transformer.requires_grad_(True) self.fake_score_transformer.train() + if self.fake_score_transformer_2 is not None: + self.fake_score_transformer_2.requires_grad_(True) + self.fake_score_transformer_2.train() if training_args.enable_gradient_checkpointing_type is not None: self.fake_score_transformer = apply_activation_checkpointing( self.fake_score_transformer, checkpointing_type=training_args. enable_gradient_checkpointing_type) + if self.fake_score_transformer_2 is not None: + self.fake_score_transformer_2 = apply_activation_checkpointing( + self.fake_score_transformer_2, + checkpointing_type=training_args. + enable_gradient_checkpointing_type) + self.real_score_transformer = apply_activation_checkpointing( self.real_score_transformer, checkpointing_type=training_args. enable_gradient_checkpointing_type) + if self.real_score_transformer_2 is not None: + self.real_score_transformer_2 = apply_activation_checkpointing( + self.real_score_transformer_2, + checkpointing_type=training_args. + enable_gradient_checkpointing_type) # Initialize optimizers fake_score_params = list( @@ -157,6 +201,27 @@ class DistillationPipeline(TrainingPipeline): last_epoch=self.init_steps - 1, ) + if self.fake_score_transformer_2 is not None: + fake_score_params_2 = list(filter(lambda p: p.requires_grad, + self.fake_score_transformer_2.parameters())) + self.fake_score_optimizer_2 = torch.optim.AdamW( + fake_score_params_2, + lr=fake_score_lr, + betas=betas, + weight_decay=training_args.weight_decay, + eps=1e-8, + ) + self.fake_score_lr_scheduler_2 = get_scheduler( + training_args.fake_score_lr_scheduler, + optimizer=self.fake_score_optimizer_2, + num_warmup_steps=training_args.lr_warmup_steps, + num_training_steps=training_args.max_train_steps, + num_cycles=training_args.lr_num_cycles, + power=training_args.lr_power, + min_lr_ratio=training_args.min_lr_ratio, + last_epoch=self.init_steps - 1, + ) + logger.info( "Distillation optimizers initialized: generator and fake_score") @@ -278,8 +343,14 @@ class DistillationPipeline(TrainingPipeline): """Prepare training environment for distillation.""" self.transformer.requires_grad_(True) self.transformer.train() + if self.transformer_2 is not None: + self.transformer_2.requires_grad_(True) + self.transformer_2.train() self.fake_score_transformer.requires_grad_(True) self.fake_score_transformer.train() + if self.fake_score_transformer_2 is not None: + self.fake_score_transformer_2.requires_grad_(True) + self.fake_score_transformer_2.train() return training_batch @@ -387,6 +458,33 @@ class DistillationPipeline(TrainingPipeline): else: logger.warning("Cannot reset EMA: EMA not initialized") + def _get_real_score_transformer(self, timestep: torch.Tensor): + """ + Get the appropriate real score transformer based on timestep and boundary logic. + """ + if self.real_score_transformer_2 is not None and self.boundary_timestep is not None: + if timestep.item() <= self.boundary_timestep: + return self.real_score_transformer_2 # Low noise expert + else: + return self.real_score_transformer # High noise expert + else: + return self.real_score_transformer + + def _get_fake_score_transformer(self, timestep: torch.Tensor): + """ + Get the appropriate fake score transformer based on timestep and boundary logic. + """ + if self.fake_score_transformer_2 is not None and self.boundary_timestep is not None: + if timestep.item() <= self.boundary_timestep: + self.train_fake_score_transformer_2 = True + return self.fake_score_transformer_2 # Low noise expert + else: + self.train_fake_score_transformer_2 = False + return self.fake_score_transformer # High noise expert + else: + self.train_fake_score_transformer_2 = False + return self.fake_score_transformer + def _build_distill_input_kwargs( self, noise_input: torch.Tensor, timestep: torch.Tensor, text_dict: dict[str, torch.Tensor] | None, @@ -433,6 +531,7 @@ class DistillationPipeline(TrainingPipeline): training_batch = self._build_distill_input_kwargs( noisy_latent, timestep, training_batch.conditional_dict, training_batch) + pred_noise = self.transformer(**training_batch.input_kwargs).permute( 0, 2, 1, 3, 4) pred_video = pred_noise_to_pred_video( @@ -575,7 +674,8 @@ class DistillationPipeline(TrainingPipeline): training_batch = self._build_distill_input_kwargs( noisy_latent, timestep, training_batch.conditional_dict, training_batch) - fake_score_pred_noise = self.fake_score_transformer( + current_fake_score_transformer = self._get_fake_score_transformer(timestep) + fake_score_pred_noise = current_fake_score_transformer( **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) faker_score_pred_video = pred_noise_to_pred_video( @@ -589,7 +689,8 @@ class DistillationPipeline(TrainingPipeline): training_batch = self._build_distill_input_kwargs( noisy_latent, timestep, training_batch.conditional_dict, training_batch) - real_score_pred_noise_cond = self.real_score_transformer( + current_real_score_transformer = self._get_real_score_transformer(timestep) + real_score_pred_noise_cond = current_real_score_transformer( **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) pred_real_video_cond = pred_noise_to_pred_video( @@ -603,7 +704,8 @@ class DistillationPipeline(TrainingPipeline): training_batch = self._build_distill_input_kwargs( noisy_latent, timestep, training_batch.unconditional_dict, training_batch) - real_score_pred_noise_uncond = self.real_score_transformer( + # Use same transformer as conditional forward for consistency + real_score_pred_noise_uncond = current_real_score_transformer( **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) pred_real_video_uncond = pred_noise_to_pred_video( @@ -686,7 +788,8 @@ class DistillationPipeline(TrainingPipeline): noisy_generator_pred_video, fake_score_timestep, training_batch.conditional_dict, training_batch) - fake_score_pred_noise = self.fake_score_transformer( + current_fake_score_transformer = self._get_fake_score_transformer(fake_score_timestep) + fake_score_pred_noise = current_fake_score_transformer( **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) target = fake_score_noise - generator_pred_video @@ -800,6 +903,8 @@ class DistillationPipeline(TrainingPipeline): attn_metadata=batch_gen.attn_metadata_vsa): (dmd_loss / gradient_accumulation_steps).backward() total_dmd_loss += dmd_loss.detach().item() + + # Only clip gradients for the model that is currently training self._clip_model_grad_norm_(batch_gen, self.transformer) for param in self.transformer.parameters(): # check if the gradient is not None and not zero @@ -822,6 +927,7 @@ class DistillationPipeline(TrainingPipeline): training_batch.generator_loss = 0.0 self.fake_score_optimizer.zero_grad() + self.fake_score_optimizer_2.zero_grad() total_fake_score_loss = 0.0 for batch in batches: batch_fake = copy.deepcopy(batch) @@ -832,14 +938,34 @@ class DistillationPipeline(TrainingPipeline): total_fake_score_loss += fake_score_loss.detach().item() fake_score_latent_vis_dict.update( batch_fake.fake_score_latent_vis_dict) - self._clip_model_grad_norm_(batch_fake, self.fake_score_transformer) + if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None: + self._clip_model_grad_norm_(batch_fake, self.fake_score_transformer_2) + else: + self._clip_model_grad_norm_(batch_fake, self.fake_score_transformer) + + # Check gradients for fake score transformer for param in self.fake_score_transformer.parameters(): - # check if the gradient is not None and not zero - assert param.grad is not None and param.grad.abs().sum() > 0 - self.fake_score_optimizer.step() - self.fake_score_lr_scheduler.step() + if param.requires_grad: + assert param.grad is not None and param.grad.abs().sum() > 0 + + # Check gradients for fake score transformer_2 if available + if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None: + for param in self.fake_score_transformer_2.parameters(): + if param.requires_grad: + assert param.grad is not None and param.grad.abs().sum() > 0 + + if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None: + self.fake_score_optimizer_2.step() + self.fake_score_lr_scheduler_2.step() + else: + self.fake_score_optimizer.step() + self.fake_score_lr_scheduler.step() + + # Step the appropriate scheduler self.lr_scheduler.step() + self.fake_score_optimizer.zero_grad(set_to_none=True) + self.fake_score_optimizer_2.zero_grad(set_to_none=True) avg_fake_score_loss = torch.tensor(total_fake_score_loss / gradient_accumulation_steps, device=self.device) @@ -886,14 +1012,38 @@ class DistillationPipeline(TrainingPipeline): logger.info(" Max gradient norm: %s", self.training_args.max_grad_norm) logger.info( - " Real score transformer parameters: %s B", + " Real score transformer (high noise expert) parameters: %s B", sum(p.numel() for p in self.real_score_transformer.parameters()) / 1e9) + + if self.real_score_transformer_2 is not None: + logger.info( + " Real score transformer_2 (low noise expert) parameters: %s B", + sum(p.numel() + for p in self.real_score_transformer_2.parameters()) / 1e9) + logger.info(" Real score MoE enabled with boundary_timestep: %s", + self.boundary_timestep) logger.info( - " Fake score transformer parameters: %s B", + " Fake score transformer (high noise expert) parameters: %s B", sum(p.numel() for p in self.fake_score_transformer.parameters()) / 1e9) + + if self.fake_score_transformer_2 is not None: + logger.info( + " Fake score transformer_2 (low noise expert) parameters: %s B", + sum(p.numel() + for p in self.fake_score_transformer_2.parameters()) / 1e9) + logger.info(" Fake score MoE enabled with boundary_timestep: %s", + self.boundary_timestep) + + if self.generator_ema is not None: + logger.info(" Generator EMA enabled with decay: %s", + self.training_args.ema_decay) + logger.info(" Generator EMA start step: %s", + self.training_args.ema_start_step) + else: + logger.info(" Generator EMA disabled") if self.generator_ema is not None: logger.info(" Generator EMA enabled with decay: %s", @@ -1395,4 +1545,4 @@ class DistillationPipeline(TrainingPipeline): self.training_args.max_train_steps) if get_sp_group(): - cleanup_dist_env_and_memory() + cleanup_dist_env_and_memory() \ No newline at end of file diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 0cdfda56a..878f4335a 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -184,10 +184,13 @@ class SelfForcingDistillationPipeline(DistillationPipeline): noisy_latent, timestep, training_batch.conditional_dict, training_batch) + # Use the appropriate transformer + current_model = self.transformer_2 if self.train_transformer_2 else self.transformer + # Step 5: Forward pass with KV cache if available - if hasattr(self.transformer, '_forward_inference'): + if hasattr(current_model, '_forward_inference'): # Use causal inference forward with KV cache - pred_noise = self.transformer( + pred_noise = current_model( hidden_states=training_batch.input_kwargs['hidden_states'], encoder_hidden_states=training_batch. input_kwargs['encoder_hidden_states'], @@ -200,7 +203,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): cache_start=0).permute(0, 2, 1, 3, 4) else: # Fallback to regular forward - pred_noise = self.transformer( + pred_noise = current_model( **training_batch.input_kwargs).permute(0, 2, 1, 3, 4) # Step 6: Convert noise prediction to video prediction @@ -289,6 +292,24 @@ class SelfForcingDistillationPipeline(DistillationPipeline): device=noise.device, dtype=noise.dtype) + # Get device from model parameters (FSDP models don't have .device attribute) + def get_model_device(model): + if model is None: + return "None" + try: + return next(model.parameters()).device + except (StopIteration, AttributeError): + return "Unknown" + + # offload real and fake score transformers + + + logger.info(f"self.transformer device: {get_model_device(self.transformer)}") + logger.info(f"self.transformer_2 device: {get_model_device(self.transformer_2)}") + logger.info(f"self.real_score_transformer device: {get_model_device(self.real_score_transformer)}") + logger.info(f"self.real_score_transformer_2 device: {get_model_device(self.real_score_transformer_2)}") + logger.info(f"self.fake_score_transformer device: {get_model_device(self.fake_score_transformer)}") + logger.info(f"self.fake_score_transformer_2 device: {get_model_device(self.fake_score_transformer_2)}") # Step 1: Initialize KV cache to all zeros self.kv_cache1, self.crossattn_cache = self._initialize_simulation_caches( batch_size, dtype, self.device) @@ -309,8 +330,9 @@ class SelfForcingDistillationPipeline(DistillationPipeline): training_batch_temp = self._build_distill_input_kwargs( initial_latent, timestep * 0, training_batch.conditional_dict, training_batch) - - self.transformer( + # we process the image latent with self.transformer_2 (low-noise expert) + current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer + current_model( hidden_states=training_batch_temp. input_kwargs['hidden_states'], encoder_hidden_states=training_batch_temp. @@ -349,6 +371,17 @@ class SelfForcingDistillationPipeline(DistillationPipeline): timestep = torch.ones([batch_size, current_num_frames], device=noise.device, dtype=torch.int64) * current_timestep + + if self.boundary_timestep is not None and current_timestep <= self.boundary_timestep and self.transformer_2 is not None: + current_model = self.transformer_2 + self._enable_training(self.transformer_2, self.optimizer_2) + self._disable_training(self.transformer, self.optimizer) + + else: + current_model = self.transformer + self._enable_training(self.transformer, self.optimizer) + if self.boundary_timestep is not None and self.transformer_2 is not None: + self._disable_training(self.transformer_2, self.optimizer_2) if not exit_flag: with torch.no_grad(): @@ -357,7 +390,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): noisy_input, timestep, training_batch.conditional_dict, training_batch) - pred_flow = self.transformer( + pred_flow = current_model( hidden_states=training_batch_temp. input_kwargs['hidden_states'], encoder_hidden_states=training_batch_temp. @@ -397,7 +430,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): noisy_input, timestep, training_batch.conditional_dict, training_batch) - pred_flow = self.transformer( + pred_flow = current_model( hidden_states=training_batch_temp. input_kwargs['hidden_states'], encoder_hidden_states=training_batch_temp. @@ -417,7 +450,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): noisy_input, timestep, training_batch.conditional_dict, training_batch) - pred_flow = self.transformer( + pred_flow = current_model( hidden_states=training_batch_temp. input_kwargs['hidden_states'], encoder_hidden_states=training_batch_temp. @@ -457,7 +490,9 @@ class SelfForcingDistillationPipeline(DistillationPipeline): denoised_pred, context_timestep, training_batch.conditional_dict, training_batch) - self.transformer( + # context_timestep is 0 so we use transformer_2 + current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer + current_model( hidden_states=training_batch_temp. input_kwargs['hidden_states'], encoder_hidden_states=training_batch_temp. @@ -781,7 +816,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline): training_batch.fake_score_latent_vis_dict = {} if train_generator: + logger.debug(f"Training generator at step {self.current_trainstep}") self.optimizer.zero_grad() + if self.transformer_2 is not None: + self.optimizer_2.zero_grad() total_generator_loss = 0.0 generator_log_dict = {} @@ -813,12 +851,23 @@ class SelfForcingDistillationPipeline(DistillationPipeline): training_batch.dmd_latent_vis_dict.update( batch_gen.dmd_latent_vis_dict) - self._clip_model_grad_norm_(batch_gen, self.transformer) - self.optimizer.step() - self.lr_scheduler.step() + # Only clip gradients and step optimizer for the model that is currently training + if hasattr(self, 'train_transformer_2') and self.train_transformer_2 and self.transformer_2 is not None: + self._clip_model_grad_norm_(batch_gen, self.transformer_2) + self.optimizer_2.step() + self.lr_scheduler_2.step() + else: + self._clip_model_grad_norm_(batch_gen, self.transformer) + self.optimizer.step() + self.lr_scheduler.step() if self.generator_ema is not None: - self.generator_ema.update(self.transformer) + if hasattr(self, 'train_transformer_2') and self.train_transformer_2 and self.transformer_2 is not None: + # Note: EMA currently only supports the main transformer + # Could be extended to support transformer_2 in the future + pass + else: + self.generator_ema.update(self.transformer) avg_generator_loss = torch.tensor(total_generator_loss / gradient_accumulation_steps, @@ -830,6 +879,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): else: training_batch.generator_loss = 0.0 + logger.debug(f"Training critic at step {self.current_trainstep}") self.fake_score_optimizer.zero_grad() total_critic_loss = 0.0 critic_log_dict = {} @@ -862,9 +912,14 @@ class SelfForcingDistillationPipeline(DistillationPipeline): training_batch.fake_score_latent_vis_dict.update( batch_critic.fake_score_latent_vis_dict) - self._clip_model_grad_norm_(batch_critic, self.fake_score_transformer) - self.fake_score_optimizer.step() - self.fake_score_lr_scheduler.step() + if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None: + self._clip_model_grad_norm_(batch_critic, self.fake_score_transformer_2) + self.fake_score_optimizer_2.step() + self.fake_score_lr_scheduler_2.step() + else: + self._clip_model_grad_norm_(batch_critic, self.fake_score_transformer) + self.fake_score_optimizer.step() + self.fake_score_lr_scheduler.step() avg_critic_loss = torch.tensor(total_critic_loss / gradient_accumulation_steps, diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index c963765f8..ef855766f 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +import gc import dataclasses import math import os @@ -22,7 +23,7 @@ from tqdm.auto import tqdm import fastvideo.envs as envs from fastvideo.attention.backends.video_sparse_attn import ( VideoSparseAttentionMetadataBuilder) -from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder +# from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder from fastvideo.configs.sample import SamplingParam from fastvideo.dataset import build_parquet_map_style_dataloader from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v @@ -39,20 +40,26 @@ from fastvideo.training.activation_checkpoint import ( apply_activation_checkpointing) from fastvideo.training.training_utils import ( clip_grad_norm_while_handling_failing_dtensor_cases, - compute_density_for_timestep_sampling, count_trainable, get_scheduler, - get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint, + compute_density_for_timestep_sampling, get_scheduler, get_sigmas, + load_checkpoint, normalize_dit_input, save_checkpoint, shard_latents_across_sp) -from fastvideo.utils import (is_vmoba_available, is_vsa_available, +# from fastvideo.utils import (is_vmoba_available, is_vsa_available, +# set_random_seed, shallow_asdict) +from fastvideo.utils import (is_vsa_available, set_random_seed, shallow_asdict) import wandb # isort: skip vsa_available = is_vsa_available() -vmoba_available = is_vmoba_available() +# vmoba_available = is_vmoba_available() logger = init_logger(__name__) +def _get_trainable_params(model: torch.nn.Module) -> int: + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + class TrainingPipeline(LoRAPipeline, ABC): """ A pipeline for training a model. All training pipelines should inherit from this class. @@ -63,6 +70,7 @@ class TrainingPipeline(LoRAPipeline, ABC): train_dataloader: StatefulDataLoader train_loader_iter: Iterator[dict[str, Any]] current_epoch: int = 0 + train_transformer_2: bool = False def __init__( self, @@ -98,6 +106,7 @@ class TrainingPipeline(LoRAPipeline, ABC): self.sp_world_size = self.sp_group.world_size self.local_rank = world_group.local_rank self.transformer = self.get_module("transformer") + self.transformer_2 = self.get_module("transformer_2", None) self.seed = training_args.seed self.set_schemas() @@ -110,17 +119,25 @@ class TrainingPipeline(LoRAPipeline, ABC): self.transformer, checkpointing_type=training_args. enable_gradient_checkpointing_type) + if self.transformer_2 is not None: + self.transformer_2 = apply_activation_checkpointing( + self.transformer_2, + checkpointing_type=training_args. + enable_gradient_checkpointing_type) noise_scheduler = self.modules["scheduler"] - # Set grads for proper modules based on the training mode (Distill, LoRA, etc.) self.set_trainable() params_to_optimize = self.transformer.parameters() params_to_optimize = list( filter(lambda p: p.requires_grad, params_to_optimize)) + # Parse betas from string format "beta1,beta2" + betas_str = training_args.betas + betas = tuple(float(x.strip()) for x in betas_str.split(",")) + self.optimizer = torch.optim.AdamW( params_to_optimize, lr=training_args.learning_rate, - betas=(0.9, 0.999), + betas=betas, weight_decay=training_args.weight_decay, eps=1e-8, ) @@ -138,6 +155,30 @@ class TrainingPipeline(LoRAPipeline, ABC): min_lr_ratio=training_args.min_lr_ratio, last_epoch=self.init_steps - 1, ) + if self.transformer_2 is not None: + # Ensure transformer_2 has trainable parameters before creating optimizer + self.transformer_2.train() + self.transformer_2.requires_grad_(True) + params_to_optimize_2 = self.transformer_2.parameters() + params_to_optimize_2 = list( + filter(lambda p: p.requires_grad, params_to_optimize_2)) + self.optimizer_2 = torch.optim.AdamW( + params_to_optimize_2, + lr=training_args.learning_rate, + betas=(0.9, 0.999), + weight_decay=training_args.weight_decay, + eps=1e-8, + ) + self.lr_scheduler_2 = get_scheduler( + training_args.lr_scheduler, + optimizer=self.optimizer_2, + num_warmup_steps=training_args.lr_warmup_steps, + num_training_steps=training_args.max_train_steps, + num_cycles=training_args.lr_num_cycles, + power=training_args.lr_power, + min_lr_ratio=training_args.min_lr_ratio, + last_epoch=self.init_steps - 1, + ) self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader( training_args.data_path, @@ -152,6 +193,16 @@ class TrainingPipeline(LoRAPipeline, ABC): seed=self.seed) self.noise_scheduler = noise_scheduler + if self.training_args.boundary_ratio is not None: + self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps + else: + self.boundary_timestep = None + + + logger.info(f"train_dataloader length: {len(self.train_dataloader)}") + logger.info(f"train_sp_batch_size: {training_args.train_sp_batch_size}") + logger.info(f"gradient_accumulation_steps: {training_args.gradient_accumulation_steps}") + logger.info(f"sp_size: {training_args.sp_size}") self.num_update_steps_per_epoch = math.ceil( len(self.train_dataloader) / @@ -178,9 +229,25 @@ class TrainingPipeline(LoRAPipeline, ABC): def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch: self.transformer.train() self.optimizer.zero_grad() + if self.transformer_2 is not None: + self.transformer_2.train() + self.optimizer_2.zero_grad() training_batch.total_loss = 0.0 return training_batch + def _enable_training(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None: + """Enable training mode and gradients for the specified model.""" + for param in model.parameters(): + param.requires_grad = True + model.train() + optimizer.zero_grad() + + def _disable_training(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None: + """Disable training mode and gradients for the specified model.""" + for param in model.parameters(): + param.requires_grad = False + optimizer.zero_grad(set_to_none=True) + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: batch = next(self.train_loader_iter, None) # type: ignore if batch is None: @@ -224,17 +291,17 @@ class TrainingPipeline(LoRAPipeline, ABC): generator=self.noise_gen_cuda, device=latents.device, dtype=latents.dtype) - u = compute_density_for_timestep_sampling( - weighting_scheme=self.training_args.weighting_scheme, - batch_size=batch_size, - generator=self.noise_random_generator, - logit_mean=self.training_args.logit_mean, - logit_std=self.training_args.logit_std, - mode_scale=self.training_args.mode_scale, - ) - indices = (u * self.noise_scheduler.config.num_train_timesteps).long() - timesteps = self.noise_scheduler.timesteps[indices].to( - device=latents.device) + timesteps = self._sample_timesteps(batch_size, latents.device) + + # Enable training for the model that will be trained next and disable the other + if self.train_transformer_2: + self._enable_training(self.transformer_2, self.optimizer_2) + self._disable_training(self.transformer, self.optimizer) + else: + self._enable_training(self.transformer, self.optimizer) + if self.transformer_2 is not None: + self._disable_training(self.transformer_2, self.optimizer_2) + if self.training_args.sp_size > 1: # Make sure that the timesteps are the same across all sp processes. sp_group = get_sp_group() @@ -257,6 +324,38 @@ class TrainingPipeline(LoRAPipeline, ABC): return training_batch + def _sample_timesteps(self, batch_size, device): + # Determine which model to train based on the boundary timestep + if (self.transformer_2 is not None and self.boundary_timestep is not None and + torch.rand(1, generator=self.noise_random_generator).item() <= self.training_args.boundary_ratio): + self.train_transformer_2 = True + else: + self.train_transformer_2 = False + + # Broadcast the decision to all processes + decision = torch.tensor(1.0 if self.train_transformer_2 else 0.0, device=self.device) + dist.broadcast(decision, src=0) + self.train_transformer_2 = decision.item() == 1.0 + + # Sample u from the appropriate range + u = compute_density_for_timestep_sampling( + weighting_scheme=self.training_args.weighting_scheme, + batch_size=batch_size, + generator=self.noise_random_generator, + logit_mean=self.training_args.logit_mean, + logit_std=self.training_args.logit_std, + mode_scale=self.training_args.mode_scale, + ) + + boundary_ratio = self.training_args.boundary_ratio + if self.train_transformer_2: + u = (1 - boundary_ratio) + u * boundary_ratio # min: 1 - boundary_ratio, max: 1 + else: + u = u * (1 - boundary_ratio) # min: 0, max: 1 - boundary_ratio + + indices = (u * self.noise_scheduler.config.num_train_timesteps).long() + return self.noise_scheduler.timesteps[indices].to(device=device) + def _build_attention_metadata( self, training_batch: TrainingBatch) -> TrainingBatch: latents_shape = training_batch.raw_latent_shape @@ -272,20 +371,6 @@ class TrainingPipeline(LoRAPipeline, ABC): patch_size=patch_size, VSA_sparsity=current_vsa_sparsity, device=get_local_torch_device()) - elif vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN": - moba_params = self.training_args.moba_config.copy() - moba_params.update({ - "current_timestep": - training_batch.timesteps, - "raw_latent_shape": - training_batch.raw_latent_shape[2:5], - "patch_size": - self.training_args.pipeline_config.dit_config.patch_size, - "device": - get_local_torch_device(), - }) - training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder( - ).build(**moba_params) else: training_batch.attn_metadata = None @@ -310,7 +395,7 @@ class TrainingPipeline(LoRAPipeline, ABC): def _transformer_forward_and_compute_loss( self, training_batch: TrainingBatch) -> TrainingBatch: - if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN": + if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": assert training_batch.attn_metadata is not None else: assert training_batch.attn_metadata is None @@ -321,11 +406,12 @@ class TrainingPipeline(LoRAPipeline, ABC): # [1000.0], # device=training_batch.noisy_model_input.device, # dtype=torch.bfloat16) + current_model = self.transformer_2 if self.train_transformer_2 else self.transformer with set_forward_context( current_timestep=training_batch.current_timestep, attn_metadata=training_batch.attn_metadata): - model_pred = self.transformer(**input_kwargs) + model_pred = current_model(**input_kwargs) if self.training_args.precondition_outputs: assert training_batch.sigmas is not None model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas @@ -356,7 +442,12 @@ class TrainingPipeline(LoRAPipeline, ABC): # the following: # grad_norm = transformer.clip_grad_norm_(max_grad_norm) if max_grad_norm is not None: - model_parts = [self.transformer] + # Only clip gradients for the model that is currently training + if self.train_transformer_2 and self.transformer_2 is not None: + model_parts = [self.transformer_2] + else: + model_parts = [self.transformer] + grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases( [p for m in model_parts for p in m.parameters()], max_grad_norm, @@ -401,9 +492,14 @@ class TrainingPipeline(LoRAPipeline, ABC): training_batch = self._clip_grad_norm(training_batch) - self.optimizer.step() - self.lr_scheduler.step() - + # Only step the optimizer and scheduler for the model that is currently training + if self.train_transformer_2 and self.transformer_2 is not None: + self.optimizer_2.step() + self.lr_scheduler_2.step() + else: + self.optimizer.step() + self.lr_scheduler.step() + training_batch.total_loss = training_batch.total_loss training_batch.grad_norm = training_batch.grad_norm return training_batch @@ -431,10 +527,15 @@ class TrainingPipeline(LoRAPipeline, ABC): local_main_process_only=False) if not self.post_init_called: self.post_init() - num_trainable_params = count_trainable(self.transformer) + num_trainable_params = _get_trainable_params(self.transformer) logger.info("Starting training with %s B trainable parameters", round(num_trainable_params / 1e9, 3)) + if getattr(self, "transformer_2", None) is not None: + num_trainable_params = _get_trainable_params(self.transformer_2) + logger.info("Transformer 2: Starting training with %s B trainable parameters", + round(num_trainable_params / 1e9, 3)) + # Set random seeds for deterministic training self.noise_random_generator = torch.Generator(device="cpu").manual_seed( self.seed) @@ -455,7 +556,7 @@ class TrainingPipeline(LoRAPipeline, ABC): self._log_training_info() - self._log_validation(self.transformer, self.training_args, + self._log_validation(self.training_args, self.init_steps) # Train! @@ -476,9 +577,6 @@ class TrainingPipeline(LoRAPipeline, ABC): current_decay_times = min(step // vsa_decay_interval_steps, vsa_sparsity // vsa_decay_rate) current_vsa_sparsity = current_decay_times * vsa_decay_rate - elif vmoba_available: - # TODO: add vmoba sparsity scheduling here - current_vsa_sparsity = 0.0 else: current_vsa_sparsity = 0.0 @@ -527,7 +625,7 @@ class TrainingPipeline(LoRAPipeline, ABC): self._log_validation(self.transformer, self.training_args, step) gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 trainable_params = round( - count_trainable(self.transformer) / 1e9, 3) + _get_trainable_params(self.transformer) / 1e9, 3) logger.info( "GPU memory usage after validation: %s MB, trainable params: %sB", gpu_memory_usage, trainable_params) @@ -563,7 +661,7 @@ class TrainingPipeline(LoRAPipeline, ABC): logger.info(" Total optimization steps = %s", self.training_args.max_train_steps) logger.info(" Total training parameters per FSDP shard = %s B", - round(count_trainable(self.transformer) / 1e9, 3)) + round(_get_trainable_params(self.transformer) / 1e9, 3)) # print dtype logger.info(" Master weight dtype: %s", self.transformer.parameters().__next__().dtype) @@ -605,12 +703,12 @@ class TrainingPipeline(LoRAPipeline, ABC): return batch @torch.no_grad() - def _log_validation(self, transformer, training_args, global_step) -> None: + def _log_validation(self, training_args, global_step) -> None: """ Generate a validation video and log it to wandb to check the quality during training. """ training_args.inference_mode = True - training_args.dit_cpu_offload = True + training_args.dit_cpu_offload = False if not training_args.log_validation: return if self.validation_pipeline is None: @@ -631,7 +729,10 @@ class TrainingPipeline(LoRAPipeline, ABC): validation_dataloader = DataLoader(validation_dataset, batch_size=None, num_workers=0) - transformer.eval() + + self.transformer.eval() + if getattr(self, "transformer_2", None) is not None: + self.transformer_2.eval() validation_steps = training_args.validation_sampling_steps.split(",") validation_steps = [int(step) for step in validation_steps] @@ -723,7 +824,9 @@ class TrainingPipeline(LoRAPipeline, ABC): # Re-enable gradients for training training_args.inference_mode = False - transformer.train() + self.transformer.train() + if getattr(self, "transformer_2", None) is not None: + self.transformer_2.train() def visualize_intermediate_latents(self, training_batch: TrainingBatch, training_args: TrainingArgs, step: int): diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 9e34bbcbc..dbcad0056 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -1323,7 +1323,10 @@ class EMA_FSDP: ema.update(model) ema.state_dict() # on rank 0 """ +<<<<<<< HEAD +======= +>>>>>>> 2b2f672c (inference works after changes added) def __init__(self, module, decay: float = 0.999, mode: str = "local_shard"): self.decay = float(decay) self.mode = mode @@ -1417,7 +1420,6 @@ class EMA_FSDP: p.data.copy_(w.to(dtype=p.dtype, device=p.device)) class _ApplyEMACtx: - def __init__(self, ema: "EMA_FSDP", module): self.ema = ema self.module = module diff --git a/wan2.2_train.sh b/wan2.2_train.sh new file mode 100644 index 000000000..a6510da0f --- /dev/null +++ b/wan2.2_train.sh @@ -0,0 +1,57 @@ +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export WANDB_API_KEY='2f25ad37933894dbf0966c838c0b8494987f9f2f' +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA +export TRITON_CACHE_DIR=/tmp/triton_cache +# DATA_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/ +DATA_DIR=/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol_processed_t2v/combined_parquet_dataset +# VALIDATION_DIR=/mnt/sharefs/users/hao.zhang/Vchitect-2M/mixkit/validation_8.json +VALIDATION_DIR=/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol-single_processed_t2v/validation.json +NUM_GPUS=8 +export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN +# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] +CHECKPOINT_PATH="outputs_train_test/wan_finetune/checkpoint-10" +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \ + fastvideo/training/wan_training_pipeline.py \ + --model_path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --inference_mode False\ + --pretrained_model_name_or_path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --cache_dir "/home/ray/.cache" \ + --data_path "$DATA_DIR" \ + --validation_dataset_file "$VALIDATION_DIR" \ + --train_batch_size 1 \ + --num_latent_t 16\ + --sp_size 4 \ + --tp_size 1 \ + --num_gpus $NUM_GPUS \ + --hsdp_replicate_dim 1 \ + --hsdp-shard-dim 8 \ + --train_sp_batch_size 1 \ + --dataloader_num_workers 4 \ + --gradient_accumulation_steps 1 \ + --max_train_steps 30000 \ + --learning_rate 1e-5 \ + --mixed_precision "bf16" \ + --checkpointing_steps 1000 \ + --validation_steps 30 \ + --validation_sampling_steps "40" \ + --log_validation True \ + --checkpoints_total_limit 3 \ + --ema_start_step 0 \ + --training_cfg_rate 0.1 \ + --seed 1024 \ + --output_dir "outputs_train_test/wan_finetune_v1" \ + --tracker_project_name VSA_finetune \ + --num_height 448 \ + --num_width 832 \ + --num_frames 61 \ + --flow_shift 5 \ + --validation_guidance_scale "5.0" \ + --num_euler_timesteps 50 \ + --master_weight_type "fp32" \ + --dit_precision "fp32" \ + --weight_decay 0.01 \ + --max_grad_norm 1.0 \ No newline at end of file