From 70f340703250428557b8016cbbc4a2df40ae0615 Mon Sep 17 00:00:00 2001
From: foreverpiano <pianoqwz@qq.com>
Date: Tue, 29 Oct 2024 12:42:04 +0000
Subject: [PATCH 1/7] update optimizer

---
 fastvideo/train.py           | 12 ++-----
 fastvideo/utils/optimizer.py | 70 ++++++++++++++++++++++++++++++++++++
 pyproject.toml               |  2 +-
 t2v_debug_multi.sh           |  1 +
 t2v_debug_single.sh          |  2 ++
 5 files changed, 77 insertions(+), 10 deletions(-)
 create mode 100644 fastvideo/utils/optimizer.py

diff --git a/fastvideo/train.py b/fastvideo/train.py
index 0f6c894d7..276b333a0 100644
--- a/fastvideo/train.py
+++ b/fastvideo/train.py
@@ -256,7 +256,7 @@ def main(args):
 
     # Enable TF32 for faster training on Ampere GPUs,
     # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
-    if args.allow_tf32:
+    if args.allow_tf32 and torch.cuda.is_available():
         torch.backends.cuda.matmul.allow_tf32 = True
 
     if args.scale_lr:
@@ -271,14 +271,8 @@ def main(args):
 
     params_to_optimize = transformer.parameters()
 
-    # TODO: Other optmizer in cogvideoX-factory
-    optimizer = torch.optim.AdamW(
-        params_to_optimize,
-        lr=args.learning_rate,
-        betas=(args.adam_beta1, args.adam_beta2),
-        weight_decay=args.adam_weight_decay,
-        eps=args.adam_epsilon,
-    )
+    from fastvideo.utlis.optimizer import get_optimizer
+    optimizer = get_optimizer(params_to_optimize, args)
 
     logger.info(f"optimizer: {optimizer}")
     
diff --git a/fastvideo/utils/optimizer.py b/fastvideo/utils/optimizer.py
new file mode 100644
index 000000000..49508d0dc
--- /dev/null
+++ b/fastvideo/utils/optimizer.py
@@ -0,0 +1,70 @@
+from accelerate.logging import get_logger
+logger = get_logger(__name__)
+
+def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):
+    # Optimizer creation
+    supported_optimizers = ["adam", "adamw", "prodigy"]
+    if args.optimizer not in supported_optimizers:
+        logger.warning(
+            f"Unsupported choice of optimizer: {args.optimizer}. Supported optimizers include {supported_optimizers}. Defaulting to AdamW"
+        )
+        args.optimizer = "adamw"
+
+    if args.use_8bit_adam and not (args.optimizer.lower() not in ["adam", "adamw"]):
+        logger.warning(
+            f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was "
+            f"set to {args.optimizer.lower()}"
+        )
+
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+            )
+
+    if args.optimizer.lower() == "adamw":
+        optimizer_class = bnb.optim.AdamW8bit if args.use_8bit_adam else torch.optim.AdamW
+
+        optimizer = optimizer_class(
+            params_to_optimize,
+            betas=(args.adam_beta1, args.adam_beta2),
+            eps=args.adam_epsilon,
+            weight_decay=args.adam_weight_decay,
+        )
+    elif args.optimizer.lower() == "adam":
+        optimizer_class = bnb.optim.Adam8bit if args.use_8bit_adam else torch.optim.Adam
+
+        optimizer = optimizer_class(
+            params_to_optimize,
+            betas=(args.adam_beta1, args.adam_beta2),
+            eps=args.adam_epsilon,
+            weight_decay=args.adam_weight_decay,
+        )
+    elif args.optimizer.lower() == "prodigy":
+        try:
+            import prodigyopt
+        except ImportError:
+            raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`")
+
+        optimizer_class = prodigyopt.Prodigy
+
+        if args.learning_rate <= 0.1:
+            logger.warning(
+                "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
+            )
+
+        optimizer = optimizer_class(
+            params_to_optimize,
+            lr=args.learning_rate,
+            betas=(args.adam_beta1, args.adam_beta2),
+            beta3=args.prodigy_beta3,
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+            decouple=args.prodigy_decouple,
+            use_bias_correction=args.prodigy_use_bias_correction,
+            safeguard_warmup=args.prodigy_safeguard_warmup,
+        )
+
+    return optimizer
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index f73a84a70..8c5ad762f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,7 +21,7 @@ dependencies = [
 "timm==1.0.11", "torchdiffeq==0.2.4", "torchmetrics==1.5.1", "tqdm==4.66.5", "urllib3==2.2.0", "uvicorn==0.32.0", 
 "scikit-video==1.1.11", "imageio-ffmpeg==0.5.1", "sentencepiece==0.2.0", "beautifulsoup4==4.12.3", "ftfy==6.3.0", 
 "moviepy==1.0.3", "wandb==0.18.5", "tensorboard==2.18.0", "pydantic==2.9.2", "gradio==5.3.0", "huggingface_hub==0.26.1", "protobuf==5.28.3",
-"watch", "gpustat"
+"watch", "gpustat", "bitsandbytes", "prodigyopt"
 ]
 
 [project.optional-dependencies]
diff --git a/t2v_debug_multi.sh b/t2v_debug_multi.sh
index 602790505..8a1e48d11 100644
--- a/t2v_debug_multi.sh
+++ b/t2v_debug_multi.sh
@@ -15,6 +15,7 @@ accelerate launch \
     --learning_rate=1e-4 \
     --lr_scheduler="constant" \
     --lr_warmup_steps=0 \
+    --optimizer="AdamW" \
     --mixed_precision="bf16" \
     --checkpointing_steps=250 \
     --allow_tf32 \
diff --git a/t2v_debug_single.sh b/t2v_debug_single.sh
index bf1eda2b4..db479429d 100644
--- a/t2v_debug_single.sh
+++ b/t2v_debug_single.sh
@@ -8,12 +8,14 @@ accelerate launch \
     --data_merge_path "data/Mochi-Synthetic-Data/merge.txt" \
     --gradient_checkpointing \
     --train_batch_size=1 \
+    --num_latent_t 14 \
     --dataloader_num_workers 1 \
     --gradient_accumulation_steps=1 \
     --max_train_steps=1000000 \
     --learning_rate=1e-4 \
     --lr_scheduler="constant" \
     --lr_warmup_steps=0 \
+    --optimizer="AdamW" \
     --mixed_precision="bf16" \
     --checkpointing_steps=250 \
     --allow_tf32 \

From 37b091f2eff47f9f4a2e6f81ad99c92fde8a2367 Mon Sep 17 00:00:00 2001
From: foreverpiano <pianoqwz@qq.com>
Date: Tue, 29 Oct 2024 12:47:48 +0000
Subject: [PATCH 2/7] load

---
 fastvideo/train.py | 33 ++++++++++++++++++++++-----------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/fastvideo/train.py b/fastvideo/train.py
index 276b333a0..026b6f599 100644
--- a/fastvideo/train.py
+++ b/fastvideo/train.py
@@ -28,7 +28,7 @@ import accelerate
 import torch
 from torch.nn import functional as F
 import transformers
-from accelerate import Accelerator
+from accelerate import Accelerator, init_empty_weights, DistributedType
 from diffusers.utils.torch_utils import is_compiled_module
 from accelerate.logging import get_logger
 from accelerate.utils import DistributedType, ProjectConfiguration, set_seed
@@ -233,23 +233,34 @@ def main(args):
 
 
     def load_model_hook(models, input_dir):
-        # TODO
         if args.use_ema:
             load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), MochiTransformer3DModel)
             ema_transformer.load_state_dict(load_model.state_dict())
             ema_transformer.to(accelerator.device)
             del load_model
 
-        for i in range(len(models)):
-            # pop models so that they are not loaded again
-            model = models.pop()
+        if not accelerator.distributed_type == DistributedType.DEEPSPEED:
+            while len(models) > 0:
+                model = models.pop()
 
-            # load diffusers style into model
-            load_model = MochiTransformer3DModel.from_pretrained(input_dir, subfolder="transformer")
-            model.register_to_config(**load_model.config)
-
-            model.load_state_dict(load_model.state_dict())
-            del load_model
+                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
+                    transformer_ = unwrap_model(model)
+                else:
+                    raise ValueError(f"Unexpected save model: {unwrap_model(model).__class__}")
+        else:
+            with init_empty_weights():
+                transformer_ = MochiTransformer3DModel.from_config(
+                    args.pretrained_model_name_or_path, subfolder="transformer"
+                )
+                init_under_meta = True
+        
+        load_model = MochiTransformer3DModel.from_pretrained(os.path.join(input_dir, "transformer"))
+        transformer_.register_to_config(**load_model.config)
+        transformer_.load_state_dict(load_model.state_dict(), assign=init_under_meta)
+        del load_model
+        
+        if args.mixed_precision == "fp16":
+            cast_training_params([transformer_])
 
     accelerator.register_save_state_pre_hook(save_model_hook)
     accelerator.register_load_state_pre_hook(load_model_hook)

From 5b3bffcca9893039932fd1181a48f11bccf92d30 Mon Sep 17 00:00:00 2001
From: foreverpiano <pianoqwz@qq.com>
Date: Tue, 29 Oct 2024 14:17:54 +0000
Subject: [PATCH 3/7] small bug

---
 fastvideo/train.py           | 6 +++---
 fastvideo/utils/optimizer.py | 1 +
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/fastvideo/train.py b/fastvideo/train.py
index 026b6f599..a882a76ec 100644
--- a/fastvideo/train.py
+++ b/fastvideo/train.py
@@ -50,7 +50,7 @@ from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_func
 
 
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.24.0")
+check_min_version("0.31.0")
 logger = get_logger(__name__)
 
 
@@ -282,8 +282,8 @@ def main(args):
 
     params_to_optimize = transformer.parameters()
 
-    from fastvideo.utlis.optimizer import get_optimizer
-    optimizer = get_optimizer(params_to_optimize, args)
+    from fastvideo.utils.optimizer import get_optimizer
+    optimizer = get_optimizer(args, params_to_optimize)
 
     logger.info(f"optimizer: {optimizer}")
     
diff --git a/fastvideo/utils/optimizer.py b/fastvideo/utils/optimizer.py
index 49508d0dc..e5aad89f8 100644
--- a/fastvideo/utils/optimizer.py
+++ b/fastvideo/utils/optimizer.py
@@ -1,4 +1,5 @@
 from accelerate.logging import get_logger
+import torch
 logger = get_logger(__name__)
 
 def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):

From 4143cd1f95d9cb4f511a01b9ca96ed1f300e36d9 Mon Sep 17 00:00:00 2001
From: foreverpiano <pianoqwz@qq.com>
Date: Tue, 29 Oct 2024 15:27:52 +0000
Subject: [PATCH 4/7] random seed args

---
 infenrece_mochi.sh    |  3 ++-
 infenrece_mochi_sp.sh | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)
 create mode 100644 infenrece_mochi_sp.sh

diff --git a/infenrece_mochi.sh b/infenrece_mochi.sh
index a5880b7b4..e04ed22f9 100644
--- a/infenrece_mochi.sh
+++ b/infenrece_mochi.sh
@@ -5,4 +5,5 @@ python fastvideo/sample/sample_t2v_mochi.py \
     --height 480 \
     --width 848 \
     --num_inference_steps 64 \
-    --guidance_scale 4.5
\ No newline at end of file
+    --guidance_scale 4.5 \
+    --seed 42
\ No newline at end of file
diff --git a/infenrece_mochi_sp.sh b/infenrece_mochi_sp.sh
new file mode 100644
index 000000000..358b36d95
--- /dev/null
+++ b/infenrece_mochi_sp.sh
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+num_gpus=4
+
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+    fastvideo/sample/sample_t2v_mochi_sp.py \
+    --model_path data/mochi \
+    --prompt "A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere." \
+    --num_frames 163 \
+    --height 480 \
+    --width 848 \
+    --num_inference_steps 64 \
+    --guidance_scale 4.5 \
+    --seed 42
\ No newline at end of file

From 48a82aba3795de51eb1242578d833a0bb9933b38 Mon Sep 17 00:00:00 2001
From: foreverpiano <pianoqwz@qq.com>
Date: Tue, 29 Oct 2024 15:28:27 +0000
Subject: [PATCH 5/7] sp enable & still has bug

---
 fastvideo/sample/monkey_patch/sp_attention.py | 131 ++++++++++++++++++
 fastvideo/sample/sample_t2v_mochi_sp.py       |  63 +++++++++
 2 files changed, 194 insertions(+)
 create mode 100644 fastvideo/sample/monkey_patch/sp_attention.py
 create mode 100644 fastvideo/sample/sample_t2v_mochi_sp.py

diff --git a/fastvideo/sample/monkey_patch/sp_attention.py b/fastvideo/sample/monkey_patch/sp_attention.py
new file mode 100644
index 000000000..d66ab9f35
--- /dev/null
+++ b/fastvideo/sample/monkey_patch/sp_attention.py
@@ -0,0 +1,131 @@
+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 diffusers.image_processor import IPAdapterMaskProcessor
+from diffusers.utils import deprecate, logging
+from diffusers.utils.import_utils import is_torch_npu_available, is_xformers_available
+from diffusers.utils.torch_utils import is_torch_version, maybe_allow_in_graph
+from diffusers.models.embeddings import apply_rotary_emb
+
+import diffusers
+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
+
+class NewMochiAttnProcessor2_0:
+    """Attention processor used in Mochi."""
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        query = query.unflatten(2, (attn.heads, -1))
+        key = key.unflatten(2, (attn.heads, -1))
+        value = value.unflatten(2, (attn.heads, -1))
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        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)
+
+        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))
+
+        if attn.norm_added_q is not None:
+            encoder_query = attn.norm_added_q(encoder_query)
+        if attn.norm_added_k is not None:
+            encoder_key = attn.norm_added_k(encoder_key)
+
+        if image_rotary_emb is not None:
+
+            def apply_rotary_emb(x, freqs_cos, freqs_sin):
+                x_even = x[..., 0::2].float()
+                x_odd = x[..., 1::2].float()
+
+                cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype)
+                sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype)
+
+                return torch.stack([cos, sin], dim=-1).flatten(-2)
+
+            query = apply_rotary_emb(query, *image_rotary_emb)
+            key = apply_rotary_emb(key, *image_rotary_emb)
+
+        query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
+        encoder_query, encoder_key, encoder_value = (
+            encoder_query.transpose(1, 2),
+            encoder_key.transpose(1, 2),
+            encoder_value.transpose(1, 2),
+        )
+
+        sequence_length = query.size(2)
+        encoder_sequence_length = encoder_query.size(2)
+
+        query = torch.cat([query, encoder_query], dim=2)
+        key = torch.cat([key, encoder_key], dim=2)
+        value = torch.cat([value, encoder_value], dim=2)
+        
+        
+        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 = hidden_states.transpose(1, 2).flatten(2, 3)
+        hidden_states = hidden_states.to(query.dtype)
+
+        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)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if hasattr(attn, "to_add_out"):
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+        return hidden_states, encoder_hidden_states
+
diff --git a/fastvideo/sample/sample_t2v_mochi_sp.py b/fastvideo/sample/sample_t2v_mochi_sp.py
new file mode 100644
index 000000000..4646bad4f
--- /dev/null
+++ b/fastvideo/sample/sample_t2v_mochi_sp.py
@@ -0,0 +1,63 @@
+import torch
+from diffusers import MochiPipeline
+import torch.distributed as dist
+
+from diffusers.utils import export_to_video, load_image, load_video
+from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info
+import argparse
+import os
+import diffusers
+from fastvideo.sample.monkey_patch.sp_attention import NewMochiAttnProcessor2_0
+
+def initialize_distributed():
+    local_rank = int(os.getenv('RANK', 0))
+    world_size = int(os.getenv('WORLD_SIZE', 1))
+    print('world_size', world_size)
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
+    initialize_sequence_parallel_state(world_size)
+    
+
+def replace_hf_ckpt_with_new_ckpt():
+    initialize_distributed()
+    diffusers.models.attention_processor.MochiAttnProcessor2_0.__call__ = NewMochiAttnProcessor2_0.__call__
+
+def main(args):
+    replace_hf_ckpt_with_new_ckpt()
+    
+    torch.manual_seed(args.seed)
+    weight_dtype = torch.bfloat16
+    device = torch.cuda.current_device()
+    
+    pipe = MochiPipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
+    pipe.enable_vae_tiling()
+    pipe.to(device)
+    # pipe.enable_model_cpu_offload()
+
+    # Generate videos from the input prompt
+    video = pipe(
+        prompt=args.prompt,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+    ).frames[0]
+
+    dist.barrier()
+    if nccl_info.rank <= 0:
+        export_to_video(video, './outputs.mp4', fps=30)
+
+if __name__ == "__main__":
+    # arg parse 
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompt", type=str, default="A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.")
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=848)
+    parser.add_argument("--num_inference_steps", type=int, default=64)
+    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)
+    args = parser.parse_args()
+    main(args)

From 2f4d6caa90cc50a2685780b4812e6ca845ceee8a Mon Sep 17 00:00:00 2001
From: foreverpiano <pianoqwz@qq.com>
Date: Wed, 30 Oct 2024 12:43:53 +0000
Subject: [PATCH 6/7] rename

---
 fastvideo/sample/{mochi_pipe.py => pipeline_mochi.py} | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)
 rename fastvideo/sample/{mochi_pipe.py => pipeline_mochi.py} (99%)

diff --git a/fastvideo/sample/mochi_pipe.py b/fastvideo/sample/pipeline_mochi.py
similarity index 99%
rename from fastvideo/sample/mochi_pipe.py
rename to fastvideo/sample/pipeline_mochi.py
index f95221662..acf4c4ab3 100644
--- a/fastvideo/sample/mochi_pipe.py
+++ b/fastvideo/sample/pipeline_mochi.py
@@ -712,14 +712,12 @@ class MochiPipeline(DiffusionPipeline):
                 latents = latents / self.vae.config.scaling_factor
 
             video = self.vae.decode(latents, return_dict=False)[0]
-            # TODO: Change pil hardcode
-            video = self.video_processor.postprocess_video(video, output_type='pil')
+            video = self.video_processor.postprocess_video(video, output_type=output_type)
 
         # Offload all models
         self.maybe_free_model_hooks()
-        if output_type == "latent_and_video":
-            return video, latents, prompt_embeds, prompt_attention_mask
-        
+
         if not return_dict:
             return (video,)
-        return MochiPipelineOutput(frames=video)
\ No newline at end of file
+
+        return MochiPipelineOutput(frames=video)

From db2c68e8792d010d481537aece5af7bf3beb9c09 Mon Sep 17 00:00:00 2001
From: foreverpiano <pianoqwz@qq.com>
Date: Wed, 30 Oct 2024 16:36:50 +0000
Subject: [PATCH 7/7] update inference sp code / can run / still has bug /
 don't output normal mp4

---
 fastvideo/sample/monkey_patch/sp_attention.py | 77 +++++++++++--------
 fastvideo/sample/monkey_patch/sp_rope.py      | 32 ++++++++
 fastvideo/sample/pipeline_mochi.py            | 17 +++-
 fastvideo/sample/sample_t2v_mochi_sp.py       | 10 ++-
 fastvideo/utils/communications.py             | 53 ++++++++++++-
 fastvideo/utils/parallel_states.py            |  1 -
 infenrece_mochi_sp.sh                         |  4 +-
 7 files changed, 153 insertions(+), 41 deletions(-)
 create mode 100644 fastvideo/sample/monkey_patch/sp_rope.py

diff --git a/fastvideo/sample/monkey_patch/sp_attention.py b/fastvideo/sample/monkey_patch/sp_attention.py
index d66ab9f35..db999d594 100644
--- a/fastvideo/sample/monkey_patch/sp_attention.py
+++ b/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)
diff --git a/fastvideo/sample/monkey_patch/sp_rope.py b/fastvideo/sample/monkey_patch/sp_rope.py
new file mode 100644
index 000000000..a1f3a1566
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/fastvideo/sample/pipeline_mochi.py b/fastvideo/sample/pipeline_mochi.py
index acf4c4ab3..fd3680338 100644
--- a/fastvideo/sample/pipeline_mochi.py
+++ b/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:
diff --git a/fastvideo/sample/sample_t2v_mochi_sp.py b/fastvideo/sample/sample_t2v_mochi_sp.py
index 4646bad4f..668e26e78 100644
--- a/fastvideo/sample/sample_t2v_mochi_sp.py
+++ b/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)
diff --git a/fastvideo/utils/communications.py b/fastvideo/utils/communications.py
index 923c9467b..94745a5a0 100644
--- a/fastvideo/utils/communications.py
+++ b/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):
diff --git a/fastvideo/utils/parallel_states.py b/fastvideo/utils/parallel_states.py
index 7b7465817..2cedb444c 100644
--- a/fastvideo/utils/parallel_states.py
+++ b/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
diff --git a/infenrece_mochi_sp.sh b/infenrece_mochi_sp.sh
index 358b36d95..1bb1f3a85 100644
--- a/infenrece_mochi_sp.sh
+++ b/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
\ No newline at end of file