update inference sp code / can run / still has bug / don't output normal mp4

fastvideo/sample/monkey_patch/sp_attention.py

@@ -18,7 +18,7 @@ from types import MethodType
 
 from diffusers.models.attention_processor import Attention
 from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-from fastvideo.utils.communications import all_to_all_SBH
+from fastvideo.utils.communications import all_to_all_BHSD, all_gather_BHSD
 
 class NewMochiAttnProcessor2_0:
     """Attention processor used in Mochi."""
@@ -31,10 +31,12 @@ class NewMochiAttnProcessor2_0:
         attention_mask: Optional[torch.Tensor] = None,
         image_rotary_emb: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
+        # [b, s, h * d]
         query = attn.to_q(hidden_states)
         key = attn.to_k(hidden_states)
         value = attn.to_v(hidden_states)
 
+        # [b, s, h=24, d=128]
         query = query.unflatten(2, (attn.heads, -1))
         key = key.unflatten(2, (attn.heads, -1))
         value = value.unflatten(2, (attn.heads, -1))
@@ -44,10 +46,12 @@ class NewMochiAttnProcessor2_0:
         if attn.norm_k is not None:
             key = attn.norm_k(key)
 
+        # [b, 256, h * d]
         encoder_query = attn.add_q_proj(encoder_hidden_states)
         encoder_key = attn.add_k_proj(encoder_hidden_states)
         encoder_value = attn.add_v_proj(encoder_hidden_states)
 
+        # [b, 256, h=24, d=128]
         encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
         encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
         encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
@@ -77,47 +81,58 @@ class NewMochiAttnProcessor2_0:
             encoder_key.transpose(1, 2),
             encoder_value.transpose(1, 2),
         )
+        # [b, h, s, d]
 
         sequence_length = query.size(2)
+        total_sequence_length = sequence_length * nccl_info.world_size
         encoder_sequence_length = encoder_query.size(2)
-
+        print("sequence_length: ", sequence_length)
+        print("encoder_sequence_length: ", encoder_sequence_length)
+        
+        # if sp query slice, encoder_query full
+        # first all-to-all query -> full query, then concat encoder_query
+        if get_sequence_parallel_state():
+            query = all_to_all_BHSD(query, scatter_dim=1, gather_dim=2).contiguous()
+            key = all_to_all_BHSD(key, scatter_dim=1, gather_dim=2).contiguous()
+            value = all_to_all_BHSD(value, scatter_dim=1, gather_dim=2).contiguous()
+            
+        if get_sequence_parallel_state():
+            def shrink_head(encoder_state):
+                world_size, rank = nccl_info.world_size, nccl_info.rank
+                return rearrange(encoder_state, 'b (n h) s d -> b n h s d', n=world_size,
+                                    h=encoder_state.shape[1] // world_size).contiguous()[:, rank, ...]
+            encoder_query = shrink_head(encoder_query)
+            encoder_key = shrink_head(encoder_key)
+            encoder_value = shrink_head(encoder_value)
+            
+        # Hint: please check encoder_query.shape
+        print("query.shape: ", query.shape)
+        print("encoder_query.shape: ", encoder_query.shape)
         query = torch.cat([query, encoder_query], dim=2)
         key = torch.cat([key, encoder_key], dim=2)
         value = torch.cat([value, encoder_value], dim=2)
         
-        
+        print("hidden_states.shape: ", hidden_states.shape)
+                
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+            
+        print("after attn before sp.shape", hidden_states.shape)
+        hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
+            (total_sequence_length, encoder_sequence_length), dim=2
+        )
+        print("after attn before sp.shape", hidden_states.shape)
         if get_sequence_parallel_state():
-            batch_size, _, _, head_dim = hidden_states.shape
-            h_size = attn.heads * head_dim
-            sp_size = nccl_info.world_size
-            h_size_sp = h_size // sp_size
-            query = all_to_all_SBH(query, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp)
-            key = all_to_all_SBH(key, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp)
-            value = all_to_all_SBH(value, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp)
-                
-            if False and self.use_rope:
-                # require the shape of (batch_size x nheads x ntokens x dim)
-                pos_thw = self.position_getter(batch_size, t=frame * sp_size, h=height, w=width, device=query.device)
-                query = self.rope(query, pos_thw)
-                key = self.rope(key, pos_thw)
-                
-                
-            hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
-            
-            hidden_states = rearrange(hidden_states, 'b h s d -> s h b d').contiguous()
-
-            # [s * b, h // sp, d] -> [s // sp * b, h, d] -> [s // sp, b, h * d]
-            hidden_states = all_to_all_SBH(hidden_states, scatter_dim=0, gather_dim=1)
-            
-            hidden_states = rearrange(hidden_states, 's h b d -> s b (h d)').contiguous()
-        else:
-            hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+            hidden_states = all_to_all_BHSD(hidden_states, scatter_dim=2, gather_dim=1).contiguous()
+            encoder_hidden_states = all_gather_BHSD(encoder_hidden_states, dim=1).contiguous()
+        print("after attn after sp.shape", hidden_states.shape)
+        
         hidden_states = hidden_states.transpose(1, 2).flatten(2, 3)
         hidden_states = hidden_states.to(query.dtype)
+        encoder_hidden_states = encoder_hidden_states.transpose(1, 2).flatten(2, 3)
+        encoder_hidden_states = encoder_hidden_states.to(query.dtype)
+        
+        print("after trans.shape", hidden_states.shape, encoder_hidden_states.shape)
 
-        hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
-            (sequence_length, encoder_sequence_length), dim=1
-        )
 
         # linear proj
         hidden_states = attn.to_out[0](hidden_states)

fastvideo/sample/monkey_patch/sp_rope.py

@@ -0,0 +1,32 @@
+import inspect
+import math
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange, repeat
+
+from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
+
+class NewMochiRoPE(nn.Module):
+    def _get_positions(
+        self,
+        num_frames: int,
+        height: int,
+        width: int,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> torch.Tensor:
+        scale = (self.target_area / (height * width)) ** 0.5
+        t = torch.arange(num_frames, device=device, dtype=dtype)
+        if get_sequence_parallel_state():
+            rank = nccl_info.rank
+            t += rank * num_frames
+        h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype)
+        w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype)
+
+        grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
+
+        positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3)
+        return positions

fastvideo/sample/pipeline_mochi.py

@@ -32,7 +32,8 @@ from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.pipelines.mochi.pipeline_output import MochiPipelineOutput
-
+from einops import rearrange
+from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
 
 if is_torch_xla_available():
     import torch_xla.core.xla_model as xm
@@ -629,7 +630,11 @@ class MochiPipeline(DiffusionPipeline):
             generator,
             latents,
         )
-
+        world_size, rank = nccl_info.world_size, nccl_info.rank
+        if get_sequence_parallel_state():
+            latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
+            latents = latents[:, :, rank, :, :, :]
+            
         # 5. Prepare timestep
         # from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
         threshold_noise = 0.025
@@ -693,6 +698,14 @@ class MochiPipeline(DiffusionPipeline):
                 if XLA_AVAILABLE:
                     xm.mark_step()
 
+        if get_sequence_parallel_state():
+            latents_shape = list(latents.shape)
+            full_shape = [latents_shape[0] * world_size] + latents_shape[1:]
+            all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device)
+            torch.distributed.all_gather_into_tensor(all_latents, latents)
+            latents_list = list(all_latents.chunk(world_size, dim=0))
+            latents = torch.cat(latents_list, dim=2)
+
         if output_type == "latent":
             video = latents
         else:

fastvideo/sample/sample_t2v_mochi_sp.py

@@ -1,5 +1,5 @@
 import torch
-from diffusers import MochiPipeline
+from fastvideo.sample.pipeline_mochi import MochiPipeline
 import torch.distributed as dist
 
 from diffusers.utils import export_to_video, load_image, load_video
@@ -8,6 +8,7 @@ import argparse
 import os
 import diffusers
 from fastvideo.sample.monkey_patch.sp_attention import NewMochiAttnProcessor2_0
+from fastvideo.sample.monkey_patch.sp_rope import NewMochiRoPE
 
 def initialize_distributed():
     local_rank = int(os.getenv('RANK', 0))
@@ -21,9 +22,11 @@ def initialize_distributed():
 def replace_hf_ckpt_with_new_ckpt():
     initialize_distributed()
     diffusers.models.attention_processor.MochiAttnProcessor2_0.__call__ = NewMochiAttnProcessor2_0.__call__
+    diffusers.models.transformers.transformer_mochi.MochiRoPE._get_positions = NewMochiRoPE._get_positions
 
 def main(args):
     replace_hf_ckpt_with_new_ckpt()
+    print(nccl_info.world_size)
     
     torch.manual_seed(args.seed)
     weight_dtype = torch.bfloat16
@@ -32,7 +35,7 @@ def main(args):
     pipe = MochiPipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
     pipe.enable_vae_tiling()
     pipe.to(device)
-    # pipe.enable_model_cpu_offload()
+    #pipe.enable_model_cpu_offload()
 
     # Generate videos from the input prompt
     video = pipe(
@@ -46,7 +49,7 @@ def main(args):
 
     dist.barrier()
     if nccl_info.rank <= 0:
-        export_to_video(video, './outputs.mp4', fps=30)
+        export_to_video(video, args.output_path, fps=30)
 
 if __name__ == "__main__":
     # arg parse 
@@ -59,5 +62,6 @@ if __name__ == "__main__":
     parser.add_argument("--guidance_scale", type=float, default=4.5)
     parser.add_argument("--model_path", type=str, default="data/mochi")
     parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--output_path", type=str, default="./outputs.mp4")
     args = parser.parse_args()
     main(args)

fastvideo/utils/communications.py

@@ -85,10 +85,10 @@ class _AllToAll(torch.autograd.Function):
             None,
         )
 
-def all_to_all_SBH(
+def all_to_all_BHSD(
     input_: torch.Tensor,
     scatter_dim: int = 1,
-    gather_dim: int = 0,
+    gather_dim: int = 2,
 ):
     return _AllToAll.apply(input_, scatter_dim, gather_dim, _single_all_to_all)
 
@@ -99,6 +99,55 @@ def all_to_all_BSND(
 ):
     return _AllToAll.apply(input_, scatter_dim, gather_dim, _all_to_all)
 
+class _AllGather(torch.autograd.Function):
+    """All-gather communication with autograd support.
+
+    Args:
+        input_: input tensor
+        dim: dimension along which to concatenate
+    """
+
+    @staticmethod
+    def forward(ctx, input_, dim):
+        ctx.dim = dim
+        world_size = nccl_info.world_size
+        group = nccl_info.group
+        input_size = list(input_.size())
+
+        ctx.input_size = input_size[dim]
+
+        tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
+        input_ = input_.contiguous()
+        dist.all_gather(tensor_list, input_, group=group)
+
+        output = torch.cat(tensor_list, dim=dim)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        world_size = nccl_info.world_size
+        rank = nccl_info.rank
+        dim = ctx.dim
+        input_size = ctx.input_size
+
+        sizes = [input_size] * world_size
+
+        grad_input_list = torch.split(grad_output, sizes, dim=dim)
+        grad_input = grad_input_list[rank]
+
+        return grad_input, None
+
+def all_gather_BHSD(input_: torch.Tensor, dim: int = 1):
+    """Performs an all-gather operation on the input tensor along the specified dimension.
+
+    Args:
+        input_ (torch.Tensor): Input tensor of shape [B, H, S, D].
+        dim (int, optional): Dimension along which to concatenate. Defaults to 1.
+
+    Returns:
+        torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'.
+    """
+    return _AllGather.apply(input_, dim)
 
 def prepare_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask):
     def all_to_all(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask):

fastvideo/utils/parallel_states.py

@@ -28,7 +28,6 @@ def initialize_sequence_parallel_group(sequence_parallel_size):
     rank = int(os.getenv('RANK', '0'))
     world_size = int(os.getenv("WORLD_SIZE", '1'))
     assert world_size % sequence_parallel_size == 0, "world_size must be divisible by sequence_parallel_size, but got world_size: {}, sequence_parallel_size: {}".format(world_size, sequence_parallel_size)
-    # hccl
     nccl_info.world_size = sequence_parallel_size
     nccl_info.rank = rank
     num_sequence_parallel_groups: int = world_size // sequence_parallel_size

infenrece_mochi_sp.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-num_gpus=4
+num_gpus=2
 
 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     fastvideo/sample/sample_t2v_mochi_sp.py \
@@ -9,6 +9,6 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     --num_frames 163 \
     --height 480 \
     --width 848 \
-    --num_inference_steps 64 \
+    --num_inference_steps 10 \
     --guidance_scale 4.5 \
     --seed 42