diff --git a/.gitignore b/.gitignore
index 141529aac..6b4c1d26c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -48,4 +48,6 @@ samples/
 *validation/
 data/
 outputs/
-outputs_video
\ No newline at end of file
+outputs_video
+sbatch.sh
+*.out
\ No newline at end of file
diff --git a/README.md b/README.md
index 6c26977ba..552608a06 100644
--- a/README.md
+++ b/README.md
@@ -23,9 +23,10 @@ We've prepared some debug data to facilitate development. To make sure the train
 
 ```
 python scripts/download_hf.py --repo_id=Stealths-Video/dummyVid --local_dir=data/dummyVid --repo_type=dataset
-python scripts/download_hf.py --repo_id=Stealths-Video/mochi --local_dir=data/mochi --repo_type=model
+python scripts/download_hf.py --repo_id=Stealths-Video/mochi_diffuser --local_dir=data/mochi --repo_type=model
 python scripts/download_hf.py --repo_id=Stealths-Video/Mochi-Synthetic-Data --local_dir=data/Mochi-Synthetic-Data --repo_type=dataset
 python scripts/download_hf.py --repo_id=Stealths-Video/Encoder_Overfit_Data --local_dir=data/Encoder_Overfit_Data --repo_type=dataset
+python scripts/download_hf.py --repo_id=Stealths-Video/BW-Finetune-Dataset --local_dir=data/BW-Finetune-Dataset --repo_type=dataset
 ```
 
 ## How to overfit
diff --git a/fastvideo/dataset/latent_datasets.py b/fastvideo/dataset/latent_datasets.py
index 7d8a43e0a..23c390bcc 100644
--- a/fastvideo/dataset/latent_datasets.py
+++ b/fastvideo/dataset/latent_datasets.py
@@ -22,6 +22,7 @@ class LatentDataset(Dataset):
         self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, "prompt_attention_mask")
         with open(self.json_path, 'r') as f:
             self.data_anno = json.load(f)
+        self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path'])
         self.num_latent_t = num_latent_t
         self.uncond_prompt_embed = torch.load(os.path.join(uncond_prompt_embed_mask_dir, "embed.pt"), map_location="cpu", weights_only=True)
         self.uncond_prompt_mask = torch.load(os.path.join(uncond_prompt_embed_mask_dir, "mask.pt"), map_location="cpu", weights_only=True)
diff --git a/fastvideo/fsdp_util.py b/fastvideo/fsdp_util.py
new file mode 100644
index 000000000..cb0044c36
--- /dev/null
+++ b/fastvideo/fsdp_util.py
@@ -0,0 +1,86 @@
+import torch
+import os
+import torch.distributed as dist
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    checkpoint_wrapper,
+    CheckpointImpl,
+    apply_activation_checkpointing,
+)
+
+from torch.distributed.fsdp import (
+    FullyShardedDataParallel as FSDP,
+    StateDictType,
+    FullStateDictConfig,  # general model non-sharded, non-flattened params
+    LocalStateDictConfig,  # flattened params, usable only by FSDP
+    # ShardedStateDictConfig, # un-flattened param but shards, usable by other parallel schemes.
+)
+
+from fastvideo.model.modeling_mochi import MochiTransformerBlock
+
+from functools import partial
+
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+
+from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
+import functools
+
+
+non_reentrant_wrapper = partial(
+    checkpoint_wrapper,
+    checkpoint_impl=CheckpointImpl.NO_REENTRANT,
+)
+
+check_fn = lambda submodule: isinstance(submodule, MochiTransformerBlock)
+
+
+def apply_fsdp_checkpointing(model):
+    """apply activation checkpointing to model
+    returns None as model is updated directly
+    """
+    print(f"--> applying fdsp activation checkpointing...")
+
+    apply_activation_checkpointing(
+        model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn
+    )
+
+bf16_mix = MixedPrecision(
+    param_dtype=torch.bfloat16,
+    # Gradient communication precision.
+    reduce_dtype=torch.bfloat16,
+    # Buffer precision.
+    buffer_dtype=torch.bfloat16,
+)
+
+
+
+def get_fsdp_kwargs( sharding_strategy, cpu_offload=False):
+    auto_wrap_policy = functools.partial(
+        transformer_auto_wrap_policy,
+        transformer_layer_cls={
+            MochiTransformerBlock,
+        },
+    )
+    
+    mixed_precision = bf16_mix
+    
+    if sharding_strategy == "full":
+        sharding_strategy = ShardingStrategy.FULL_SHARD
+    elif sharding_strategy == "none":
+        sharding_strategy = ShardingStrategy.NO_SHARD
+        auto_wrap_policy = None
+    elif sharding_strategy == "hybrid_zero2":
+        sharding_strategy  = ShardingStrategy._HYBRID_SHARD_ZERO2
+    
+    device_id = torch.cuda.current_device()
+    cpu_offload=torch.distributed.fsdp.CPUOffload(offload_params=True) if cpu_offload else None
+    return {
+        "auto_wrap_policy": auto_wrap_policy,
+        "mixed_precision": mixed_precision,
+        "sharding_strategy": sharding_strategy,
+        "device_id": device_id,
+        "limit_all_gathers": True,
+        "cpu_offload": cpu_offload,
+    }
+    
+    
+        
\ No newline at end of file
diff --git a/fastvideo/model/mochi_monkey_patches.py b/fastvideo/model/mochi_monkey_patches.py
deleted file mode 100644
index 65d7bd1a3..000000000
--- a/fastvideo/model/mochi_monkey_patches.py
+++ /dev/null
@@ -1,161 +0,0 @@
-from typing import Optional
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-import diffusers
-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_gather, all_to_all_4D
-
-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:
-        # [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))
-
-
-        if attn.norm_q is not None:
-            query = attn.norm_q(query)
-        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))
-        
-        
-        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:
-            freqs_cos, freqs_sin = image_rotary_emb[0], image_rotary_emb[1]
-        # shard the head dimension
-        if get_sequence_parallel_state():
-            # B, S, H, D to (S, B,) H, D
-            # batch_size, seq_len, attn_heads, head_dim 
-            query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
-            key = all_to_all_4D(key,  scatter_dim=2, gather_dim=1)
-            value = all_to_all_4D(value, scatter_dim=2, gather_dim=1)
-
-            
-            def shrink_head(encoder_state, dim):
-                local_heads = encoder_state.shape[dim] // nccl_info.sp_size
-                return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
-            encoder_query = shrink_head(encoder_query, dim=2)
-            encoder_key = shrink_head(encoder_key, dim=2)
-            encoder_value = shrink_head(encoder_value, dim=2)
-            if image_rotary_emb is not None:
-                freqs_cos = shrink_head(freqs_cos, dim=1)
-                freqs_sin = shrink_head(freqs_sin, dim=1)
-    
-
-    
-        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, freqs_cos, freqs_sin)
-            key = apply_rotary_emb(key, freqs_cos, freqs_sin)
-            
-        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),
-        )
-        # [b, h, s, d]
-        sequence_length = query.size(2)
-        encoder_sequence_length = encoder_query.size(2)
-
-        # Hint: please check 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)
-        
-        
-                
-        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
-
-        if get_sequence_parallel_state():
-            hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
-                (sequence_length, encoder_sequence_length), dim=2
-            )
-            # B, H, S, D
-            hidden_states = all_to_all_4D(hidden_states, scatter_dim=2, gather_dim=1)
-            encoder_hidden_states = all_gather(encoder_hidden_states, dim=1).contiguous()
-            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)
-        else:
-            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
-
-
-
-
-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 * nccl_info.sp_size, device=device, dtype=dtype)
-        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
-    
-def hf_mochi_add_sp_monkey_patch():
-    diffusers.models.attention_processor.MochiAttnProcessor2_0.__call__ = NewMochiAttnProcessor2_0.__call__
-    diffusers.models.transformers.transformer_mochi.MochiRoPE._get_positions = NewMochiRoPE._get_positions
\ No newline at end of file
diff --git a/fastvideo/model/modeling_mochi.py b/fastvideo/model/modeling_mochi.py
new file mode 100644
index 000000000..3c74b47e0
--- /dev/null
+++ b/fastvideo/model/modeling_mochi.py
@@ -0,0 +1,527 @@
+# Copyright 2024 The Genmo team and The HuggingFace Team.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import is_torch_version, logging
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.attention import FeedForward
+from diffusers.models.attention_processor import Attention
+from diffusers.models.embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, LuminaLayerNormContinuous, MochiRMSNormZero, RMSNorm
+from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
+from fastvideo.utils.communications import all_gather, all_to_all_4D
+import torch.nn.functional as F
+from diffusers.utils.torch_utils import is_torch_version, maybe_allow_in_graph
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class MochiAttnProcessor2_0:
+    """Attention processor used in Mochi."""
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
+
+    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:
+        # [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))
+
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        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))
+        
+        
+        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:
+            freqs_cos, freqs_sin = image_rotary_emb[0], image_rotary_emb[1]
+        # shard the head dimension
+        if get_sequence_parallel_state():
+            # B, S, H, D to (S, B,) H, D
+            # batch_size, seq_len, attn_heads, head_dim 
+            query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
+            key = all_to_all_4D(key,  scatter_dim=2, gather_dim=1)
+            value = all_to_all_4D(value, scatter_dim=2, gather_dim=1)
+
+            
+            def shrink_head(encoder_state, dim):
+                local_heads = encoder_state.shape[dim] // nccl_info.sp_size
+                return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
+            encoder_query = shrink_head(encoder_query, dim=2)
+            encoder_key = shrink_head(encoder_key, dim=2)
+            encoder_value = shrink_head(encoder_value, dim=2)
+            if image_rotary_emb is not None:
+                freqs_cos = shrink_head(freqs_cos, dim=1)
+                freqs_sin = shrink_head(freqs_sin, dim=1)
+    
+
+    
+        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, freqs_cos, freqs_sin)
+            key = apply_rotary_emb(key, freqs_cos, freqs_sin)
+            
+        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),
+        )
+        # [b, h, s, d]
+        sequence_length = query.size(2)
+        encoder_sequence_length = encoder_query.size(2)
+
+        # Hint: please check 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)
+        
+        
+                
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+
+        if get_sequence_parallel_state():
+            hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
+                (sequence_length, encoder_sequence_length), dim=2
+            )
+            # B, H, S, D
+            hidden_states = all_to_all_4D(hidden_states, scatter_dim=2, gather_dim=1)
+            encoder_hidden_states = all_gather(encoder_hidden_states, dim=1).contiguous()
+            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)
+        else:
+            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
+
+@maybe_allow_in_graph
+class MochiTransformerBlock(nn.Module):
+    r"""
+    Transformer block used in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
+
+    Args:
+        dim (`int`):
+            The number of channels in the input and output.
+        num_attention_heads (`int`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`):
+            The number of channels in each head.
+        qk_norm (`str`, defaults to `"rms_norm"`):
+            The normalization layer to use.
+        activation_fn (`str`, defaults to `"swiglu"`):
+            Activation function to use in feed-forward.
+        context_pre_only (`bool`, defaults to `False`):
+            Whether or not to process context-related conditions with additional layers.
+        eps (`float`, defaults to `1e-6`):
+            Epsilon value for normalization layers.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        pooled_projection_dim: int,
+        qk_norm: str = "rms_norm",
+        activation_fn: str = "swiglu",
+        context_pre_only: bool = False,
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+
+        self.context_pre_only = context_pre_only
+        self.ff_inner_dim = (4 * dim * 2) // 3
+        self.ff_context_inner_dim = (4 * pooled_projection_dim * 2) // 3
+
+        self.norm1 = MochiRMSNormZero(dim, 4 * dim, eps=eps, elementwise_affine=False)
+
+        if not context_pre_only:
+            self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim, eps=eps, elementwise_affine=False)
+        else:
+            self.norm1_context = LuminaLayerNormContinuous(
+                embedding_dim=pooled_projection_dim,
+                conditioning_embedding_dim=dim,
+                eps=eps,
+                elementwise_affine=False,
+                norm_type="rms_norm",
+                out_dim=None,
+            )
+
+        self.attn1 = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            bias=False,
+            qk_norm=qk_norm,
+            added_kv_proj_dim=pooled_projection_dim,
+            added_proj_bias=False,
+            out_dim=dim,
+            out_context_dim=pooled_projection_dim,
+            context_pre_only=context_pre_only,
+            processor=MochiAttnProcessor2_0(),
+            eps=eps,
+            elementwise_affine=True,
+        )
+
+        # TODO(aryan): norm_context layers are not needed when `context_pre_only` is True
+        self.norm2 = RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm2_context = RMSNorm(pooled_projection_dim, eps=eps, elementwise_affine=False)
+
+        self.norm3 = RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm3_context = RMSNorm(pooled_projection_dim, eps=eps, elementwise_affine=False)
+
+        self.ff = FeedForward(dim, inner_dim=self.ff_inner_dim, activation_fn=activation_fn, bias=False)
+        self.ff_context = None
+        if not context_pre_only:
+            self.ff_context = FeedForward(
+                pooled_projection_dim,
+                inner_dim=self.ff_context_inner_dim,
+                activation_fn=activation_fn,
+                bias=False,
+            )
+
+        self.norm4 = RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm4_context = RMSNorm(pooled_projection_dim, eps=eps, elementwise_affine=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+        output_attn = False, 
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
+
+        if not self.context_pre_only:
+            norm_encoder_hidden_states, enc_gate_msa, enc_scale_mlp, enc_gate_mlp = self.norm1_context(
+                encoder_hidden_states, temb
+            )
+        else:
+            norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb)
+
+        attn_hidden_states, context_attn_hidden_states = self.attn1(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        hidden_states = hidden_states + self.norm2(attn_hidden_states) * torch.tanh(gate_msa).unsqueeze(1)
+        norm_hidden_states = self.norm3(hidden_states) * (1 + scale_mlp.unsqueeze(1))
+        ff_output = self.ff(norm_hidden_states)
+        hidden_states = hidden_states + self.norm4(ff_output) * torch.tanh(gate_mlp).unsqueeze(1)
+
+        if not self.context_pre_only:
+            encoder_hidden_states = encoder_hidden_states + self.norm2_context(
+                context_attn_hidden_states
+            ) * torch.tanh(enc_gate_msa).unsqueeze(1)
+            norm_encoder_hidden_states = self.norm3_context(encoder_hidden_states) * (1 + enc_scale_mlp.unsqueeze(1))
+            context_ff_output = self.ff_context(norm_encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states + self.norm4_context(context_ff_output) * torch.tanh(
+                enc_gate_mlp
+            ).unsqueeze(1)
+
+        if not output_attn:
+            attn_hidden_states = None
+        return hidden_states, encoder_hidden_states, attn_hidden_states
+
+
+class MochiRoPE(nn.Module):
+    r"""
+    RoPE implementation used in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
+
+    Args:
+        base_height (`int`, defaults to `192`):
+            Base height used to compute interpolation scale for rotary positional embeddings.
+        base_width (`int`, defaults to `192`):
+            Base width used to compute interpolation scale for rotary positional embeddings.
+    """
+
+    def __init__(self, base_height: int = 192, base_width: int = 192) -> None:
+        super().__init__()
+
+        self.target_area = base_height * base_width
+
+    def _centers(self, start, stop, num, device, dtype) -> torch.Tensor:
+        edges = torch.linspace(start, stop, num + 1, device=device, dtype=dtype)
+        return (edges[:-1] + edges[1:]) / 2
+
+    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 * nccl_info.sp_size, device=device, dtype=dtype)
+        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
+    
+    def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        freqs = torch.einsum("nd,dhf->nhf", pos, freqs.float())
+        freqs_cos = torch.cos(freqs)
+        freqs_sin = torch.sin(freqs)
+        return freqs_cos, freqs_sin
+
+    def forward(
+        self,
+        pos_frequencies: torch.Tensor,
+        num_frames: int,
+        height: int,
+        width: int,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        pos = self._get_positions(num_frames, height, width, device, dtype)
+        rope_cos, rope_sin = self._create_rope(pos_frequencies, pos)
+        return rope_cos, rope_sin
+
+
+@maybe_allow_in_graph
+class MochiTransformer3DModel(ModelMixin, ConfigMixin):
+    r"""
+    A Transformer model for video-like data introduced in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
+
+    Args:
+        patch_size (`int`, defaults to `2`):
+            The size of the patches to use in the patch embedding layer.
+        num_attention_heads (`int`, defaults to `24`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`, defaults to `128`):
+            The number of channels in each head.
+        num_layers (`int`, defaults to `48`):
+            The number of layers of Transformer blocks to use.
+        in_channels (`int`, defaults to `12`):
+            The number of channels in the input.
+        out_channels (`int`, *optional*, defaults to `None`):
+            The number of channels in the output.
+        qk_norm (`str`, defaults to `"rms_norm"`):
+            The normalization layer to use.
+        text_embed_dim (`int`, defaults to `4096`):
+            Input dimension of text embeddings from the text encoder.
+        time_embed_dim (`int`, defaults to `256`):
+            Output dimension of timestep embeddings.
+        activation_fn (`str`, defaults to `"swiglu"`):
+            Activation function to use in feed-forward.
+        max_sequence_length (`int`, defaults to `256`):
+            The maximum sequence length of text embeddings supported.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 2,
+        num_attention_heads: int = 24,
+        attention_head_dim: int = 128,
+        num_layers: int = 48,
+        pooled_projection_dim: int = 1536,
+        in_channels: int = 12,
+        out_channels: Optional[int] = None,
+        qk_norm: str = "rms_norm",
+        text_embed_dim: int = 4096,
+        time_embed_dim: int = 256,
+        activation_fn: str = "swiglu",
+        max_sequence_length: int = 256,
+    ) -> None:
+        super().__init__()
+
+        inner_dim = num_attention_heads * attention_head_dim
+        out_channels = out_channels or in_channels
+
+        self.patch_embed = PatchEmbed(
+            patch_size=patch_size,
+            in_channels=in_channels,
+            embed_dim=inner_dim,
+            pos_embed_type=None,
+        )
+
+        self.time_embed = MochiCombinedTimestepCaptionEmbedding(
+            embedding_dim=inner_dim,
+            pooled_projection_dim=pooled_projection_dim,
+            text_embed_dim=text_embed_dim,
+            time_embed_dim=time_embed_dim,
+            num_attention_heads=8,
+        )
+
+        self.pos_frequencies = nn.Parameter(torch.full((3, num_attention_heads, attention_head_dim // 2), 0.0))
+        self.rope = MochiRoPE()
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                MochiTransformerBlock(
+                    dim=inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    pooled_projection_dim=pooled_projection_dim,
+                    qk_norm=qk_norm,
+                    activation_fn=activation_fn,
+                    context_pre_only=i == num_layers - 1,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(
+            inner_dim, inner_dim, elementwise_affine=False, eps=1e-6, norm_type="layer_norm"
+        )
+        self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        encoder_attention_mask: torch.Tensor,
+        output_attn = False,
+        return_dict: bool = False,
+    ) -> torch.Tensor:
+        assert return_dict is False, "return_dict is not supported in MochiTransformer3DModel"
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p = self.config.patch_size
+
+        post_patch_height = height // p
+        post_patch_width = width // p
+
+        temb, encoder_hidden_states = self.time_embed(
+            timestep, encoder_hidden_states, encoder_attention_mask, hidden_dtype=hidden_states.dtype
+        )
+
+        hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        hidden_states = self.patch_embed(hidden_states)
+        hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2)
+
+        image_rotary_emb = self.rope(
+            self.pos_frequencies,
+            num_frames,
+            post_patch_height,
+            post_patch_width,
+            device=hidden_states.device,
+            dtype=torch.float32,
+        )
+        attn_outputs_list = []
+        for i, block in enumerate(self.transformer_blocks):
+            if self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states, encoder_hidden_states, attn_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    output_attn,
+                    **ckpt_kwargs,
+                )
+            else:
+                hidden_states, encoder_hidden_states, attn_outputs = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    output_attn = output_attn,
+                )
+            attn_outputs_list.append(attn_outputs)
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1)
+        hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5)
+        output = hidden_states.reshape(batch_size, -1, num_frames, height, width)
+        
+        if not output_attn :
+            attn_outputs_list = None 
+        else:
+            attn_outputs_list = torch.stack(attn_outputs_list, dim=0)
+        return (output, attn_outputs_list)
\ No newline at end of file
diff --git a/fastvideo/model/pipeline_mochi.py b/fastvideo/model/pipeline_mochi.py
index b78d889a3..87467c5d8 100644
--- a/fastvideo/model/pipeline_mochi.py
+++ b/fastvideo/model/pipeline_mochi.py
@@ -21,7 +21,8 @@ from transformers import T5EncoderModel, T5TokenizerFast
 
 from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
 from diffusers.models.autoencoders import AutoencoderKL
-from diffusers.models.transformers import MochiTransformer3DModel
+from fastvideo.model.modeling_mochi import MochiTransformer3DModel
+
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import (
     is_torch_xla_available,
diff --git a/fastvideo/model/test.py b/fastvideo/model/test.py
index 572aa5a56..312689e87 100644
--- a/fastvideo/model/test.py
+++ b/fastvideo/model/test.py
@@ -7,8 +7,8 @@ from fastvideo.utils.parallel_states import initialize_sequence_parallel_state,
 import argparse
 import os
 from diffusers.models.transformers.transformer_mochi import MochiTransformerBlock
-from fastvideo.model.mochi_monkey_patches import hf_mochi_add_sp_monkey_patch
-from diffusers import MochiTransformer3DModel
+from fastvideo.model.modeling_mochi import MochiTransformer3DModel
+
 
 
 import sys
@@ -123,7 +123,6 @@ if __name__ == "__main__":
         single_no_patch_bs_2 = test_single_block(2)
         # check all close 
         assert torch.allclose(single_no_patch_bs_1, single_no_patch_bs_2)
-        hf_mochi_add_sp_monkey_patch()
         single_patch_bs_1 = test_single_block(1)
         single_patch_bs_2 = test_single_block(2)
         assert torch.allclose(single_patch_bs_1, single_patch_bs_2)
@@ -141,7 +140,6 @@ if __name__ == "__main__":
         single_no_patch_bs_1 = test_DiT(1, transformer, seed)
         single_no_patch_bs_2_a, single_no_patch_bs_2_b = test_DiT(2, transformer, seed)
         
-        hf_mochi_add_sp_monkey_patch()
         
         single_patch_bs_1 = test_DiT(1, transformer, seed)
         single_patch_bs_2_a, single_patch_bs_2_b = test_DiT(2, transformer, seed)
diff --git a/fastvideo/sample/sample_t2v_mochi.py b/fastvideo/sample/sample_t2v_mochi.py
index 66106f156..64b9378eb 100644
--- a/fastvideo/sample/sample_t2v_mochi.py
+++ b/fastvideo/sample/sample_t2v_mochi.py
@@ -6,8 +6,7 @@ from diffusers.utils import export_to_video
 from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info
 import argparse
 import os
-from diffusers import MochiTransformer3DModel
-from fastvideo.model.mochi_monkey_patches import hf_mochi_add_sp_monkey_patch
+from fastvideo.model.modeling_mochi import MochiTransformer3DModel
 
 def initialize_distributed():
     local_rank = int(os.getenv('RANK', 0))
@@ -21,7 +20,6 @@ def initialize_distributed():
 
 def main(args):
     initialize_distributed()
-    hf_mochi_add_sp_monkey_patch()
     print(nccl_info.sp_size)
     device = torch.cuda.current_device()
     generator = torch.Generator(device).manual_seed(args.seed)
diff --git a/fastvideo/sample/sample_t2v_mochi_no_sp.py b/fastvideo/sample/sample_t2v_mochi_no_sp.py
index 5a2738618..e67bb00ba 100644
--- a/fastvideo/sample/sample_t2v_mochi_no_sp.py
+++ b/fastvideo/sample/sample_t2v_mochi_no_sp.py
@@ -1,5 +1,6 @@
 import torch
-from diffusers import MochiPipeline, MochiTransformer3DModel
+from fastvideo.model.pipeline_mochi import MochiPipeline
+from fastvideo.model.modeling_mochi import MochiTransformer3DModel
 from diffusers.utils import export_to_video, load_image, load_video
 import argparse
 
diff --git a/fastvideo/train.py b/fastvideo/train.py
index 7db02a64c..8fe601a55 100644
--- a/fastvideo/train.py
+++ b/fastvideo/train.py
@@ -4,39 +4,38 @@ import math
 import os
 import shutil
 from pathlib import Path
-from tqdm import tqdm
 from diffusers.training_utils import cast_training_params, compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3
 from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \
     destroy_sequence_parallel_group, get_sequence_parallel_state, nccl_info
-from fastvideo.utils.communications import prepare_sequence_parallel_data, broadcast
-from fastvideo.model.mochi_monkey_patches import hf_mochi_add_sp_monkey_patch
+from fastvideo.utils.communications import sp_parallel_dataloader_wrapper, broadcast
 from fastvideo.model.mochi_latents_stat import mochi_stat
 from fastvideo.utils.validation import log_validation
 import time
 from torch.utils.data import DataLoader
-from copy import deepcopy
 import torch
-import transformers
-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
+from torch.distributed.fsdp import (
+    FullyShardedDataParallel as FSDP,
+    StateDictType,
+    FullStateDictConfig,
+)
+import json
+from torch.utils.data.distributed import DistributedSampler
+import wandb
+from accelerate.utils import set_seed
 from tqdm.auto import tqdm
-
+from fastvideo.fsdp_util import get_fsdp_kwargs, apply_fsdp_checkpointing
 import diffusers
 from diffusers import (
     FlowMatchEulerDiscreteScheduler,
-    MochiTransformer3DModel,
 )
-from diffusers.optimization import get_scheduler
-from diffusers.utils import check_min_version, is_wandb_available
+from fastvideo.model.modeling_mochi import MochiTransformer3DModel
+from diffusers.utils import check_min_version
 from fastvideo.utils.ema import EMAModel
 from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_function
-
-
+import torch.distributed as dist
+from safetensors.torch import save_file
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
 check_min_version("0.31.0")
-logger = get_logger(__name__)
 
 
 import sys
@@ -57,191 +56,178 @@ class ForkedPdb(pdb.Pdb):
             
 
 
+def main_print(content):
+    if int(os.environ['LOCAL_RANK']) <= 0: 
+        print(content)
+
+def save_checkpoint(transformer: MochiTransformer3DModel, rank, output_dir, step):
+    main_print(f"--> saving checkpoint at step {step}")
+    with FSDP.state_dict_type(
+        transformer, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    ):
+        cpu_state = transformer.state_dict()
+
+    if rank <= 0:
+        save_dir = os.path.join(output_dir, f"checkpoint-{step}")
+        os.makedirs(save_dir, exist_ok=True)
+        # save using safetensors 
+        weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors")
+        save_file(cpu_state, weight_path)
+        config_dict = dict(transformer.config)
+        config_path = os.path.join(save_dir, "config.json")
+        # save dict as json
+        with open(config_path, "w", indent=4) as f:
+            json.dump(config_dict, f)
+    main_print(f"--> checkpoint saved at step {step}")
+    
+                
+def get_sigmas(noise_scheduler, device, timesteps, n_dim=4, dtype=torch.float32):
+    sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype)
+    schedule_timesteps = noise_scheduler.timesteps.to(device)
+    timesteps = timesteps.to(device)
+    step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+    sigma = sigmas[step_indices].flatten()
+    while len(sigma.shape) < n_dim:
+        sigma = sigma.unsqueeze(-1)
+    return sigma
 
 
-class ProgressInfo:
-    def __init__(self, global_step, train_loss=0.0):
-        self.global_step = global_step
-        self.train_loss = train_loss
+def train_one_step_mochi(transformer, optimizer, loader,noise_scheduler, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm):
+    total_loss = 0.0
+    optimizer.zero_grad()
+    for _ in range(gradient_accumulation_steps):
+        latents, encoder_hidden_states, latents_attention_mask, encoder_attention_mask = next(loader)
+        
+        latents = latents.to(dtype=torch.bfloat16)
+        encoder_hidden_states = encoder_hidden_states.to(dtype=torch.bfloat16)
+        latents_mean = (
+            torch.tensor(mochi_stat.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
+        )
+        latents_std = (
+            torch.tensor(mochi_stat.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
+        )
+        latents = (latents - latents_mean) / latents_std
+        
+        batch_size = latents.shape[0]
+        noise = torch.randn_like(latents)
+        u = torch.rand(size=(batch_size,), device="cpu")
+        indices = (u * noise_scheduler.config.num_train_timesteps).long()
+        timesteps = noise_scheduler.timesteps[indices].to(device=latents.device)
+        if sp_size > 1:
+            # Make sure that the timesteps are the same across all sp processes.
+            broadcast(timesteps)
+
+        sigmas = get_sigmas(noise_scheduler, latents.device, timesteps, n_dim=latents.ndim, dtype=latents.dtype)
+        noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
+
+        model_pred = transformer(
+            noisy_model_input,
+            encoder_hidden_states,
+            noise_scheduler.config.num_train_timesteps - timesteps,
+            encoder_attention_mask, # B, L
+            return_dict= False
+        )[0]
+
+        if precondition_outputs:
+            model_pred = model_pred * sigmas + noisy_model_input
 
 
-#################################################################################
-#                                  Training Loop                                #
-#################################################################################
 
+
+        if precondition_outputs:
+            target = latents
+        else:
+            target = latents - noise
+
+        loss = torch.mean((model_pred.float() - target.float()) ** 2)
+
+        loss.backward()
+        
+        avg_loss = loss.detach().clone()
+        dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
+        total_loss += avg_loss.item() / gradient_accumulation_steps
+        
+
+    transformer.clip_grad_norm_(max_grad_norm)
+    optimizer.step()
+    return total_loss
+        
+    
 def main(args):
-    logging_dir = Path(args.output_dir, args.logging_dir)
-
     # use LayerNorm, GeLu, SiLu always as fp32 mode
     # TODO: 
     if args.enable_stable_fp32:
         raise NotImplementedError("enable_stable_fp32 is not supported now.")
-
-    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
-
-    accelerator = Accelerator(
-        gradient_accumulation_steps=args.gradient_accumulation_steps,
-        mixed_precision=args.mixed_precision,
-        log_with="wandb",
-        project_config=accelerator_project_config,
-    )
-
-    initialize_sequence_parallel_state(args.sp_size)
-    hf_mochi_add_sp_monkey_patch()
+    torch.backends.cuda.matmul.allow_tf32 = True
     
-    if not is_wandb_available():
-        raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
-    # Make one log on every process with the configuration for debugging.
-    logging.basicConfig(
-        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
-        datefmt="%m/%d/%Y %H:%M:%S",
-        level=logging.INFO,
-    )
-    logger.info(accelerator.state, main_process_only=False)
-    if accelerator.is_local_main_process:
-        transformers.utils.logging.set_verbosity_warning()
-        diffusers.utils.logging.set_verbosity_info()
-    else:
-        transformers.utils.logging.set_verbosity_error()
-        diffusers.utils.logging.set_verbosity_error()
+    local_rank = int(os.environ['LOCAL_RANK'])
+    rank = int(os.environ['RANK'])
+    world_size = int(os.environ['WORLD_SIZE'])
+    dist.init_process_group("nccl")
+    torch.cuda.set_device(local_rank)
+    device = torch.cuda.current_device()
+    initialize_sequence_parallel_state(args.sp_size)
+
+
+
 
     # If passed along, set the training seed now. On GPU...
     if args.seed is not None:
         # TODO: t within the same seq parallel group should be the same. Noise should be different.
-        set_seed(args.seed + accelerator.process_index)
+        set_seed(args.seed + rank)
 
     # Handle the repository creation
-    if accelerator.is_main_process:
-        if args.output_dir is not None:
-            os.makedirs(args.output_dir, exist_ok=True)
+    if rank <=0 and args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
 
     # For mixed precision training we cast all non-trainable weigths to half-precision
     # as these weights are only used for inference, keeping weights in full precision is not required.
-    weight_dtype = torch.float32
-    if accelerator.mixed_precision == "fp16":
-        weight_dtype = torch.float16
-    elif accelerator.mixed_precision == "bf16":
-        weight_dtype = torch.bfloat16
+
 
     # Create model:
     
-    
-    load_dtype = torch.bfloat16 # TODO
+    main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
     transformer = MochiTransformer3DModel.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="transformer",
-        torch_dtype=load_dtype,
     )
     
+    
+    main_print(f"--> Initializing FSDP with sharding strategy: full")
+    fsdp_kwargs = get_fsdp_kwargs("full", args.use_cpu_offload)
+    transformer = FSDP(
+        transformer,
+        **fsdp_kwargs,
+    )
+    
+    main_print(f"--> model loaded")
     if args.gradient_checkpointing:
-        transformer.enable_gradient_checkpointing()
-
+        apply_fsdp_checkpointing(transformer)
+        
+        
     # Set model as trainable.
     transformer.train()
 
     noise_scheduler = FlowMatchEulerDiscreteScheduler()
 
 
-    # Create EMA for the unet.
-    if args.use_ema:
-        ema_transformer = deepcopy(transformer)
-        ema_transformer = EMAModel(ema_transformer.parameters(), decay=args.ema_decay, update_after_step=args.ema_start_step,
-                             model_cls=MochiTransformer3DModel, model_config=ema_transformer.config)
-
-    def unwrap_model(model):
-        model = accelerator.unwrap_model(model)
-        model = model._orig_mod if is_compiled_module(model) else model
-        return model
-    # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
-    def save_model_hook(models, weights, output_dir):
-        if accelerator.is_main_process:
-            if args.use_ema:
-                ema_transformer.save_pretrained(os.path.join(output_dir, "model_ema"))
-
-            for model in models:
-                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
-                    model: MochiTransformer3DModel = unwrap_model(model)
-                    model.save_pretrained(
-                        os.path.join(output_dir, "transformer"), safe_serialization=True, max_shard_size="5GB"
-                    )
-                else:
-                    raise ValueError(f"Unexpected save model: {model.__class__}")
-                # make sure to pop weight so that corresponding model is not saved again
-                if weights:
-                    weights.pop()
-
-
-    def load_model_hook(models, input_dir):
-        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
-
-        if not accelerator.distributed_type == DistributedType.DEEPSPEED:
-            while len(models) > 0:
-                model = models.pop()
-
-                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)
-
-    # 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 and torch.cuda.is_available():
-        torch.backends.cuda.matmul.allow_tf32 = True
-
-    if args.scale_lr:
-        args.learning_rate = (
-                args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
-        )
-
-    # Make sure the trainable params are in float32.
-    # if args.mixed_precision == "fp16" or args.mixed_precision == "bf16":
-    #     cast_training_params([transformer], dtype=torch.float32)
-        
-
-    params_to_optimize = transformer.parameters()
-
     optimizer = torch.optim.AdamW(
-        params_to_optimize,
+        transformer.parameters(),
         lr=args.learning_rate,
         betas=(0.9,0.999),
         weight_decay=0.01,
         eps=1e-8,
     )
     
-    lr_scheduler = get_scheduler(
-        "constant",
-        optimizer=optimizer,
-        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
-        num_training_steps=args.max_train_steps * accelerator.num_processes,
-    )
 
-    logger.info(f"optimizer: {optimizer}")
+    main_print(f"optimizer: {optimizer}")
     
     train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
+    sampler = DistributedSampler(train_dataset, rank=rank, num_replicas=world_size, shuffle=True)
     train_dataloader = DataLoader(
         train_dataset,
-        shuffle=True,
+        sampler=sampler,
         collate_fn=latent_collate_function,
         pin_memory=True,
         batch_size=args.train_batch_size,
@@ -251,17 +237,6 @@ def main(args):
 
 
 
-    # Prepare everything with our `accelerator`.
-    # model.requires_grad_(False)
-    # model.pos_embed.requires_grad_(True)
-    logger.info(f'before accelerator.prepare')
-    transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        transformer, optimizer, train_dataloader, lr_scheduler
-    )
-    logger.info(f'after accelerator.prepare')
-    if args.use_ema:
-        ema_transformer.to(accelerator.device)
-
     # We need to recalculate our total training steps as the size of the training dataloader may have changed.
     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
     # Afterwards we recalculate our number of training epochs
@@ -269,237 +244,56 @@ def main(args):
 
     # We need to initialize the trackers we use, and also store our configuration.
     # The trackers initializes automatically on the main process.
-    if accelerator.is_main_process:
-        tracker_project_name = args.tracker_project_name or "fastvideo"
-        accelerator.init_trackers(tracker_project_name, config=vars(args))
+    if rank <= 0:
+        project = args.tracker_project_name or "fastvideo"
+        wandb.init(project=project, config=args)
 
     # Train!
-    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size
-    logger.info("***** Running training *****")
-    logger.info(f"  Num examples = {len(train_dataset)}")
-    logger.info(f"  Dataloader size = {len(train_dataloader)}")
-    logger.info(f"  Num Epochs = {args.num_train_epochs}")
-    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
-    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
-    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
-    logger.info(f"  Total optimization steps = {args.max_train_steps}")
-    logger.info(f"  Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B")
-    global_step = 0
-    first_epoch = 0
+    total_batch_size = args.train_batch_size * world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size
+    main_print("***** Running training *****")
+    main_print(f"  Num examples = {len(train_dataset)}")
+    main_print(f"  Dataloader size = {len(train_dataloader)}")
+    main_print(f"  Num Epochs = {args.num_train_epochs}")
+    main_print(f"  Instantaneous batch size per device = {args.train_batch_size}")
+    main_print(f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}")
+    main_print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    main_print(f"  Total optimization steps = {args.max_train_steps}")
+    main_print(f"  Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B")
+    # print dtype
+    main_print(f"  Master weight dtype: {transformer.parameters().__next__().dtype}")
+
 
     # Potentially load in the weights and states from a previous save
     if args.resume_from_checkpoint:
-        if args.resume_from_checkpoint != "latest":
-            path = os.path.basename(args.resume_from_checkpoint)
-        else:
-            # Get the most recent checkpoint
-            dirs = os.listdir(args.output_dir)
-            dirs = [d for d in dirs if d.startswith("checkpoint")]
-            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
-            path = dirs[-1] if len(dirs) > 0 else None
+        assert NotImplementedError("resume_from_checkpoint is not supported now.")
+        # TODO 
 
-        if path is None:
-            accelerator.print(
-                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
-            )
-            args.resume_from_checkpoint = None
-            initial_global_step = 0
-        else:
-            accelerator.print(f"Resuming from checkpoint {path}")
-            accelerator.load_state(os.path.join(args.output_dir, path))
-            global_step = int(path.split("-")[1])
-
-            initial_global_step = global_step
-            first_epoch = global_step // num_update_steps_per_epoch
-    else:
-        initial_global_step = 0
 
     progress_bar = tqdm(
         range(0, args.max_train_steps),
-        initial=initial_global_step,
+        initial=0,
         desc="Steps",
         # Only show the progress bar once on each machine.
-        disable=not accelerator.is_local_main_process,
+        disable= local_rank > 0,
     )
-    progress_info = ProgressInfo(global_step, train_loss=0.0)
 
-    def sync_gradients_info():
-        # Checks if the accelerator has performed an optimization step behind the scenesv
-        if args.use_ema:
-            ema_transformer.step(transformer.parameters())
-        progress_bar.update(1)
-        progress_info.global_step += 1
-        end_time = time.time()
-        one_step_duration = end_time - start_time
-        accelerator.log(
-            {
-            "train_loss": progress_info.train_loss, 
-            "learning_rate": lr_scheduler.get_last_lr()[0]
-            },
-            step=progress_info.global_step
-            )
-        logs = {"step_loss": progress_info.train_loss, "lr": lr_scheduler.get_last_lr()[0]}
-        progress_bar.set_postfix(**logs)
-        
-        progress_info.train_loss = 0.0
-
-        # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues.
-        if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
-            if progress_info.global_step % args.checkpointing_steps == 0:
-                # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
-                if accelerator.is_main_process and args.checkpoints_total_limit is not None:
-                    checkpoints = os.listdir(args.output_dir)
-                    checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
-                    checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
-
-                    # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
-                    if len(checkpoints) >= args.checkpoints_total_limit:
-                        num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
-                        removing_checkpoints = checkpoints[0:num_to_remove]
-
-                        logger.info(
-                            f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
-                        )
-                        logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
-
-                        for removing_checkpoint in removing_checkpoints:
-                            removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
-                            shutil.rmtree(removing_checkpoint)
-
-                save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}")
-                accelerator.save_state(save_path)
-                logger.info(f"Saved state to {save_path}")
-                
-    def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
-        sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype)
-        schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
-        timesteps = timesteps.to(accelerator.device)
-        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
-
-        sigma = sigmas[step_indices].flatten()
-        while len(sigma.shape) < n_dim:
-            sigma = sigma.unsqueeze(-1)
-        return sigma
-    def run(latents, encoder_hidden_states,encoder_attention_mask):
-        global start_time
-        start_time = time.time()
-        bsz = latents.shape[0]
-        noise = torch.randn_like(latents)
-        u = compute_density_for_timestep_sampling(
-            weighting_scheme=args.weighting_scheme,
-            batch_size=bsz,
-            logit_mean=args.logit_mean,
-            logit_std=args.logit_std,
-            mode_scale=args.mode_scale,
-        )
-        indices = (u * noise_scheduler.config.num_train_timesteps).long()
-        timesteps = noise_scheduler.timesteps[indices].to(device=latents.device)
-        if args.sp_size > 1:
-            # Make sure that the timesteps are the same across all sp processes.
-            broadcast(timesteps)
-        # Add noise according to flow matching.
-        # zt = (1 - texp) * x + texp * z1
-        sigmas = get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype)
-        noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
-        # accelerator.print(latents.shape, encoder_hidden_states.shape, encoder_attention_mask.shape)
-        # accelerator.print(encoder_attention_mask.tolist())
-        # print(timesteps)
-        # accelerator.print(sigmas, timesteps)
-        model_pred = transformer(
-            noisy_model_input,
-            encoder_hidden_states,
-            noise_scheduler.config.num_train_timesteps - timesteps,
-            encoder_attention_mask, # B, L
-            return_dict= False
-        )[0]
-
-        if args.precondition_outputs:
-            model_pred = model_pred * sigmas + noisy_model_input
-
-        # these weighting schemes use a uniform timestep sampling
-        # and instead post-weight the loss
-        weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas)
-
-                # flow matching loss
-        if args.precondition_outputs:
-            target = latents
-        else:
-            target = latents - noise
-
-        loss = torch.mean(weighting.float() * (model_pred.float() - target.float()) ** 2)
-
-        # Gather the losses across all processes for logging (if we use distributed training).
-        # TODO Why repeat here? Weird
-        avg_loss = accelerator.reduce(loss.clone().detach(), "mean")
-        # accelerator.print(model_pred.shape)
-        # accelerator.print(avg_loss)
-        progress_info.train_loss += avg_loss.item() / args.gradient_accumulation_steps
-        # Backpropagate
-        accelerator.backward(loss)
-        if accelerator.sync_gradients:
-            accelerator.clip_grad_norm_(transformer.parameters(), args.max_grad_norm)
-        optimizer.step()
-        lr_scheduler.step()
-        optimizer.zero_grad()
-
-        if accelerator.sync_gradients:
-            # accelerator.print("Syncing gradients")
-            sync_gradients_info()
-
-        # TODO: add indent
-        #Group ID indicates which SP group the process belongs to. We only need SP group 0 to log validation.
-            if progress_info.global_step % args.validation_steps == 0 and args.log_validation:
-                log_validation(args, transformer, accelerator,
-                                weight_dtype, progress_info.global_step)
-                if args.use_ema:
-                    # Store the UNet parameters temporarily and load the EMA parameters to perform inference.
-                    ema_transformer.store(transformer.parameters())
-                    ema_transformer.copy_to(transformer.parameters())
-                    log_validation(args, transformer,  accelerator,
-                                weight_dtype, progress_info.global_step, ema=True)
-                    # Switch back to the original UNet parameters.
-                    ema_transformer.restore(transformer.parameters())
             
 
+    loader = sp_parallel_dataloader_wrapper(train_dataloader, device, args.train_batch_size, args.sp_size, args.train_sp_batch_size)
+    
+    for step in range(1, args.max_train_steps+1):
+        loss = train_one_step_mochi(transformer, optimizer, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, args.max_grad_norm)
+        progress_bar.set_postfix({"loss": loss})
+        progress_bar.update(1)
+        if rank <= 0:
+            wandb.log({"train_loss": loss}, step=step)
+        if step  % args.checkpointing_steps == 0:
+            save_checkpoint(transformer, rank, args.output_dir, step)
+            
+        if args.log_validation and step  % args.validation_steps == 0:
+            log_validation(args, transformer, device,
+                            torch.bfloat16, step)
 
-        return loss
-
-    def sp_parallel_dataloader_wrapper(dataloader):
-        for data_item in dataloader:
-            latents, cond,attn_mask, cond_mask = data_item    
-            latents, cond, attn_mask, cond_mask = prepare_sequence_parallel_data(latents, cond, attn_mask, cond_mask)
-            # accelerator.print(latents.shape, cond.shape, attn_mask.shape, cond_mask.shape)
-            for iter in range(args.train_batch_size * args.sp_size // args.train_sp_batch_size):
-                st_idx = iter * args.train_sp_batch_size
-                ed_idx = (iter + 1) * args.train_sp_batch_size
-                encoder_hidden_states=cond[st_idx: ed_idx]
-                attention_mask=attn_mask[st_idx: ed_idx]
-                encoder_attention_mask=cond_mask[st_idx: ed_idx]
-                yield latents[st_idx: ed_idx], encoder_hidden_states, attention_mask, encoder_attention_mask
-                    
-                    
-    def train():
-        for epoch in range(first_epoch, args.num_train_epochs):
-            progress_info.train_loss = 0.0
-            for latents, cond, attn_mask, cond_mask in sp_parallel_dataloader_wrapper(train_dataloader):
-                latents = latents.to(dtype=torch.bfloat16)
-                cond = cond.to(dtype=torch.bfloat16)
-                latents_mean = (
-                    torch.tensor(mochi_stat.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
-                )
-                latents_std = (
-                    torch.tensor(mochi_stat.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
-                )
-                latents = (latents - latents_mean) / latents_std
-                with accelerator.accumulate(transformer):
-                    run(latents, cond,  cond_mask)
-                
-                if progress_info.global_step >= args.max_train_steps:
-                    return
-
-    train()
-    accelerator.wait_for_everyone()
-    accelerator.end_training()
     if get_sequence_parallel_state():
         destroy_sequence_parallel_group()
 
@@ -522,30 +316,9 @@ if __name__ == "__main__":
     parser.add_argument('--enable_stable_fp32', action='store_true') # TODO
 
     # diffusion setting
-    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
     parser.add_argument("--ema_decay", type=float, default=0.999)
     parser.add_argument("--ema_start_step", type=int, default=0)
     parser.add_argument('--cfg', type=float, default=0.1)
-    parser.add_argument("--prediction_type", type=str, default=None, help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.")
-            
-    parser.add_argument(
-        "--weighting_scheme",
-        type=str,
-        default="uniform",
-        choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"],
-    )
-    parser.add_argument(
-        "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme."
-    )
-    parser.add_argument(
-        "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
-    )
-    parser.add_argument(
-        "--mode_scale",
-        type=float,
-        default=1.29,
-        help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
-    )
     parser.add_argument("--precondition_outputs", action="store_true", help="Whether to precondition the outputs of the model.")
     
     # validation & logs
@@ -601,6 +374,7 @@ if __name__ == "__main__":
                             " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
                         ),
                         )
+    parser.add_argument("--use_cpu_offload", action="store_true", help="Whether to use CPU offload for param & gradient & optimizer states.")
 
     parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel")
     parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training")
diff --git a/fastvideo/utils/communications.py b/fastvideo/utils/communications.py
index 6de4dd5d9..b05b35a6f 100644
--- a/fastvideo/utils/communications.py
+++ b/fastvideo/utils/communications.py
@@ -274,4 +274,21 @@ def prepare_sequence_parallel_data(hidden_states, encoder_hidden_states, attenti
     return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask
 
 
-    
\ No newline at end of file
+
+def sp_parallel_dataloader_wrapper(dataloader, device, train_batch_size, sp_size, train_sp_batch_size):
+        while True:
+            for data_item in dataloader:
+                latents, cond,attn_mask, cond_mask = data_item    
+                latents = latents.to(device)
+                cond = cond.to(device)
+                attn_mask = attn_mask.to(device)
+                cond_mask = cond_mask.to(device)
+                latents, cond, attn_mask, cond_mask = prepare_sequence_parallel_data(latents, cond, attn_mask, cond_mask)
+
+                for iter in range(train_batch_size * sp_size // train_sp_batch_size):
+                    st_idx = iter * train_sp_batch_size
+                    ed_idx = (iter + 1) * train_sp_batch_size
+                    encoder_hidden_states=cond[st_idx: ed_idx]
+                    attention_mask=attn_mask[st_idx: ed_idx]
+                    encoder_attention_mask=cond_mask[st_idx: ed_idx]
+                    yield latents[st_idx: ed_idx], encoder_hidden_states, attention_mask, encoder_attention_mask
\ No newline at end of file
diff --git a/fastvideo/utils/data_preprocess/finetune_data_T5.py b/fastvideo/utils/data_preprocess/finetune_data_T5.py
index 7c00d0bca..796e73ae7 100644
--- a/fastvideo/utils/data_preprocess/finetune_data_T5.py
+++ b/fastvideo/utils/data_preprocess/finetune_data_T5.py
@@ -1,7 +1,7 @@
 import argparse
 import torch
 from accelerate.logging import get_logger
-from diffusers import MochiPipeline
+from fastvideo.model.pipeline_mochi import MochiPipeline
 from diffusers.utils import export_to_video
 import json
 import os
diff --git a/fastvideo/utils/encode_vae_t5.py b/fastvideo/utils/encode_vae_t5.py
index 527564262..ce908ca69 100644
--- a/fastvideo/utils/encode_vae_t5.py
+++ b/fastvideo/utils/encode_vae_t5.py
@@ -38,9 +38,9 @@ import diffusers
 from diffusers import (
     AutoencoderKL,
     FlowMatchEulerDiscreteScheduler,
-    MochiTransformer3DModel,
-    MochiPipeline,
 )
+from fastvideo.model.modeling_mochi import MochiTransformer3DModel
+from fastvideo.model.pipeline_mochi import MochiPipeline
 from diffusers.optimization import get_scheduler
 from diffusers.training_utils import compute_snr
 from diffusers.utils import check_min_version, is_wandb_available
diff --git a/fastvideo/utils/validation.py b/fastvideo/utils/validation.py
index 443b627c8..141b6132c 100644
--- a/fastvideo/utils/validation.py
+++ b/fastvideo/utils/validation.py
@@ -177,12 +177,11 @@ def sample_validation_video(
 
 
 @torch.inference_mode()
-def log_validation(args, transformer,accelerator, weight_dtype, global_step,  ema=False):
+def log_validation(args, transformer, device, weight_dtype, global_step,  ema=False):
     #TODO
-    accelerator.print(f"Running validation....\n")
-    transformer = accelerator.unwrap_model(transformer)
-    generator = torch.Generator(device=accelerator.device).manual_seed(12345)
-    vae = AutoencoderKLMochi.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", torch_dtype=weight_dtype).to(accelerator.device)
+    print(f"Running validation....\n")
+    generator = torch.Generator(device="cuda").manual_seed(12345)
+    vae = AutoencoderKLMochi.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", torch_dtype=weight_dtype).to("cuda")
     vae.enable_tiling()
     scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler", torch_dtype=weight_dtype)
     
@@ -191,10 +190,10 @@ def log_validation(args, transformer,accelerator, weight_dtype, global_step,  em
     prompt_mask_path = os.path.join(args.validation_prompt_dir, "mask.pt")
     negative_prompt_embed_path = os.path.join(args.uncond_prompt_dir, "embed.pt")
     negative_prompt_mask_path = os.path.join(args.uncond_prompt_dir, "mask.pt")
-    prompt_embeds = torch.load(prompt_embed_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
-    prompt_attention_mask = torch.load(prompt_mask_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
-    negative_prompt_embeds = torch.load(negative_prompt_embed_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
-    negative_prompt_attention_mask = torch.load(negative_prompt_mask_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
+    prompt_embeds = torch.load(prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0)
+    prompt_attention_mask = torch.load(prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0)
+    negative_prompt_embeds = torch.load(negative_prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0)
+    negative_prompt_attention_mask = torch.load(negative_prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0)
     
 
     videos = []
@@ -219,20 +218,20 @@ def log_validation(args, transformer,accelerator, weight_dtype, global_step,  em
     gc.collect()
     torch.cuda.empty_cache()
     # log if main process
-    if accelerator.is_main_process :
+    if int(os.environ['RANK']) <= 0:
         video_filenames = []
         for i, video in enumerate(videos):
             filename = os.path.join(args.output_dir, f"validation_step_{global_step}_video_{i}.mp4")
             export_to_video(video, filename, fps=30)
             video_filenames.append(filename)
-        for tracker in accelerator.trackers:
-            logs = {
-                f"{'ema_' if ema else ''}validation": [
-                    wandb.Video(filename)
-                    for i, filename in enumerate(video_filenames)
-                ]
-            }
-            tracker.log(logs, step=global_step)
+
+        logs = {
+            f"{'ema_' if ema else ''}validation": [
+                wandb.Video(filename)
+                for i, filename in enumerate(video_filenames)
+            ]
+        }
+        wandb.log(logs, step=global_step)
 
     del vae
     del prompt_embeds
diff --git a/inference_mochi.sh b/inference_mochi.sh
index e37ede3ba..bd87bc8e7 100644
--- a/inference_mochi.sh
+++ b/inference_mochi.sh
@@ -17,18 +17,18 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
 
 
 
-num_gpus=2
+num_gpus=4
 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     fastvideo/sample/sample_t2v_mochi.py \
     --model_path data/mochi \
-    --prompt_embed_path data/Encoder_Overfit_Data/prompt_embed/1.pt \
-    --encoder_attention_mask_path data/Encoder_Overfit_Data/prompt_attention_mask/1.pt \
-    --num_frames 79 \
+    --prompt_embed_path data/Encoder_Overfit_Data/prompt_embed/0.pt \
+    --encoder_attention_mask_path data/Encoder_Overfit_Data/prompt_attention_mask/0.pt \
+    --num_frames 163 \
     --height 480 \
     --width 848 \
     --num_inference_steps 64 \
     --guidance_scale 4.5 \
-    --transformer_path data/outputs/debug_uniform_14_sp2_spb2_4videos_2K/checkpoint-1000/transformer \
+    --transformer_path data/outputs/BW_Testrun/checkpoint-0 \
     --output_path outputs_video/overfit/uniform_cfg_train_1000_cfg_f79_sp2_spb2_4videos_second \
     --seed 12345
 
diff --git a/pretty_json.py b/pretty_json.py
index 380522342..50f98da4a 100644
--- a/pretty_json.py
+++ b/pretty_json.py
@@ -1,7 +1,7 @@
 import json
 import os 
 
-path = "data/Mochi-Synthetic-Data/videos2caption.json"
+path = "data/outputs/BW_Testrun/checkpoint-0/config.json"
 
 with open(path, 'r') as f:
     data = json.load(f)
diff --git a/scripts/accelerate_configs/deepspeed_zero2_config.yaml b/scripts/accelerate_configs/deepspeed_zero2_config.yaml
deleted file mode 100644
index 3622f2ece..000000000
--- a/scripts/accelerate_configs/deepspeed_zero2_config.yaml
+++ /dev/null
@@ -1,13 +0,0 @@
-compute_environment: LOCAL_MACHINE
-distributed_type: DEEPSPEED
-deepspeed_config:
- deepspeed_config_file: scripts/accelerate_configs/zero2.json
-fsdp_config: {}
-machine_rank: 0
-main_process_ip: null
-main_process_port: 29503
-main_training_function: main
-num_machines: 1
-num_processes: 4
-# gpu_ids: 0
-use_cpu: false
diff --git a/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml b/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml
deleted file mode 100644
index 9c96bae67..000000000
--- a/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml
+++ /dev/null
@@ -1,13 +0,0 @@
-compute_environment: LOCAL_MACHINE
-distributed_type: DEEPSPEED
-deepspeed_config:
- deepspeed_config_file: scripts/accelerate_configs/zero2_offload.json
-fsdp_config: {}
-machine_rank: 0
-main_process_ip: null
-main_process_port: 29501
-main_training_function: main
-num_machines: 1
-num_processes: 8
-gpu_ids: 0,1,2,3,4,5,6,7
-use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/deepspeed_zero3_config.yaml b/scripts/accelerate_configs/deepspeed_zero3_config.yaml
deleted file mode 100644
index bfd8fa308..000000000
--- a/scripts/accelerate_configs/deepspeed_zero3_config.yaml
+++ /dev/null
@@ -1,13 +0,0 @@
-compute_environment: LOCAL_MACHINE
-distributed_type: DEEPSPEED
-deepspeed_config:
- deepspeed_config_file: scripts/accelerate_configs/zero3.json
- zero3_init_flag: true
-fsdp_config: {}
-machine_rank: 0
-main_process_ip: null
-main_process_port: 29501
-main_training_function: main
-num_machines: 1
-num_processes: 4
-use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml b/scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml
deleted file mode 100644
index b01631caf..000000000
--- a/scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml
+++ /dev/null
@@ -1,13 +0,0 @@
-compute_environment: LOCAL_MACHINE
-distributed_type: DEEPSPEED
-deepspeed_config:
- deepspeed_config_file: scripts/accelerate_configs/zero3_offload.json
-fsdp_config: {}
-machine_rank: 0
-main_process_ip: null
-main_process_port: 29501
-main_training_function: main
-num_machines: 1
-num_processes: 8
-gpu_ids: 0,1,2,3,4,5,6,7
-use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/default_config.yaml b/scripts/accelerate_configs/default_config.yaml
deleted file mode 100644
index ac3bb99b5..000000000
--- a/scripts/accelerate_configs/default_config.yaml
+++ /dev/null
@@ -1,12 +0,0 @@
-compute_environment: LOCAL_MACHINE
-distributed_type: MULTI_GPU
-fsdp_config: {}
-machine_rank: 0
-main_process_ip: null
-main_process_port: 29501
-main_training_function: main
-mixed_precision: bf16
-num_machines: 1
-num_processes: 8
-gpu_ids: 0,1,2,3,4,5,6,7
-use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/hostfile b/scripts/accelerate_configs/hostfile
deleted file mode 100644
index becb02e09..000000000
--- a/scripts/accelerate_configs/hostfile
+++ /dev/null
@@ -1,16 +0,0 @@
-node30 slots=8
-node005 slots=8
-node006 slots=8
-node007 slots=8
-node008 slots=8
-node010 slots=8
-node012 slots=8
-node013 slots=8
-node019 slots=8
-node026 slots=8
-node027 slots=8
-node028 slots=8
-node029 slots=8
-node031 slots=8
-node059 slots=8
-node060 slots=8
\ No newline at end of file
diff --git a/scripts/accelerate_configs/multi_node_example.yaml b/scripts/accelerate_configs/multi_node_example.yaml
deleted file mode 100644
index 43a9f7f48..000000000
--- a/scripts/accelerate_configs/multi_node_example.yaml
+++ /dev/null
@@ -1,18 +0,0 @@
-compute_environment: LOCAL_MACHINE
-distributed_type: DEEPSPEED
-deepspeed_config:
- deepspeed_config_file: scripts/accelerate_configs/zero2.json
- deepspeed_hostfile: scripts/accelerate_configs/hostfile
-fsdp_config: {}
-machine_rank: 0
-main_process_ip: 100.64.24.30
-main_process_port: 29522
-main_training_function: main
-num_machines: 16
-num_processes: 128
-rdzv_backend: static
-same_network: true
-tpu_env: []
-tpu_use_cluster: false
-tpu_use_sudo: false
-use_cpu: false
diff --git a/scripts/accelerate_configs/multi_node_example_by_ddp.yaml b/scripts/accelerate_configs/multi_node_example_by_ddp.yaml
deleted file mode 100644
index 17fd420ac..000000000
--- a/scripts/accelerate_configs/multi_node_example_by_ddp.yaml
+++ /dev/null
@@ -1,13 +0,0 @@
-compute_environment: LOCAL_MACHINE
-distributed_type: MULTI_GPU
-fsdp_config: {}
-main_process_port: 29501
-main_training_function: main
-num_machines: 2
-num_processes: 16
-rdzv_backend: static
-same_network: true
-tpu_env: []
-tpu_use_cluster: false
-tpu_use_sudo: false
-use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/single.yaml b/scripts/accelerate_configs/single.yaml
deleted file mode 100644
index 5c22b3f30..000000000
--- a/scripts/accelerate_configs/single.yaml
+++ /dev/null
@@ -1,13 +0,0 @@
-compute_environment: LOCAL_MACHINE
-distributed_type: DEEPSPEED
-deepspeed_config:
- deepspeed_config_file: scripts/accelerate_configs/zero2.json
-fsdp_config: {}
-machine_rank: 0
-main_process_ip: null
-main_process_port: 29503
-main_training_function: main
-num_machines: 1
-num_processes: 1
-# gpu_ids: 0
-use_cpu: false
diff --git a/scripts/accelerate_configs/zero2.json b/scripts/accelerate_configs/zero2.json
deleted file mode 100644
index 8dcdd154d..000000000
--- a/scripts/accelerate_configs/zero2.json
+++ /dev/null
@@ -1,25 +0,0 @@
-{
-    "fp16": {
-        "enabled": false,
-        "loss_scale": 0,
-        "loss_scale_window": 1000,
-        "initial_scale_power": 16,
-        "hysteresis": 2,
-        "min_loss_scale": 1
-    },
-    "bf16": {
-        "enabled": "auto"
-    },
-    "communication_data_type": "fp32",
-    "gradient_clipping": 1.0,
-    "train_micro_batch_size_per_gpu": "auto",
-    "train_batch_size": "auto",
-    "gradient_accumulation_steps": "auto",
-    "zero_optimization": {
-        "stage": 2,
-        "overlap_comm": true,
-        "contiguous_gradients": true,
-        "sub_group_size": 1e9,
-        "reduce_bucket_size": 5e8
-    }
-}
\ No newline at end of file
diff --git a/scripts/accelerate_configs/zero2_offload.json b/scripts/accelerate_configs/zero2_offload.json
deleted file mode 100644
index c315ba5fb..000000000
--- a/scripts/accelerate_configs/zero2_offload.json
+++ /dev/null
@@ -1,26 +0,0 @@
-{
-    "fp16": {
-        "enabled": "auto",
-        "loss_scale": 0,
-        "loss_scale_window": 1000,
-        "initial_scale_power": 16,
-        "hysteresis": 2,
-        "min_loss_scale": 1
-    },
-    "bf16": {
-        "enabled": "auto"
-    },
-    "train_micro_batch_size_per_gpu": "auto",
-    "train_batch_size": "auto",
-    "gradient_accumulation_steps": "auto",
-    "zero_optimization": {
-        "stage": 2,
-        "offload_optimizer": {
-            "device": "cpu"
-        },
-        "overlap_comm": true,
-        "contiguous_gradients": true,
-        "sub_group_size": 1e9,
-        "reduce_bucket_size": 5e8
-    }
-}
\ No newline at end of file
diff --git a/scripts/accelerate_configs/zero3.json b/scripts/accelerate_configs/zero3.json
deleted file mode 100644
index 90a219f80..000000000
--- a/scripts/accelerate_configs/zero3.json
+++ /dev/null
@@ -1,28 +0,0 @@
-{
-    "fp16": {
-        "enabled": "auto",
-        "loss_scale": 0,
-        "loss_scale_window": 1000,
-        "initial_scale_power": 16,
-        "hysteresis": 2,
-        "min_loss_scale": 1
-    },
-    "bf16": {
-        "enabled": "auto"
-    },
-    "train_micro_batch_size_per_gpu": "auto",
-    "train_batch_size": "auto",
-    "gradient_accumulation_steps": "auto",
-    "zero_optimization": {
-        "stage": 3,
-        "overlap_comm": true,
-        "contiguous_gradients": true,
-        "sub_group_size": 1e9,
-        "reduce_bucket_size": 5e8,
-        "stage3_prefetch_bucket_size": 5e8,
-        "stage3_param_persistence_threshold": 1e6,
-        "stage3_max_live_parameters": 1e9,
-        "stage3_max_reuse_distance": 1e9,
-        "stage3_gather_16bit_weights_on_model_save": true
-    }
-}
\ No newline at end of file
diff --git a/scripts/accelerate_configs/zero3_offload.json b/scripts/accelerate_configs/zero3_offload.json
deleted file mode 100644
index c15bcf91b..000000000
--- a/scripts/accelerate_configs/zero3_offload.json
+++ /dev/null
@@ -1,39 +0,0 @@
-{
-  "fp16": {
-    "enabled": "auto",
-    "loss_scale": 0,
-    "loss_scale_window": 1000,
-    "initial_scale_power": 16,
-    "hysteresis": 2,
-    "min_loss_scale": 1
-  },
-  "bf16": {
-    "enabled": "auto"
-  },
-  "zero_optimization": {
-    "stage": 3,
-    "offload_optimizer": {
-      "device": "cpu",
-      "pin_memory": true
-    },
-    "offload_param": {
-      "device": "cpu",
-      "pin_memory": true
-    },
-    "overlap_comm": true,
-    "contiguous_gradients": true,
-    "sub_group_size": 1e9,
-    "reduce_bucket_size": 5e8,
-    "stage3_prefetch_bucket_size": "auto",
-    "stage3_param_persistence_threshold": "auto",
-    "stage3_max_live_parameters": 1e9,
-    "stage3_max_reuse_distance": 1e9,
-    "gather_16bit_weights_on_model_save": true
-  },
-  "gradient_accumulation_steps": "auto",
-  "gradient_clipping": "auto",
-  "train_batch_size": "auto",
-  "train_micro_batch_size_per_gpu": "auto",
-  "steps_per_print": 1e5,
-  "wall_clock_breakdown": false
-}
\ No newline at end of file
diff --git a/t2v_debug_multi.sh b/t2v_debug_multi.sh
index 06e848f0b..9219ad33c 100644
--- a/t2v_debug_multi.sh
+++ b/t2v_debug_multi.sh
@@ -1,53 +1,23 @@
-
-accelerate launch \
-    --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
+torchrun --nnodes 1 --nproc_per_node 4 \
     fastvideo/train.py \
     --seed 42 \
     --pretrained_model_name_or_path data/mochi \
     --cache_dir "data/.cache" \
     --data_json_path "data/Encoder_Overfit_Data/videos2caption.json" \
+    --validation_prompt_dir "data/Encoder_Overfit_Data/validation_prompt_embed_mask" \
+    --uncond_prompt_dir "data/Encoder_Overfit_Data/uncond_prompt_embed_mask" \
     --gradient_checkpointing \
     --train_batch_size=1 \
-    --num_latent_t 14 \
-    --sp_size 2 \
-    --train_sp_batch_size 2 \
+    --num_latent_t 1 \
+    --sp_size 1 \
+    --train_sp_batch_size 1 \
     --dataloader_num_workers 1 \
-    --gradient_accumulation_steps=2 \
+    --gradient_accumulation_steps=1 \
     --max_train_steps=2000 \
     --learning_rate=1e-5 \
     --mixed_precision="bf16" \
-    --weighting_scheme "uniform" \
-    --checkpointing_steps=500 \
-    --checkpoints_total_limit 3 \
-    --allow_tf32 \
-    --ema_start_step 0 \
-    --cfg 0.1 \
-    --ema_decay 0.999 \
-    --output_dir="data/outputs/debug_uniform_14_sp2_spb2_4videos_2K"
-
-
-
-
-accelerate launch \
-    --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
-    fastvideo/train.py \
-    --seed 42 \
-    --pretrained_model_name_or_path data/mochi \
-    --cache_dir "data/.cache" \
-    --data_json_path "data/Encoder_Overfit_Data/videos2caption.json" \
-    --gradient_checkpointing \
-    --train_batch_size=1 \
-    --num_latent_t 20 \
-    --sp_size 4 \
-    --train_sp_batch_size 1 \
-    --dataloader_num_workers 1 \
-    --gradient_accumulation_steps=2 \
-    --max_train_steps=500 \
-    --learning_rate=0.0 \
-    --mixed_precision="bf16" \
-    --weighting_scheme "uniform" \
-    --checkpointing_steps=100 \
-    --validation_steps 2 \
+    --checkpointing_steps=200 \
+    --validation_steps 200 \
     --validation_sampling_steps 64 \
     --checkpoints_total_limit 3 \
     --allow_tf32 \
@@ -55,48 +25,19 @@ accelerate launch \
     --cfg 0.1 \
     --ema_decay 0.999 \
     --log_validation \
-    --validation_prompt_dir "data/Encoder_Overfit_Data/validation_prompt_embed_mask" \
-    --uncond_prompt_dir "data/Encoder_Overfit_Data/uncond_prompt_embed_mask" \
-    --output_dir="data/outputs/debug_logvalidation_sp2"
+    --output_dir="data/outputs/BW_Testrun" \
+    --use_cpu_offload
 
 
 
-torchrun --nnodes=1 --nproc_per_node=4 --master_port 29503 debug_mochi_sp.py
-torchrun --nnodes=1 --nproc_per_node=4 --master_port 29503  debug_OSP_A2A.py
 
 
-for seed in 42 
-do
-torchrun --nnodes=1 --nproc_per_node=4 --master_port 29503 fastvideo/model/test.py --seed $seed
-done
+
+# SP single node: 13
+# SP multi node: 20
+# SP multi node with larger batch size: 15
+# No SP Single node: 10.5
+# No SP Multe node: 11
 
 
-num_gpus=1
 
-torchrun --nproc_per_node=$num_gpus fastvideo/sample/generate_synthetic.py \
-    --model_path data/mochi \
-    --num_frames 163 \
-    --height 480 \
-    --width 848 \
-    --num_inference_steps 2 \
-    --guidance_scale 4.5 \
-    --prompt_path "data/prompt.txt" \
-    --dataset_output_dir data/synthetic_debug3
-
-    
-
-
-"A person is sawing a piece of wood placed on a support, using a large hand saw.
-The saw moves rapidly back and forth across the top of the wood, generating small wood chips and dust. 
-The person's hands are visible, gripping the saw firmly. The background shows a workshop environment with indistinct objects and a bench.
-The focus remains on the action of sawing, with no significant change in the camera angle or background.
-The sawing action continues with the saw moving through the wood, which shows a visible cut deepening as the saw progresses. 
-Wood chips continue to scatter from the cutting area. 
-The person's grip adjusts slightly for better control and force.
-The background remains consistent with the workshop setting, and the camera stays focused on the sawing action without any shifts in perspective
-No new objects or persons enter the scene. 
-The sawing action reaches a deeper cut into the wood, indicating progress in the task. 
-The person's movements become more pronounced as they apply more force. 
-The saw blade appears more prominently as it moves back and forth. 
-The background and setting remain unchanged, maintaining the focus on the sawing activity.
-The camera continues to capture the action closely, emphasizing the interaction between the saw and the wood."