From eddc7efc6486de1ad3805a495f873ccd0546fca8 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 --- examples/inference/basic/basic.py | 14 +- fastvideo/layers/layernorm.py | 10 +- fastvideo/models/dits/causal_wanvideo.py | 98 +++--- fastvideo/models/dits/wanvideo.py | 153 ++++----- .../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 | 3 +- fastvideo/training/training_pipeline.py | 69 +++-- fastvideo/training/training_utils.py | 2 - 10 files changed, 760 insertions(+), 158 deletions(-) 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 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/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 0656b9740..64366a411 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( @@ -378,7 +375,7 @@ class CausalWanTransformer3DModel(BaseDiT): # Causal-specific self.block_mask = None - self.num_frame_per_block = 1 + self.num_frame_per_block = 3 self.independent_first_frame = False self.__post_init__() @@ -490,12 +487,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) @@ -542,14 +543,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, @@ -590,8 +586,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: @@ -604,8 +600,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) @@ -640,14 +640,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, @@ -658,3 +653,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..3262a9a1e 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,20 @@ 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 = 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) 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 +367,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 +397,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 +429,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 +449,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 +469,22 @@ 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() + 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 = (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 +509,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 +516,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 @@ -598,8 +582,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 +642,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 +657,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: @@ -728,14 +715,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 +835,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/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..fedc8ee62 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 @@ -1395,4 +1396,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/training_pipeline.py b/fastvideo/training/training_pipeline.py index 4e0b4b917..ce0420582 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -22,7 +22,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 +39,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. @@ -112,15 +118,18 @@ class TrainingPipeline(LoRAPipeline, ABC): 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, ) @@ -272,20 +281,20 @@ 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) + # 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 +319,8 @@ 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" 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 @@ -431,7 +441,7 @@ 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)) @@ -476,9 +486,9 @@ 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 + # elif vmoba_available: + # # TODO: add vmoba sparsity scheduling here + # pass else: current_vsa_sparsity = 0.0 @@ -523,7 +533,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) @@ -559,7 +569,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) @@ -627,6 +637,7 @@ class TrainingPipeline(LoRAPipeline, ABC): validation_dataloader = DataLoader(validation_dataset, batch_size=None, num_workers=0) + transformer.eval() validation_steps = training_args.validation_sampling_steps.split(",") @@ -719,4 +730,4 @@ class TrainingPipeline(LoRAPipeline, ABC): # Re-enable gradients for training training_args.inference_mode = False - transformer.train() + transformer.train() \ No newline at end of file diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 9e34bbcbc..a58850fb9 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -1323,7 +1323,6 @@ class EMA_FSDP: ema.update(model) ema.state_dict() # on rank 0 """ - def __init__(self, module, decay: float = 0.999, mode: str = "local_shard"): self.decay = float(decay) self.mode = mode @@ -1417,7 +1416,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