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

@@ -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)
 
 

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

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

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
-        
+            return hidden_states + self.previous_residual_odd

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

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()}"

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()}"

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()

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()

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