run lint

cosmos2 dit
2025-07-06 01:16:42 -07:00 · 2025-07-06 01:12:05 -07:00
80 changed files with 1845 additions and 584 deletions
@@ -58,8 +58,6 @@ steps:
                  queue: "default"
            - path: 
                - "fastvideo/v1/**/*.py"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 30m .buildkite/scripts/pr_test.sh"
                label: "SSIM Tests"
@@ -67,21 +65,6 @@ steps:
                  - TEST_TYPE=ssim
                agents:
                  queue: "default"
-            - path:
-                - "fastvideo/v1/tests/lora/**"
-                - "fastvideo/v1/models/loader/**"
-                - "fastvideo/v1/tests/transformers/**"
-                - "fastvideo/v1/pipelines/**"
-                - "fastvideo/v1/layers/lora/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "LoRA Inference Tests"
-                env:
-                  - TEST_TYPE=inference_lora
-                agents:
-                  queue: "default"
            - path:
                - "fastvideo/v1/**"
                - "pyproject.toml"
@@ -97,10 +97,6 @@ case "$TEST_TYPE" in
        log "Running precision VSA tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_VSA"
        ;;
-    "inference_lora")
-        log "Running LoRA tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_lora_tests"
-        ;;
    *)
        log "Error: Unknown test type: $TEST_TYPE"
        exit 1
@@ -13,4 +13,4 @@
      ]
    }
  ]
-}
+}
@@ -372,4 +372,4 @@ jobs:
          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12", "training-test", "training-test-VSA", "inference-test-STA", "precision-test-STA", "precision-test-VSA"]'
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
-        run: python .github/scripts/runpod_cleanup.py
+        run: python .github/scripts/runpod_cleanup.py
@@ -60,7 +60,7 @@ repos:
  rev: v1.15.0
  hooks:
  - id: mypy
-    args: [--python-version, '3.10', --follow-imports, "skip" ]
+    args: [--python-version, '3.10', --follow-imports, "skip", ]
    additional_dependencies: [types-cachetools, types-setuptools, types-PyYAML, types-requests]
 - repo: local
  hooks:
@@ -69,7 +69,7 @@ repos:
    entry: bash
    args:
      - -c
-      - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/" | grep -v "^fastvideo/v1/tests/inference/lora/L40S_reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
+      - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
    language: system
    always_run: true
    pass_filenames: false
@@ -8,7 +8,7 @@ It features a clean, consistent API that works across popular video models, maki
 With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.

 <p align="center">
-    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |
+    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
 </p>

 <div align="center">
@@ -8,14 +8,12 @@ You can easily use the FastVideo Docker image as a custom container on [RunPod](

 Choose a GPU that supports CUDA 12.4

-Pick 1 or 2 L40S GPU(s)
-
 ![RunPod CUDA selection](../../_static/images/runpod_cuda.png)

 When creating your pod template, use this image:

 ```
-ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest
+ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
 ```

 Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.io/pods/configuration/use-ssh)):
@@ -117,4 +117,4 @@ If you're planning to contribute to FastVideo please see the following page:

 If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).

-You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg) for additional support.
@@ -12,7 +12,7 @@ This guide explains how to implement a custom diffusion pipeline in FastVideo, l
 4. **Register Your Pipeline** - Make it discoverable by the framework
 5. **Configure Your Pipeline** - (Coming soon)

-Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
+Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).

 ## Step 1: Pipeline Modules

@@ -27,7 +27,7 @@ fastvideo generate --help
 ### Hardware Configuration

 - `--num-gpus {NUM_GPUS}`: Number of GPUs to use
- `--tp-size {TP_SIZE}`: Tensor parallelism size (only for the encoder, should not be larger than 1 if text encoder offload is enabled, as layerwise offload + prefetch is faster)
+- `--tp-size {TP_SIZE}`: Tensor parallelism size (Typically should match the number of GPUs)
 - `--sp-size {SP_SIZE}`: Sequence parallelism size (Typically should match the number of GPUs)

 #### Video Configuration
@@ -68,7 +68,7 @@ Example configuration file (config.json):
    "output_path": "outputs/",
    "num_gpus": 2,
    "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
    "num_frames": 45,
    "height": 720,
    "width": 1280,
@@ -102,7 +102,7 @@ prompt: "A beautiful woman in a red dress walking down a street"
 output_path: "outputs/"
 num_gpus: 2
 sp_size: 2
-tp_size: 1
+tp_size: 2
 num_frames: 45
 height: 720
 width: 1280
@@ -121,4 +121,4 @@ If the generated video doesn't match your prompt:
 - Learn about using [Optimizations](#inference-optimizations)
 - See [Examples](../examples/examples_inference_index.md) for more usage scenarios
 - Join our [Community Discord](https://discord.gg/JA7cksDz86).
- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).
@@ -6,7 +6,7 @@ def main():
    # Initialize VideoGenerator with the Wan model
    generator = VideoGenerator.from_pretrained(
        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        num_gpus=1,
+        num_gpus=2,
        lora_path="benjamin-paine/steamboat-willie-1.3b",
        lora_nickname="steamboat"
    )
@@ -16,7 +16,6 @@ def main():
        "num_frames": 81,
        "guidance_scale": 5.0,
        "num_inference_steps": 32,
-        "seed": 42,
    }
    # Generate video with LoRA style
    prompt = "steamboat willie style, golden era animation, close-up of a short fluffy monster  kneeling beside a melting red candle. the mood is one of wonder and curiosity,  as the monster gazes at the flame with wide eyes and open mouth. Its pose and expression  convey a sense of innocence and playfulness, as if it is exploring the world around it for the first time.  The use of warm colors and dramatic lighting further enhances the cozy atmosphere of the image."
@@ -30,17 +29,8 @@ def main():
        negative_prompt=negative_prompt,
        **kwargs
    )
-    del generator
-    
-    # Until FSDP resharding bug is fixed, multi-lora requires reloading the model
-    # see https://github.com/pytorch/pytorch/issues/157209
-    generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        num_gpus=1,
-        lora_path="motimalu/wan-flat-color-1.3b-v2",
-        lora_nickname="flat_color"
-    )
-    # generator.set_lora_adapter(lora_nickname="flat_color", lora_path="motimalu/wan-flat-color-1.3b-v2")
+
+    generator.set_lora_adapter(lora_nickname="flat_color", lora_path="motimalu/wan-flat-color-1.3b-v2")
    prompt = "flat color, no lineart, blending, negative space, artist:[john kafka|ponsuke kaikai|hara id 21|yoneyama mai|fuzichoco],  1girl, sakura miko, pink hair, cowboy shot, white shirt, floral print, off shoulder, outdoors, cherry blossom, tree shade, wariza, looking up, falling petals, half-closed eyes, white sky, clouds,  live2d animation, upper body, high quality cinematic video of a woman sitting under a sakura tree. Dreamy and lonely, the camera close-ups on the face of the woman as she turns towards the viewer. The Camera is steady, This is a cowboy shot. The animation is smooth and fluid."
    negative_prompt = "bad quality video,色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
    video = generator.generate_video(
@@ -24,7 +24,6 @@ training_args=(
  --num_height 480
  --num_width 832
  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
@@ -1,11 +1,13 @@
 #!/bin/bash
 #SBATCH --job-name=i2v
 #SBATCH --partition=main
+#SBATCH --qos=hao
 #SBATCH --nodes=4
 #SBATCH --ntasks=4
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
+#SBATCH --nodelist=fs-mbz-gpu-[100-850]
 #SBATCH --mem=1440G
 #SBATCH --output=i2v_output/i2v_%j.out
 #SBATCH --error=i2v_output/i2v_%j.err
@@ -58,7 +60,6 @@ training_args=(
  --num_height 480
  --num_width 832
  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
@@ -24,14 +24,13 @@ training_args=(
  --num_height 480
  --num_width 832
  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size 8
-  --tp_size 1
+  --tp_size 8
  --hsdp_replicate_dim 1
  --hsdp_shard_dim 8
 )
@@ -1,11 +1,13 @@
 #!/bin/bash
 #SBATCH --job-name=i2v
 #SBATCH --partition=main
+#SBATCH --qos=hao
 #SBATCH --nodes=4
 #SBATCH --ntasks=4
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
+#SBATCH --nodelist=fs-mbz-gpu-[100-850]
 #SBATCH --mem=1440G
 #SBATCH --output=i2v_output/i2v_%j.out
 #SBATCH --error=i2v_output/i2v_%j.err
@@ -58,14 +60,13 @@ training_args=(
  --num_height 480
  --num_width 832
  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size $NUM_GPUS
-  --tp_size 1
+  --tp_size $NUM_GPUS
  --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
  --hsdp_shard_dim $NUM_GPUS
 )
@@ -24,14 +24,13 @@ training_args=(
  --num_height 480
  --num_width 832
  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
 parallel_args=(
  --num_gpus $NUM_GPUS 
  --sp_size $NUM_GPUS 
-  --tp_size 1
+  --tp_size $NUM_GPUS
  --hsdp_replicate_dim 1
  --hsdp_shard_dim $NUM_GPUS
 )
@@ -1,11 +1,13 @@
 #!/bin/bash
 #SBATCH --job-name=t2v
 #SBATCH --partition=main
+#SBATCH --qos=hao
 #SBATCH --nodes=1
 #SBATCH --ntasks=1
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
+#SBATCH --nodelist=fs-mbz-gpu-[100-850]
 #SBATCH --mem=1440G
 #SBATCH --output=t2v_output/t2v_%j.out
 #SBATCH --error=t2v_output/t2v_%j.err
@@ -55,14 +57,13 @@ training_args=(
  --num_height 480
  --num_width 832
  --num_frames 77
-  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size 4
-  --tp_size 1
+  --tp_size 4
  --hsdp_replicate_dim 2
  --hsdp_shard_dim 4
 )
@@ -62,7 +62,6 @@ SystemEnv = namedtuple(
 DEFAULT_CONDA_PATTERNS = {
    "torch",
    "numpy",
-    "mypy"
    "cudatoolkit",
    "soumith",
    "mkl",
@@ -81,6 +80,7 @@ DEFAULT_CONDA_PATTERNS = {
 DEFAULT_PIP_PATTERNS = {
    "torch",
    "numpy",
+    "mypy",
    "flake8",
    "triton",
    "optree",
@@ -4,7 +4,7 @@
  "use_cpu_offload": false,
  "disable_autocast": false,
  "precision": "bf16",
-  "vae_precision": "fp32",
+  "vae_precision": "fp16",
  "vae_tiling": true,
  "vae_sp": true,
  "vae_config": {
@@ -1,5 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.v1.configs.models.dits.cosmos import CosmosConfig
 from fastvideo.v1.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.v1.configs.models.dits.stepvideo import StepVideoConfig
 from fastvideo.v1.configs.models.dits.wanvideo import WanVideoConfig

-__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig"]
+__all__ = [
+    "HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig", "CosmosConfig"
+]
@@ -11,9 +11,9 @@ from fastvideo.v1.platforms import AttentionBackendEnum
 class DiTArchConfig(ArchConfig):
    _fsdp_shard_conditions: list = field(default_factory=list)
    _compile_conditions: list = field(default_factory=list)
-    param_names_mapping: dict = field(default_factory=dict)
-    reverse_param_names_mapping: dict = field(default_factory=dict)
-    lora_param_names_mapping: dict = field(default_factory=dict)
+    _param_names_mapping: dict = field(default_factory=dict)
+    _reverse_param_names_mapping: dict = field(default_factory=dict)
+    _lora_param_names_mapping: dict = field(default_factory=dict)
    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
        AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,
        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_transformer_blocks(n: str, m) -> bool:
+    return "transformer_blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class CosmosArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_transformer_blocks])
+
+    _param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^patch_embed\.(.*)$": r"patch_embed.\1",
+            r"^time_embed\.time_proj\.(.*)$": r"time_embed.time_proj.\1",
+            r"^time_embed\.t_embedder\.(.*)$": r"time_embed.t_embedder.\1",
+            r"^time_embed\.norm\.(.*)$": r"time_embed.norm.\1",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
+            r"transformer_blocks.\1.attn1.norm_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
+            r"transformer_blocks.\1.attn1.norm_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_q\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_k\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_v\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.norm_q\.(.*)$":
+            r"transformer_blocks.\1.attn2.norm_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.norm_k\.(.*)$":
+            r"transformer_blocks.\1.attn2.norm_k.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$":
+            r"transformer_blocks.\1.ff.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$":
+            r"transformer_blocks.\1.ff.fc_out.\2",
+            r"^norm_out\.(.*)$": r"norm_out.\1",
+            r"^proj_out\.(.*)$": r"proj_out.\1",
+        })
+
+    _lora_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_out\.(.*)$":
+            r"transformer_blocks.\1.attn1.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_q\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_k\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_v\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_out\.(.*)$":
+            r"transformer_blocks.\1.attn2.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.(.*)$":
+            r"transformer_blocks.\1.ff.\2",
+        })
+
+    # Cosmos-specific config parameters based on transformer_cosmos.py
+    in_channels: int = 16
+    out_channels: int = 16
+    num_attention_heads: int = 16
+    attention_head_dim: int = 128
+    num_layers: int = 28
+    mlp_ratio: float = 4.0
+    text_embed_dim: int = 1024
+    adaln_lora_dim: int = 256
+    max_size: tuple[int, int, int] = (128, 240, 240)
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    rope_scale: tuple[float, float, float] = (1.0, 4.0, 4.0)
+    concat_padding_mask: bool = True
+    extra_pos_embed_type: str | None = None
+    qk_norm: str = "rms_norm"
+    eps: float = 1e-6
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.in_channels
+
+
+@dataclass
+class CosmosConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=CosmosArchConfig)
+    prefix: str = "Cosmos"
@@ -31,7 +31,7 @@ class HunyuanVideoArchConfig(DiTArchConfig):
    _compile_conditions: list = field(
        default_factory=lambda: [is_double_block, is_single_block, is_txt_in])

-    param_names_mapping: dict = field(
+    _param_names_mapping: dict = field(
        default_factory=lambda: {
            # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
            r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
@@ -146,8 +146,8 @@ class HunyuanVideoArchConfig(DiTArchConfig):
            r"final_layer.linear.\1",
        })

-    # Reverse mapping for saving checkpoints: custom -> hf
-    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+    # Reverse mapping for saving checkpoints: training -> diffusers
+    _reverse_param_names_mapping: dict = field(default_factory=lambda: {})

    patch_size: int = 2
    patch_size_t: int = 1
@@ -10,7 +10,7 @@ class StepVideoArchConfig(DiTArchConfig):
        default_factory=lambda:
        [lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit()])

-    param_names_mapping: dict = field(
+    _param_names_mapping: dict = field(
        default_factory=lambda: {
            # transformer block
            r"^transformer_blocks\.(\d+)\.norm1\.(weight|bias)$":
@@ -12,7 +12,7 @@ def is_blocks(n: str, m) -> bool:
 class WanVideoArchConfig(DiTArchConfig):
    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])

-    param_names_mapping: dict = field(
+    _param_names_mapping: dict = field(
        default_factory=lambda: {
            r"^patch_embedding\.(.*)$":
            r"patch_embedding.proj.\1",
@@ -52,12 +52,12 @@ class WanVideoArchConfig(DiTArchConfig):
            r"blocks.\1.self_attn_residual_norm.norm.\2",
        })

-    # Reverse mapping for saving checkpoints: custom -> hf
-    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+    # Reverse mapping for saving checkpoints: training -> diffusers
+    _reverse_param_names_mapping: dict = field(default_factory=lambda: {})

    # Some LoRA adapters use the original official layer names instead of hf layer names,
    # so apply this before the param_names_mapping
-    lora_param_names_mapping: dict = field(
+    _lora_param_names_mapping: dict = field(
        default_factory=lambda: {
            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2",
            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2",
@@ -62,11 +62,11 @@ class PipelineConfig:
    image_encoder_precision: str = "fp32"

    # Text encoder configuration
-    DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp32", )
+    DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp16", )
    text_encoder_configs: tuple[EncoderConfig, ...] = field(
        default_factory=lambda: (EncoderConfig(), ))
    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("fp32", ))
+        default_factory=lambda: ("fp16", ))
    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
        default_factory=lambda: (preprocess_text, ))
    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor],
@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.sample.base import CacheParams
+
+
+@dataclass
+class CosmosTeaCacheParams(CacheParams):
+    cache_type: str = "teacache"
+    teacache_thresh: float = 0.0
+    use_ret_steps: bool = True
+    ret_steps_coeffs: list[float] = field(default_factory=list)
+    non_ret_steps_coeffs: list[float] = field(default_factory=list)
+
+    @property
+    def coefficients(self) -> list[float]:
+        if self.use_ret_steps:
+            return self.ret_steps_coeffs
+        else:
+            return self.non_ret_steps_coeffs
+
+    @property
+    def ret_steps(self) -> int:
+        if self.use_ret_steps:
+            return 5 * 2
+        else:
+            return 1 * 2
+
+    def get_cutoff_steps(self, num_inference_steps: int) -> int:
+        if self.use_ret_steps:
+            return num_inference_steps * 2
+        else:
+            return num_inference_steps * 2 - 2
@@ -23,6 +23,7 @@ If you only need to use the distributed environment without model parallelism,
 you can skip the model parallel initialization and destruction steps.
 """
 import contextlib
+import gc
 import os
 import pickle
 import weakref
@@ -1015,6 +1016,15 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
    if shutdown_ray:
        import ray  # Lazy import Ray
        ray.shutdown()
+    gc.collect()
+    from fastvideo.v1.platforms import current_platform
+    if not current_platform.is_cpu():
+        torch.cuda.empty_cache()
+    try:
+        torch._C._host_emptyCache()
+    except AttributeError:
+        logger.warning(
+            "torch._C._host_emptyCache() only available in Pytorch >=2.5")


 def in_the_same_node_as(pg: ProcessGroup | StatelessProcessGroup,
@@ -6,6 +6,7 @@ This module provides a consolidated interface for generating videos using
 diffusion models.
 """

+import gc
 import math
 import os
 import time
@@ -70,7 +71,7 @@ class VideoGenerator:
        """
        # If users also provide some kwargs, it will override the FastVideoArgs and PipelineConfig.
        kwargs['model_path'] = model_path
-        fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
+        fastvideo_args = FastVideoArgs.from_kwargs(kwargs)

        return cls.from_fastvideo_args(fastvideo_args)

@@ -109,7 +110,6 @@ class VideoGenerator:
            prompt: The prompt to use for generation
            negative_prompt: The negative prompt to use (overrides the one in fastvideo_args)
            output_path: Path to save the video (overrides the one in fastvideo_args)
-            output_video_name: Name of the video file to save. Default is the first 100 characters of the prompt.
            save_video: Whether to save the video to disk
            return_frames: Whether to return the raw frames
            num_inference_steps: Number of denoising steps (overrides fastvideo_args)
@@ -229,7 +229,6 @@ class VideoGenerator:
            n_tokens=n_tokens,
            VSA_sparsity=fastvideo_args.VSA_sparsity,
            extra={},
-            output_video_name=kwargs.get("output_video_name", prompt[:100]),
        )

        # Run inference
@@ -253,8 +252,7 @@ class VideoGenerator:
            output_path = batch.output_path
            if output_path:
                os.makedirs(output_path, exist_ok=True)
-                video_path = os.path.join(output_path,
-                                          f"{batch.output_video_name}.mp4")
+                video_path = os.path.join(output_path, f"{prompt[:100]}.mp4")
                imageio.mimsave(video_path, frames, fps=batch.fps, format="mp4")
                logger.info("Saved video to %s", video_path)
            else:
@@ -270,9 +268,7 @@ class VideoGenerator:
                "generation_time": gen_time
            }

-    def set_lora_adapter(self,
-                         lora_nickname: str,
-                         lora_path: str | None = None) -> None:
+    def set_lora_adapter(self, lora_nickname: str, lora_path: str) -> None:
        self.executor.set_lora_adapter(lora_nickname, lora_path)

    def shutdown(self):
@@ -281,3 +277,5 @@ class VideoGenerator:
        """
        self.executor.shutdown()
        del self.executor
+        gc.collect()
+        torch.cuda.empty_cache()
@@ -280,7 +280,7 @@ class FastVideoArgs:
        return cls(**kwargs)  # type: ignore

    @classmethod
-    def from_kwargs(cls, **kwargs: Any) -> "FastVideoArgs":
+    def from_kwargs(cls, kwargs: dict[str, Any]) -> "FastVideoArgs":
        kwargs['pipeline_config'] = PipelineConfig.from_kwargs(kwargs)
        return cls(**kwargs)

@@ -292,7 +292,7 @@ class FastVideoArgs:
            assert self.sp_size != -1, "sp_size must be set for training"

        if self.tp_size == -1:
-            self.tp_size = 1
+            self.tp_size = self.num_gpus
        if self.sp_size == -1:
            self.sp_size = self.num_gpus
        if self.hsdp_shard_dim == -1:
@@ -305,6 +305,11 @@ class FastVideoArgs:
        if self.num_gpus < max(self.tp_size, self.sp_size):
            self.num_gpus = max(self.tp_size, self.sp_size)

+        if self.tp_size != self.sp_size:
+            raise ValueError(
+                f"tp_size ({self.tp_size}) must be equal to sp_size ({self.sp_size})"
+            )
+
        if self.enable_torch_compile and self.num_gpus > 1:
            logger.warning(
                "Currently torch compile does not work with multi-gpu. Setting enable_torch_compile to False"
@@ -3,12 +3,9 @@

 import torch
 from torch import nn
-from torch.distributed._composable.fsdp import (CPUOffloadPolicy, OffloadPolicy,
-                                                fully_shard)
-from torch.distributed.tensor import DTensor
+from torch.distributed.tensor import DTensor, distribute_tensor

-from fastvideo.v1.distributed import (get_local_torch_device, get_tp_rank,
-                                      split_tensor_along_last_dim,
+from fastvideo.v1.distributed import (get_tp_rank, split_tensor_along_last_dim,
                                      tensor_model_parallel_all_gather,
                                      tensor_model_parallel_all_reduce)
 from fastvideo.v1.layers.linear import (ColumnParallelLinear, LinearBase,
@@ -16,7 +13,6 @@ from fastvideo.v1.layers.linear import (ColumnParallelLinear, LinearBase,
                                        QKVParallelLinear, ReplicatedLinear,
                                        RowParallelLinear)
 from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from fastvideo.v1.utils import get_mixed_precision_state


 class BaseLayerWithLoRA(nn.Module):
@@ -30,11 +26,12 @@ class BaseLayerWithLoRA(nn.Module):
        self.lora_A: torch.Tensor = None
        self.lora_B: torch.Tensor = None
        self.merged: bool = False
+        self.weight = base_layer.weight
        self.cpu_weight = base_layer.weight.to("cpu")
+        self.unmerge_count = 0
        # indicates adapter weights don't contain this layer
        # (which shouldn't normally happen, but we want to separate it from the case of erroneous merging)
        self.disable_lora: bool = False
-        self.lora_path: str | None = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.base_layer.forward(x)
@@ -48,14 +45,12 @@ class BaseLayerWithLoRA(nn.Module):
    def set_lora_weights(self,
                         A: torch.Tensor,
                         B: torch.Tensor,
-                         training_mode: bool = False,
-                         lora_path: str | None = None) -> None:
+                         training_mode: bool = False) -> None:
        self.lora_A = A  # share storage with weights in the pipeline
        self.lora_B = B
        self.disable_lora = False
        if not training_mode:
            self.merge_lora_weights()
-        self.lora_path = lora_path

    @torch.no_grad()
    def merge_lora_weights(self) -> None:
@@ -63,44 +58,27 @@ class BaseLayerWithLoRA(nn.Module):
            return

        if self.merged:
-            self.unmerge_lora_weights()
+            raise ValueError(
+                "LoRA weights already merged. Please unmerge them first.")
        assert self.lora_A is not None and self.lora_B is not None, "LoRA weights not set. Please set them first."
        if isinstance(self.base_layer.weight, DTensor):
            mesh = self.base_layer.weight.data.device_mesh
-            # Using offload param is on CPU, so current_device is for "CPU -> GPU -> merge -> CPU"
+            placements = self.base_layer.weight.data.placements
            current_device = self.base_layer.weight.data.device
            data = self.base_layer.weight.data.to(
-                get_local_torch_device()).full_tensor()
-            data += (self.slice_lora_b_weights(self.lora_B).to(data)
-                     @ self.slice_lora_a_weights(self.lora_A).to(data))
-
-            # Must re-register updated weights for FSDP to recognize them
-            self.base_layer.weight = nn.Parameter(data.to(current_device))
-            if isinstance(getattr(self.base_layer, "bias", None), DTensor):
-                self.base_layer.bias = nn.Parameter(
-                    self.base_layer.bias.to(
-                        get_local_torch_device(),
-                        non_blocking=True).full_tensor().to(current_device))
-
-            offload_policy = CPUOffloadPolicy() if "cpu" in str(
-                current_device) else OffloadPolicy()
-            # see https://github.com/pytorch/torchtune/pull/2714/files#diff-909ee7ef184b0d834c40a1980ca4149afc38612ec7a4b344d8e2fc27641758c9R69-R79
-            # After the 1st forward, self.base_layer becomes a FSDP module and needs to be resharded
-            if hasattr(self.base_layer, "unshard"):
-                self.base_layer.unshard()
-            mp_policy = get_mixed_precision_state().mp_policy
-            fully_shard(self.base_layer,
-                        mesh=mesh,
-                        mp_policy=mp_policy,
-                        offload_policy=offload_policy)
+                f"cuda:{torch.cuda.current_device()}").full_tensor()
+            data += (self.slice_lora_b_weights(self.lora_B)
+                     @ self.slice_lora_a_weights(self.lora_A)).to(data)
+            self.base_layer.weight = nn.Parameter(
+                distribute_tensor(data, mesh,
+                                  placements=placements).to(current_device))
        else:
            current_device = self.base_layer.weight.data.device
-            data = self.base_layer.weight.data.to(get_local_torch_device())
+            data = self.base_layer.weight.to(
+                f"cuda:{torch.cuda.current_device()}")
            data += \
-                (self.slice_lora_b_weights(self.lora_B.to(data)) @ self.slice_lora_a_weights(self.lora_A.to(data)))
-            self.base_layer.weight.data = data.to(current_device,
-                                                  non_blocking=True)
-
+                (self.slice_lora_b_weights(self.lora_B) @ self.slice_lora_a_weights(self.lora_A)).to(data)
+            self.base_layer.weight = nn.Parameter(data.to(current_device))
        self.merged = True

    @torch.no_grad()
@@ -112,15 +90,28 @@ class BaseLayerWithLoRA(nn.Module):
            raise ValueError(
                "LoRA weights not merged. Please merge them first before unmerging."
            )
+        self.unmerge_count += 1

-        # To avoid precision loss we do not subtract the LoRA weights here
-        if isinstance(self.base_layer.weight, DTensor):
-            device = self.base_layer.weight.data.device
-            self.base_layer.weight = nn.Parameter(self.cpu_weight.to(device))
-        else:
+        # Avoid precision loss
+        if self.unmerge_count % 3 == 0:
            self.base_layer.weight.data = self.cpu_weight.data.to(
                self.base_layer.weight)

+        if isinstance(self.base_layer.weight, DTensor):
+            mesh = self.base_layer.weight.data.device_mesh
+            placement = self.base_layer.weight.data.placements
+            device = self.base_layer.weight.data.device
+            data = self.base_layer.weight.data.to(
+                f"cuda:{torch.cuda.current_device()}").full_tensor()
+            data -= self.slice_lora_b_weights(
+                self.lora_B) @ self.slice_lora_a_weights(self.lora_A)
+            self.base_layer.weight = nn.Parameter(
+                distribute_tensor(data, mesh, placements=placement).to(device))
+        else:
+            self.base_layer.weight.data -= \
+                self.slice_lora_b_weights(self.lora_B) @\
+                self.slice_lora_a_weights(self.lora_A)
+
        self.merged = False


@@ -44,6 +44,60 @@ def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
    return x.flatten(-2)


+def apply_rotary_emb(
+    x: torch.Tensor,
+    freqs_cis: torch.Tensor | tuple[torch.Tensor],
+    use_real: bool = True,
+    use_real_unbind_dim: int = -1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+    to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+    reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+    tensors contain rotary embeddings and are returned as real tensors.
+
+    Args:
+        x (`torch.Tensor`):
+            Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
+        freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+    """
+    if use_real:
+        cos, sin = freqs_cis  # [S, D]
+        cos = cos[None, None]
+        sin = sin[None, None]
+        cos, sin = cos.to(x.device), sin.to(x.device)
+
+        if use_real_unbind_dim == -1:
+            # Used for flux, cogvideox, hunyuan-dit
+            x_real, x_imag = x.reshape(*x.shape[:-1], -1,
+                                       2).unbind(-1)  # [B, S, H, D//2]
+            x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+        elif use_real_unbind_dim == -2:
+            # Used for Stable Audio, OmniGen, CogView4 and Cosmos
+            x_real, x_imag = x.reshape(*x.shape[:-1], 2,
+                                       -1).unbind(-2)  # [B, S, H, D//2]
+            x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+        else:
+            raise ValueError(
+                f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2."
+            )
+
+        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+        return out
+    else:
+        # used for lumina
+        x_rotated = torch.view_as_complex(x.float().reshape(
+            *x.shape[:-1], -1, 2))
+        freqs_cis = freqs_cis.unsqueeze(2)
+        x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+        return x_out.type_as(x)
+
+
 def _apply_rotary_emb(
    x: torch.Tensor,
    cos: torch.Tensor,
@@ -173,3 +173,80 @@ def unpatchify(x, t, h, w, patch_size, channels) -> torch.Tensor:
    imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))

    return imgs
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1,
+    scale: float = 1,
+    max_period: int = 10000,
+) -> torch.Tensor:
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+
+    Args
+        timesteps (torch.Tensor):
+            a 1-D Tensor of N indices, one per batch element. These may be fractional.
+        embedding_dim (int):
+            the dimension of the output.
+        flip_sin_to_cos (bool):
+            Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
+        downscale_freq_shift (float):
+            Controls the delta between frequencies between dimensions
+        scale (float):
+            Scaling factor applied to the embeddings.
+        max_period (int):
+            Controls the maximum frequency of the embeddings
+    Returns
+        torch.Tensor: an [N x dim] Tensor of positional embeddings.
+    """
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+    half_dim = embedding_dim // 2
+    exponent = -math.log(max_period) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    # scale embeddings
+    emb = scale * emb
+
+    # concat sine and cosine embeddings
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+    # flip sine and cosine embeddings
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+    # zero pad
+    if embedding_dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+class Timesteps(nn.Module):
+
+    def __init__(self,
+                 num_channels: int,
+                 flip_sin_to_cos: bool,
+                 downscale_freq_shift: float,
+                 scale: int = 1):
+        super().__init__()
+        self.num_channels = num_channels
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+        self.scale = scale
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        t_emb = get_timestep_embedding(
+            timesteps,
+            self.num_channels,
+            flip_sin_to_cos=self.flip_sin_to_cos,
+            downscale_freq_shift=self.downscale_freq_shift,
+            scale=self.scale,
+        )
+        return t_emb
@@ -13,8 +13,8 @@ from fastvideo.v1.platforms import AttentionBackendEnum
 class BaseDiT(nn.Module, ABC):
    _fsdp_shard_conditions: list = []
    _compile_conditions: list = []
-    param_names_mapping: dict
-    reverse_param_names_mapping: dict
+    _param_names_mapping: dict
+    _reverse_param_names_mapping: dict
    hidden_size: int
    num_attention_heads: int
    num_channels_latents: int
@@ -24,7 +24,7 @@ class BaseDiT(nn.Module, ABC):

    def __init_subclass__(cls) -> None:
        required_class_attrs = [
-            "_fsdp_shard_conditions", "param_names_mapping",
+            "_fsdp_shard_conditions", "_param_names_mapping",
            "_compile_conditions"
        ]
        super().__init_subclass__()
@@ -78,9 +78,9 @@ class CachableDiT(BaseDiT):
    """
    # These are required class attributes that should be overridden by concrete implementations
    _fsdp_shard_conditions = []
-    param_names_mapping = {}
-    reverse_param_names_mapping = {}
-    lora_param_names_mapping: dict = {}
+    _param_names_mapping = {}
+    _reverse_param_names_mapping = {}
+    _lora_param_names_mapping: dict = {}
    # Ensure these instance attributes are properly defined in subclasses
    hidden_size: int
    num_attention_heads: int
@@ -0,0 +1,731 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from fastvideo.v1.attention import DistributedAttention, LocalAttention
+from fastvideo.v1.configs.models.dits.cosmos import CosmosConfig
+from fastvideo.v1.forward_context import get_forward_context
+from fastvideo.v1.layers.layernorm import RMSNorm
+from fastvideo.v1.layers.linear import ReplicatedLinear
+from fastvideo.v1.layers.mlp import MLP
+from fastvideo.v1.layers.rotary_embedding import apply_rotary_emb
+from fastvideo.v1.layers.visual_embedding import Timesteps
+from fastvideo.v1.models.dits.base import BaseDiT
+from fastvideo.v1.platforms import AttentionBackendEnum
+
+
+class CosmosPatchEmbed(nn.Module):
+
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 patch_size: tuple[int, int, int],
+                 bias: bool = True) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+
+        self.proj = nn.Linear(in_channels * patch_size[0] * patch_size[1] *
+                              patch_size[2],
+                              out_channels,
+                              bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        hidden_states = hidden_states.reshape(batch_size, num_channels,
+                                              num_frames // p_t, p_t,
+                                              height // p_h, p_h, width // p_w,
+                                              p_w)
+        hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5,
+                                              7).flatten(4, 7)
+        hidden_states = self.proj(hidden_states)
+        return hidden_states
+
+
+class CosmosTimestepEmbedding(nn.Module):
+
+    def __init__(self, in_features: int, out_features: int) -> None:
+        super().__init__()
+        self.linear_1 = nn.Linear(in_features, out_features, bias=False)
+        self.activation = nn.SiLU()
+        self.linear_2 = nn.Linear(out_features, 3 * out_features, bias=False)
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        emb = self.linear_1(timesteps)
+        emb = self.activation(emb)
+        emb = self.linear_2(emb)
+        return emb
+
+
+class CosmosEmbedding(nn.Module):
+
+    def __init__(self, embedding_dim: int, condition_dim: int) -> None:
+        super().__init__()
+
+        self.time_proj = Timesteps(embedding_dim,
+                                   flip_sin_to_cos=True,
+                                   downscale_freq_shift=0.0)
+        self.t_embedder = CosmosTimestepEmbedding(embedding_dim, condition_dim)
+        self.norm = RMSNorm(embedding_dim, eps=1e-6)
+
+    def forward(self, hidden_states: torch.Tensor,
+                timestep: torch.LongTensor) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep).type_as(hidden_states)
+        temb = self.t_embedder(timesteps_proj)
+        embedded_timestep = self.norm(timesteps_proj)
+        return temb, embedded_timestep
+
+
+class CosmosAdaLayerNorm(nn.Module):
+
+    def __init__(self, in_features: int, hidden_features: int) -> None:
+        super().__init__()
+        self.embedding_dim = in_features
+
+        self.activation = nn.SiLU()
+        self.norm = nn.LayerNorm(in_features,
+                                 elementwise_affine=False,
+                                 eps=1e-6)
+        self.linear_1 = nn.Linear(in_features, hidden_features, bias=False)
+        self.linear_2 = nn.Linear(hidden_features, 2 * in_features, bias=False)
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                embedded_timestep: torch.Tensor,
+                temb: torch.Tensor | None = None) -> torch.Tensor:
+        embedded_timestep = self.activation(embedded_timestep)
+        embedded_timestep = self.linear_1(embedded_timestep)
+        embedded_timestep = self.linear_2(embedded_timestep)
+
+        if temb is not None:
+            embedded_timestep = embedded_timestep + temb[..., :2 *
+                                                         self.embedding_dim]
+
+        shift, scale = embedded_timestep.chunk(2, dim=-1)
+        hidden_states = self.norm(hidden_states)
+
+        if embedded_timestep.ndim == 2:
+            shift, scale = (x.unsqueeze(1) for x in (shift, scale))
+
+        hidden_states = hidden_states * (1 + scale) + shift
+        return hidden_states
+
+
+class CosmosAdaLayerNormZero(nn.Module):
+
+    def __init__(self,
+                 in_features: int,
+                 hidden_features: int | None = None) -> None:
+        super().__init__()
+
+        self.norm = nn.LayerNorm(in_features,
+                                 elementwise_affine=False,
+                                 eps=1e-6)
+        self.activation = nn.SiLU()
+
+        if hidden_features is None:
+            self.linear_1 = nn.Identity()
+        else:
+            self.linear_1 = nn.Linear(in_features, hidden_features, bias=False)
+
+        self.linear_2 = nn.Linear(
+            hidden_features if hidden_features is not None else in_features,
+            3 * in_features,
+            bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        embedded_timestep: torch.Tensor,
+        temb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        embedded_timestep = self.activation(embedded_timestep)
+        embedded_timestep = self.linear_1(embedded_timestep)
+        embedded_timestep = self.linear_2(embedded_timestep)
+
+        if temb is not None:
+            embedded_timestep = embedded_timestep + temb
+
+        shift, scale, gate = embedded_timestep.chunk(3, dim=-1)
+        hidden_states = self.norm(hidden_states)
+
+        if embedded_timestep.ndim == 2:
+            shift, scale, gate = (x.unsqueeze(1) for x in (shift, scale, gate))
+
+        hidden_states = hidden_states * (1 + scale) + shift
+        return hidden_states, gate
+
+
+class CosmosSelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 qk_norm=True,
+                 eps=1e-6,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
+                 | None = None,
+                 prefix: str = "") -> None:
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.to_q = ReplicatedLinear(dim, dim, bias=False)
+        self.to_k = ReplicatedLinear(dim, dim, bias=False)
+        self.to_v = ReplicatedLinear(dim, dim, bias=False)
+        self.to_out = ReplicatedLinear(dim, dim, bias=False)
+        self.norm_q = RMSNorm(self.head_dim,
+                              eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = RMSNorm(self.head_dim,
+                              eps=eps) if qk_norm else nn.Identity()
+
+        # Attention mechanism
+        self.attn = DistributedAttention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+            prefix=prefix)
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: torch.Tensor | None = None,
+                attention_mask: torch.Tensor | None = None,
+                image_rotary_emb: torch.Tensor | None = None) -> torch.Tensor:
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+
+        # Get QKV
+        query, _ = self.to_q(hidden_states)
+        key, _ = self.to_k(encoder_hidden_states)
+        value, _ = self.to_v(encoder_hidden_states)
+
+        # Reshape for multi-head attention
+        query = query.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
+        key = key.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
+        value = value.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
+
+        # Apply normalization
+        if self.norm_q is not None:
+            query = self.norm_q.forward_native(query)
+        if self.norm_k is not None:
+            key = self.norm_k.forward_native(key)
+
+        # Apply RoPE if provided
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query,
+                                     image_rotary_emb,
+                                     use_real=True,
+                                     use_real_unbind_dim=-2)
+            key = apply_rotary_emb(key,
+                                   image_rotary_emb,
+                                   use_real=True,
+                                   use_real_unbind_dim=-2)
+
+        # Attention computation
+        attn_output, _ = self.attn(query, key, value)
+        # attn_output = attn_output.flatten(2)
+        attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query)
+
+        # Output projection
+        attn_output, _ = self.to_out(attn_output)
+        return attn_output
+
+
+class CosmosCrossAttention(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 cross_attention_dim: int,
+                 num_heads: int,
+                 qk_norm=True,
+                 eps=1e-6,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
+                 | None = None,
+                 prefix: str = "") -> None:
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.cross_attention_dim = cross_attention_dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.to_q = ReplicatedLinear(dim, dim, bias=False)
+        self.to_k = ReplicatedLinear(cross_attention_dim, dim, bias=False)
+        self.to_v = ReplicatedLinear(cross_attention_dim, dim, bias=False)
+        self.to_out = ReplicatedLinear(dim, dim, bias=False)
+        self.norm_q = RMSNorm(self.head_dim,
+                              eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = RMSNorm(self.head_dim,
+                              eps=eps) if qk_norm else nn.Identity()
+
+        # Attention mechanism
+        self.attn = LocalAttention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=supported_attention_backends)
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: torch.Tensor,
+                attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+
+        # Get QKV
+        query, _ = self.to_q(hidden_states)
+        key, _ = self.to_k(encoder_hidden_states)
+        value, _ = self.to_v(encoder_hidden_states)
+
+        # Reshape for multi-head attention
+        query = query.unflatten(2, (self.num_heads, -1))
+        key = key.unflatten(2, (self.num_heads, -1))
+        value = value.unflatten(2, (self.num_heads, -1))
+
+        # Apply normalization
+        if self.norm_q is not None:
+            query = self.norm_q.forward_native(query)
+        if self.norm_k is not None:
+            key = self.norm_k.forward_native(key)
+
+        # Attention computation
+        attn_output = self.attn(query, key, value)
+        attn_output = attn_output.flatten(2, 3).type_as(query)
+
+        # Output projection
+        attn_output, _ = self.to_out(attn_output)
+        return attn_output
+
+
+class CosmosTransformerBlock(nn.Module):
+
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        cross_attention_dim: int,
+        mlp_ratio: float = 4.0,
+        adaln_lora_dim: int = 256,
+        qk_norm: str = "rms_norm",
+        out_bias: bool = False,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...]
+        | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        hidden_size = num_attention_heads * attention_head_dim
+
+        self.norm1 = CosmosAdaLayerNormZero(in_features=hidden_size,
+                                            hidden_features=adaln_lora_dim)
+        self.attn1 = CosmosSelfAttention(
+            dim=hidden_size,
+            num_heads=num_attention_heads,
+            qk_norm=(qk_norm == "rms_norm"),
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn1")
+
+        self.norm2 = CosmosAdaLayerNormZero(in_features=hidden_size,
+                                            hidden_features=adaln_lora_dim)
+        self.attn2 = CosmosCrossAttention(
+            dim=hidden_size,
+            cross_attention_dim=cross_attention_dim,
+            num_heads=num_attention_heads,
+            qk_norm=(qk_norm == "rms_norm"),
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn2")
+
+        self.norm3 = CosmosAdaLayerNormZero(in_features=hidden_size,
+                                            hidden_features=adaln_lora_dim)
+        self.ff = MLP(hidden_size,
+                      int(hidden_size * mlp_ratio),
+                      act_type="gelu",
+                      bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        embedded_timestep: torch.Tensor,
+        temb: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        extra_pos_emb: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if extra_pos_emb is not None:
+            hidden_states = hidden_states + extra_pos_emb
+
+        # 1. Self Attention
+        norm_hidden_states, gate = self.norm1(hidden_states, embedded_timestep,
+                                              temb)
+        attn_output = self.attn1(norm_hidden_states,
+                                 image_rotary_emb=image_rotary_emb)
+        hidden_states = hidden_states + gate * attn_output
+
+        # 2. Cross Attention
+        norm_hidden_states, gate = self.norm2(hidden_states, embedded_timestep,
+                                              temb)
+        attn_output = self.attn2(norm_hidden_states,
+                                 encoder_hidden_states=encoder_hidden_states,
+                                 attention_mask=attention_mask)
+        hidden_states = hidden_states + gate * attn_output
+
+        # 3. Feed Forward
+        norm_hidden_states, gate = self.norm3(hidden_states, embedded_timestep,
+                                              temb)
+        ff_output = self.ff(norm_hidden_states)
+        hidden_states = hidden_states + gate * ff_output
+
+        return hidden_states
+
+
+class CosmosRotaryPosEmbed(nn.Module):
+
+    def __init__(
+            self,
+            hidden_size: int,
+            max_size: tuple[int, int, int] = (128, 240, 240),
+            patch_size: tuple[int, int, int] = (1, 2, 2),
+            base_fps: int = 24,
+            rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0),
+    ) -> None:
+        super().__init__()
+
+        self.max_size = [
+            size // patch
+            for size, patch in zip(max_size, patch_size, strict=False)
+        ]
+        self.patch_size = patch_size
+        self.base_fps = base_fps
+
+        self.dim_h = hidden_size // 6 * 2
+        self.dim_w = hidden_size // 6 * 2
+        self.dim_t = hidden_size - self.dim_h - self.dim_w
+
+        self.h_ntk_factor = rope_scale[1]**(self.dim_h / (self.dim_h - 2))
+        self.w_ntk_factor = rope_scale[2]**(self.dim_w / (self.dim_w - 2))
+        self.t_ntk_factor = rope_scale[0]**(self.dim_t / (self.dim_t - 2))
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        pe_size = [
+            num_frames // self.patch_size[0], height // self.patch_size[1],
+            width // self.patch_size[2]
+        ]
+        device = hidden_states.device
+
+        h_theta = 10000.0 * self.h_ntk_factor
+        w_theta = 10000.0 * self.w_ntk_factor
+        t_theta = 10000.0 * self.t_ntk_factor
+
+        seq = torch.arange(max(self.max_size),
+                           device=device,
+                           dtype=torch.float32)
+        dim_h_range = (
+            torch.arange(0, self.dim_h, 2, device=device,
+                         dtype=torch.float32)[:(self.dim_h // 2)] / self.dim_h)
+        dim_w_range = (
+            torch.arange(0, self.dim_w, 2, device=device,
+                         dtype=torch.float32)[:(self.dim_w // 2)] / self.dim_w)
+        dim_t_range = (
+            torch.arange(0, self.dim_t, 2, device=device,
+                         dtype=torch.float32)[:(self.dim_t // 2)] / self.dim_t)
+        h_spatial_freqs = 1.0 / (h_theta**dim_h_range)
+        w_spatial_freqs = 1.0 / (w_theta**dim_w_range)
+        temporal_freqs = 1.0 / (t_theta**dim_t_range)
+
+        emb_h = torch.outer(seq[:pe_size[1]],
+                            h_spatial_freqs)[None, :, None, :].repeat(
+                                pe_size[0], 1, pe_size[2], 1)
+        emb_w = torch.outer(seq[:pe_size[2]],
+                            w_spatial_freqs)[None, None, :, :].repeat(
+                                pe_size[0], pe_size[1], 1, 1)
+
+        # Apply sequence scaling in temporal dimension
+        if fps is None:
+            # Images
+            emb_t = torch.outer(seq[:pe_size[0]], temporal_freqs)
+        else:
+            # Videos
+            emb_t = torch.outer(seq[:pe_size[0]] / fps * self.base_fps,
+                                temporal_freqs)
+
+        emb_t = emb_t[:, None, None, :].repeat(1, pe_size[1], pe_size[2], 1)
+        freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1).flatten(0,
+                                                                     2).float()
+        cos = torch.cos(freqs)
+        sin = torch.sin(freqs)
+        return cos, sin
+
+
+class CosmosLearnablePositionalEmbed(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        max_size: tuple[int, int, int],
+        patch_size: tuple[int, int, int],
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+
+        self.max_size = [
+            size // patch
+            for size, patch in zip(max_size, patch_size, strict=False)
+        ]
+        self.patch_size = patch_size
+        self.eps = eps
+
+        self.pos_emb_t = nn.Parameter(torch.zeros(self.max_size[0],
+                                                  hidden_size))
+        self.pos_emb_h = nn.Parameter(torch.zeros(self.max_size[1],
+                                                  hidden_size))
+        self.pos_emb_w = nn.Parameter(torch.zeros(self.max_size[2],
+                                                  hidden_size))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        pe_size = [
+            num_frames // self.patch_size[0], height // self.patch_size[1],
+            width // self.patch_size[2]
+        ]
+
+        emb_t = self.pos_emb_t[:pe_size[0]][None, :, None, None, :].repeat(
+            batch_size, 1, pe_size[1], pe_size[2], 1)
+        emb_h = self.pos_emb_h[:pe_size[1]][None, None, :, None, :].repeat(
+            batch_size, pe_size[0], 1, pe_size[2], 1)
+        emb_w = self.pos_emb_w[:pe_size[2]][None, None, None, :, :].repeat(
+            batch_size, pe_size[0], pe_size[1], 1, 1)
+        emb = emb_t + emb_h + emb_w
+        emb = emb.flatten(1, 3)
+
+        norm = torch.linalg.vector_norm(emb,
+                                        dim=-1,
+                                        keepdim=True,
+                                        dtype=torch.float32)
+        norm = torch.add(self.eps,
+                         norm,
+                         alpha=np.sqrt(norm.numel() / emb.numel()))
+        return (emb / norm).type_as(hidden_states)
+
+
+class CosmosTransformer3DModel(BaseDiT):
+    _fsdp_shard_conditions = CosmosConfig()._fsdp_shard_conditions
+    _compile_conditions = CosmosConfig()._compile_conditions
+    _supported_attention_backends = CosmosConfig()._supported_attention_backends
+    _param_names_mapping = CosmosConfig()._param_names_mapping
+    _lora_param_names_mapping = CosmosConfig()._lora_param_names_mapping
+
+    def __init__(self, config: CosmosConfig, hf_config: dict[str, Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        inner_dim = config.num_attention_heads * config.attention_head_dim
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.in_channels = config.in_channels
+        self.out_channels = config.out_channels
+        self.num_channels_latents = config.num_channels_latents
+        self.patch_size = config.patch_size
+        self.max_size = config.max_size
+        self.rope_scale = config.rope_scale
+        self.concat_padding_mask = config.concat_padding_mask
+        self.extra_pos_embed_type = config.extra_pos_embed_type
+
+        # 1. Patch Embedding
+        patch_embed_in_channels = config.in_channels + 1 if config.concat_padding_mask else config.in_channels
+        self.patch_embed = CosmosPatchEmbed(patch_embed_in_channels,
+                                            inner_dim,
+                                            config.patch_size,
+                                            bias=False)
+
+        # 2. Positional Embedding
+        self.rope = CosmosRotaryPosEmbed(hidden_size=config.attention_head_dim,
+                                         max_size=config.max_size,
+                                         patch_size=config.patch_size,
+                                         rope_scale=config.rope_scale)
+
+        self.learnable_pos_embed = None
+        if config.extra_pos_embed_type == "learnable":
+            self.learnable_pos_embed = CosmosLearnablePositionalEmbed(
+                hidden_size=inner_dim,
+                max_size=config.max_size,
+                patch_size=config.patch_size,
+            )
+
+        # 3. Time Embedding
+        self.time_embed = CosmosEmbedding(inner_dim, inner_dim)
+
+        # 4. Transformer Blocks
+        self.transformer_blocks = nn.ModuleList([
+            CosmosTransformerBlock(
+                num_attention_heads=config.num_attention_heads,
+                attention_head_dim=config.attention_head_dim,
+                cross_attention_dim=config.text_embed_dim,
+                mlp_ratio=config.mlp_ratio,
+                adaln_lora_dim=config.adaln_lora_dim,
+                qk_norm=config.qk_norm,
+                out_bias=False,
+                supported_attention_backends=self._supported_attention_backends,
+                prefix=f"{config.prefix}.transformer_blocks.{i}",
+            ) for i in range(config.num_layers)
+        ])
+
+        # 5. Output norm & projection
+        self.norm_out = CosmosAdaLayerNorm(inner_dim, config.adaln_lora_dim)
+        self.proj_out = nn.Linear(inner_dim,
+                                  config.out_channels *
+                                  math.prod(config.patch_size),
+                                  bias=False)
+
+        self.gradient_checkpointing = False
+
+        # For TeaCache
+        self.previous_e0_even = None
+        self.previous_e0_odd = None
+        self.previous_residual_even = None
+        self.previous_residual_odd = None
+        self.is_even = True
+        self.should_calc_even = True
+        self.should_calc_odd = True
+        self.accumulated_rel_l1_distance_even = 0
+        self.accumulated_rel_l1_distance_odd = 0
+        self.cnt = 0
+        self.__post_init__()
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                timestep: torch.Tensor,
+                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+                attention_mask: torch.Tensor | None = None,
+                fps: int | None = None,
+                condition_mask: torch.Tensor | None = None,
+                padding_mask: torch.Tensor | None = None,
+                **kwargs) -> torch.Tensor:
+        forward_batch = get_forward_context().forward_batch
+        enable_teacache = forward_batch is not None and forward_batch.enable_teacache
+
+        orig_dtype = hidden_states.dtype
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+
+        # 1. Concatenate padding mask if needed & prepare attention mask
+        if condition_mask is not None:
+            hidden_states = torch.cat([hidden_states, condition_mask], dim=1)
+
+        if self.concat_padding_mask and padding_mask is not None:
+            from torchvision import transforms
+            padding_mask = transforms.functional.resize(
+                padding_mask,
+                list(hidden_states.shape[-2:]),
+                interpolation=transforms.InterpolationMode.NEAREST)
+            hidden_states = torch.cat([
+                hidden_states,
+                padding_mask.unsqueeze(2).repeat(batch_size, 1, num_frames, 1,
+                                                 1)
+            ],
+                                      dim=1)
+            # # Resize padding mask to match hidden states spatial dimensions
+            # padding_mask_resized = F.interpolate(
+            #     padding_mask.float().unsqueeze(1),
+            #     size=(height, width),
+            #     mode='nearest'
+            # ).squeeze(1)
+            # hidden_states = torch.cat(
+            #     [hidden_states, padding_mask_resized.unsqueeze(1).unsqueeze(2).repeat(1, 1, num_frames, 1, 1)], dim=1
+            # )
+
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(
+                1)  # [B, 1, 1, S]
+
+        # 2. Generate positional embeddings
+        image_rotary_emb = self.rope(hidden_states, fps=fps)
+        extra_pos_emb = self.learnable_pos_embed(
+            hidden_states) if self.extra_pos_embed_type == "learnable" else None
+
+        # 3. Patchify input
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+        hidden_states = self.patch_embed(hidden_states)
+        hidden_states = hidden_states.flatten(
+            1, 3)  # [B, T, H, W, C] -> [B, THW, C] codespell:ignore
+
+        # 4. Timestep embeddings
+        if timestep.ndim == 1:
+            temb, embedded_timestep = self.time_embed(hidden_states, timestep)
+        elif timestep.ndim == 5:
+            assert timestep.shape == (batch_size, 1, num_frames, 1, 1), (
+                f"Expected timestep to have shape [B, 1, T, 1, 1], but got {timestep.shape}"
+            )
+            timestep = timestep.flatten()
+            temb, embedded_timestep = self.time_embed(hidden_states, timestep)
+            # We can do this because num_frames == post_patch_num_frames, as p_t is 1
+            temb, embedded_timestep = (
+                x.view(batch_size, post_patch_num_frames, 1, 1,
+                       -1).expand(-1, -1, post_patch_height, post_patch_width,
+                                  -1).flatten(1, 3)
+                for x in (temb, embedded_timestep)
+            )  # [BT, C] -> [B, T, 1, 1, C] -> [B, T, H, W, C] -> [B, THW, C] codespell:ignore
+        else:
+            raise ValueError(f"Unsupported timestep shape: {timestep.shape}")
+
+        # 6. Transformer blocks
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for block in self.transformer_blocks:
+                hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    embedded_timestep,
+                    temb,
+                    image_rotary_emb,
+                    extra_pos_emb,
+                    attention_mask,
+                )
+        else:
+            for block in self.transformer_blocks:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    embedded_timestep=embedded_timestep,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    extra_pos_emb=extra_pos_emb,
+                    attention_mask=attention_mask,
+                )
+
+        # 7. Output norm & projection & unpatchify
+        hidden_states = self.norm_out(hidden_states, embedded_timestep, temb)
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1))
+        hidden_states = hidden_states.unflatten(
+            1, (post_patch_num_frames, post_patch_height, post_patch_width))
+        # NOTE: The permutation order here is not the inverse operation of what happens when patching as usually expected.
+        # It might be a source of confusion to the reader, but this is correct
+        hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5)
+        hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        return hidden_states
@@ -441,10 +441,10 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
    _compile_conditions = HunyuanVideoConfig()._compile_conditions
    _supported_attention_backends = HunyuanVideoConfig(
    )._supported_attention_backends
-    param_names_mapping = HunyuanVideoConfig().param_names_mapping
-    reverse_param_names_mapping = HunyuanVideoConfig(
-    ).reverse_param_names_mapping
-    lora_param_names_mapping = HunyuanVideoConfig().lora_param_names_mapping
+    _param_names_mapping = HunyuanVideoConfig()._param_names_mapping
+    _reverse_param_names_mapping = HunyuanVideoConfig(
+    )._reverse_param_names_mapping
+    _lora_param_names_mapping = HunyuanVideoConfig()._lora_param_names_mapping

    def __init__(self, config: HunyuanVideoConfig, hf_config: dict[str, Any]):
        super().__init__(config=config, hf_config=hf_config)
@@ -457,13 +457,11 @@ class StepVideoTransformerBlock(nn.Module):

 class StepVideoModel(BaseDiT):
    # (Optional) Keep the same attribute for compatibility with splitting, etc.
-    _fsdp_shard_conditions = [
-        lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit(),
-        # lambda n, m: "pos_embed" in n  # If needed for the patch embedding.
-    ]
-    param_names_mapping = StepVideoConfig().param_names_mapping
-    reverse_param_names_mapping = StepVideoConfig().reverse_param_names_mapping
-    lora_param_names_mapping = StepVideoConfig().lora_param_names_mapping
+    _fsdp_shard_conditions = StepVideoConfig()._fsdp_shard_conditions
+    _param_names_mapping = StepVideoConfig()._param_names_mapping
+    _reverse_param_names_mapping = StepVideoConfig(
+    )._reverse_param_names_mapping
+    _lora_param_names_mapping = StepVideoConfig()._lora_param_names_mapping
    _supported_attention_backends = StepVideoConfig(
    )._supported_attention_backends

@@ -515,9 +515,9 @@ class WanTransformer3DModel(CachableDiT):
    _compile_conditions = WanVideoConfig()._compile_conditions
    _supported_attention_backends = WanVideoConfig(
    )._supported_attention_backends
-    param_names_mapping = WanVideoConfig().param_names_mapping
-    reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
-    lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping
+    _param_names_mapping = WanVideoConfig()._param_names_mapping
+    _reverse_param_names_mapping = WanVideoConfig()._reverse_param_names_mapping
+    _lora_param_names_mapping = WanVideoConfig()._lora_param_names_mapping

    def __init__(self, config: WanVideoConfig, hf_config: dict[str,
                                                               Any]) -> None:
@@ -429,6 +429,7 @@ class TransformerLoader(ComponentLoader):
            hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
            cpu_offload=fastvideo_args.use_cpu_offload,
            fsdp_inference=fastvideo_args.use_fsdp_inference,
+            default_dtype=default_dtype,
            # TODO(will): make these configurable
            param_dtype=torch.bfloat16,
            reduce_dtype=torch.float32,
@@ -5,6 +5,7 @@
 # Copyright 2025 The FastVideo Authors.

 import contextlib
+from collections import defaultdict
 from collections.abc import Callable, Generator
 from itertools import chain
 from typing import Any
@@ -18,8 +19,7 @@ from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule,
 from torch.nn.modules.module import _IncompatibleKeys

 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.loader.utils import (get_param_names_mapping,
-                                              hf_to_custom_state_dict)
+from fastvideo.v1.models.loader.utils import get_param_names_mapping
 from fastvideo.v1.models.loader.weight_utils import safetensors_weights_iterator
 from fastvideo.v1.utils import set_mixed_precision_policy

@@ -62,6 +62,7 @@ def maybe_load_fsdp_model(
    device: torch.device,
    hsdp_replicate_dim: int,
    hsdp_shard_dim: int,
+    default_dtype: torch.dtype,
    param_dtype: torch.dtype,
    reduce_dtype: torch.dtype,
    cpu_offload: bool = False,
@@ -80,14 +81,12 @@ def maybe_load_fsdp_model(
                                     output_dtype,
                                     cast_forward_inputs=False)

-    set_mixed_precision_policy(
-        param_dtype=param_dtype,
-        reduce_dtype=reduce_dtype,
-        output_dtype=output_dtype,
-        mp_policy=mp_policy,
-    )
+    set_mixed_precision_policy(master_dtype=default_dtype,
+                               param_dtype=param_dtype,
+                               reduce_dtype=reduce_dtype,
+                               output_dtype=output_dtype)

-    with set_default_dtype(param_dtype), torch.device("meta"):
+    with set_default_dtype(default_dtype), torch.device("meta"):
        model = model_cls(**init_params)
    world_size = hsdp_replicate_dim * hsdp_shard_dim
    if not training_mode and not fsdp_inference:
@@ -107,8 +106,9 @@ def maybe_load_fsdp_model(
                fsdp_shard_conditions=model._fsdp_shard_conditions,
                pin_cpu_memory=pin_cpu_memory)

-    weight_iterator = safetensors_weights_iterator(weight_dir_list)
-    param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
+    weight_iterator = safetensors_weights_iterator(weight_dir_list,
+                                                   to_cpu=cpu_offload)
+    param_names_mapping_fn = get_param_names_mapping(model._param_names_mapping)
    load_model_from_full_model_state_dict(
        model,
        weight_iterator,
@@ -233,14 +233,36 @@ def load_model_from_full_model_state_dict(
        NotImplementedError: If got FSDP with more than 1D.
    """
    meta_sd = model.state_dict()
+    # Find new params
+    used_keys = set()
    sharded_sd = {}
-    custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
-        full_sd_iterator, param_names_mapping)  # type: ignore
-    for target_param_name, full_tensor in custom_param_sd.items():
+    to_merge_params: defaultdict[str, dict[Any, Any]] = defaultdict(dict)
+    reverse_param_names_mapping = {}
+    assert param_names_mapping is not None
+    for source_param_name, full_tensor in full_sd_iterator:
+        target_param_name, merge_index, num_params_to_merge = param_names_mapping(
+            source_param_name)
+        reverse_param_names_mapping[target_param_name] = (source_param_name,
+                                                          merge_index,
+                                                          num_params_to_merge)
+        used_keys.add(target_param_name)
+        if merge_index is not None:
+            to_merge_params[target_param_name][merge_index] = full_tensor
+            if len(to_merge_params[target_param_name]) == num_params_to_merge:
+                # cat at output dim according to the merge_index order
+                sorted_tensors = [
+                    to_merge_params[target_param_name][i]
+                    for i in range(num_params_to_merge)
+                ]
+                full_tensor = torch.cat(sorted_tensors, dim=0)
+                del to_merge_params[target_param_name]
+            else:
+                continue
+
        meta_sharded_param = meta_sd.get(target_param_name)
        if meta_sharded_param is None:
            raise ValueError(
-                f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect."
+                f"Parameter {source_param_name}-->{target_param_name} not found in meta sharded state dict"
            )
        if not hasattr(meta_sharded_param, "device_mesh"):
            full_tensor = full_tensor.to(device=device, dtype=param_dtype)
@@ -257,10 +279,10 @@ def load_model_from_full_model_state_dict(
                sharded_tensor = sharded_tensor.cpu()
        sharded_sd[target_param_name] = nn.Parameter(sharded_tensor)

-    model.reverse_param_names_mapping = reverse_param_names_mapping
-    unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys())
+    model._reverse_param_names_mapping = reverse_param_names_mapping
+    unused_keys = set(meta_sd.keys()) - used_keys
    if unused_keys:
-        logger.warning("Found unloaded parameters in meta state dict: %s",
+        logger.warning("Found new parameters in meta state dict: %s",
                       unused_keys)

    # List of allowed parameter name patterns
@@ -2,8 +2,7 @@
 """Utilities for selecting and loading models."""
 import contextlib
 import re
-from collections import defaultdict
-from collections.abc import Callable, Iterator
+from collections.abc import Callable
 from typing import Any

 import torch
@@ -36,6 +35,7 @@ def get_param_names_mapping(
    """

    def mapping_fn(name: str) -> tuple[str, Any, Any]:
+
        # Try to match and transform the name using the regex patterns in mapping_dict
        for pattern, replacement in mapping_dict.items():
            match = re.match(pattern, name)
@@ -52,46 +52,4 @@ def get_param_names_mapping(
        # If no pattern matches, return the original name
        return name, None, None

-    return mapping_fn
-
-
-def hf_to_custom_state_dict(
-    hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]],
-    param_names_mapping: Callable[[str], tuple[str, Any, Any]]
-) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]:
-    """
-    Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary.
-    
-    Args:
-        hf_param_sd (Dict[str, torch.Tensor]): The Hugging Face parameter state dictionary
-        param_names_mapping (Callable[[str], tuple[str, Any, Any]]): A function that maps parameter names from source to target format
-        
-    Returns:
-        custom_param_sd (Dict[str, torch.Tensor]): The custom formatted parameter state dict
-        reverse_param_names_mapping (Dict[str, Tuple[str, Any, Any]]): Maps back from custom to hf
-    """
-    custom_param_sd = {}
-    to_merge_params = defaultdict(dict)  # type: ignore
-    reverse_param_names_mapping = {}
-    if isinstance(hf_param_sd, dict):
-        hf_param_sd = hf_param_sd.items()  # type: ignore
-    for source_param_name, full_tensor in hf_param_sd:  # type: ignore
-        target_param_name, merge_index, num_params_to_merge = param_names_mapping(
-            source_param_name)
-        reverse_param_names_mapping[target_param_name] = (source_param_name,
-                                                          merge_index,
-                                                          num_params_to_merge)
-        if merge_index is not None:
-            to_merge_params[target_param_name][merge_index] = full_tensor
-            if len(to_merge_params[target_param_name]) == num_params_to_merge:
-                # cat at output dim according to the merge_index order
-                sorted_tensors = [
-                    to_merge_params[target_param_name][i]
-                    for i in range(num_params_to_merge)
-                ]
-                full_tensor = torch.cat(sorted_tensors, dim=0)
-                del to_merge_params[target_param_name]
-            else:
-                continue
-        custom_param_sd[target_param_name] = full_tensor
-    return custom_param_sd, reverse_param_names_mapping
+    return mapping_fn
@@ -23,7 +23,8 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
    "HunyuanVideoTransformer3DModel":
    ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
    "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
-    "StepVideoModel": ("dits", "stepvideo", "StepVideoModel")
+    "StepVideoModel": ("dits", "stepvideo", "StepVideoModel"),
+    "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel")
 }

 _IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -239,6 +239,7 @@ class ParallelTiledVAE(ABC):

        results = torch.cat(local_results, dim=0).contiguous()
        del local_results
+        torch.cuda.empty_cache()
        # first gather size to pad the results
        local_size = torch.tensor([results.size(0)],
                                  device=results.device,
@@ -252,7 +253,7 @@ class ParallelTiledVAE(ABC):
        padded_results = torch.zeros(max_size, device=results.device)
        padded_results[:results.size(0)] = results
        del results
-
+        torch.cuda.empty_cache()
        # Gather all results
        gathered_dim_metadata = [None] * world_size
        gathered_results = torch.zeros_like(padded_results).repeat(
@@ -113,7 +113,7 @@ class ComposedPipelineBase(ABC):
        if args is None or args.inference_mode:

            kwargs['model_path'] = model_path
-            fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
+            fastvideo_args = FastVideoArgs.from_kwargs(kwargs)
        else:
            assert args is not None, "args must be provided for training mode"
            fastvideo_args = TrainingArgs.from_cli_args(args)
@@ -7,7 +7,6 @@ import torch
 import torch.distributed as dist
 from safetensors.torch import load_file

-from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.layers.lora.linear import (BaseLayerWithLoRA, get_lora_layer,
                                             replace_submodule)
@@ -30,13 +29,13 @@ class LoRAPipeline(ComposedPipelineBase):
    lora_layers: dict[str, BaseLayerWithLoRA] = {}
    fastvideo_args: FastVideoArgs
    exclude_lora_layers: list[str] = []
-    device: torch.device | None = None
+    device: torch.device = torch.device(f"cuda:{torch.cuda.current_device()}")

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.exclude_lora_layers = self.modules[
            "transformer"].config.arch_config.exclude_lora_layers
-        self.device = get_local_torch_device()
+
        self.convert_to_lora_layers()
        if self.fastvideo_args.pipeline_config.lora_path is not None:
            self.set_lora_adapter(
@@ -54,6 +53,7 @@ class LoRAPipeline(ComposedPipelineBase):
        """
        Converts the transformer to a LoRA transformer.
        """
+
        for name, layer in self.modules["transformer"].named_modules():
            if not self.is_target_layer(name):
                continue
@@ -85,18 +85,16 @@ class LoRAPipeline(ComposedPipelineBase):
            raise ValueError(
                f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it."
            )
-
        adapter_updated = False
        rank = dist.get_rank()
        if lora_path is not None:
            lora_local_path = maybe_download_lora(lora_path)
-            lora_state_dict = load_file(lora_local_path,
-                                        device=str(self.device))
+            lora_state_dict = load_file(lora_local_path)
            # Map the hf layer names to our custom layer names
            param_names_mapping_fn = get_param_names_mapping(
-                self.modules["transformer"].param_names_mapping)
+                self.modules["transformer"]._param_names_mapping)
            lora_param_names_mapping_fn = get_param_names_mapping(
-                self.modules["transformer"].lora_param_names_mapping)
+                self.modules["transformer"]._lora_param_names_mapping)

            to_merge_params: defaultdict[Hashable,
                                         dict[Any, Any]] = defaultdict(dict)
@@ -121,11 +119,6 @@ class LoRAPipeline(ComposedPipelineBase):
                        del to_merge_params[target_name]
                    else:
                        continue
-
-                if target_name in self.lora_adapters[lora_nickname]:
-                    raise ValueError(
-                        f"Target name {target_name} already exists in lora_adapters[{lora_nickname}]"
-                    )
                self.lora_adapters[lora_nickname][target_name] = weight.to(
                    self.device)
            adapter_updated = True
@@ -141,11 +134,12 @@ class LoRAPipeline(ComposedPipelineBase):
            lora_B_name = name + ".lora_B"
            if lora_A_name in self.lora_adapters[lora_nickname]\
                and lora_B_name in self.lora_adapters[lora_nickname]:
+                if layer.merged:
+                    layer.unmerge_lora_weights()
                layer.set_lora_weights(
                    self.lora_adapters[lora_nickname][lora_A_name],
                    self.lora_adapters[lora_nickname][lora_B_name],
-                    training_mode=self.fastvideo_args.training_mode,
-                    lora_path=lora_path)
+                    training_mode=self.fastvideo_args.training_mode)
                adapted_count += 1
            else:
                if rank == 0:
@@ -155,5 +149,4 @@ class LoRAPipeline(ComposedPipelineBase):
                layer.disable_lora = True
        logger.info("Rank %d: LoRA adapter %s applied to %d layers", rank,
                    lora_path, adapted_count)
-
        self.cur_adapter_name = lora_nickname
@@ -46,7 +46,7 @@ class ForwardBatch:
    negative_prompt: str | list[str] | None = None
    prompt_path: str | None = None
    output_path: str = "outputs/"
-    output_video_name: str | None = None
+
    # Primary encoder embeddings
    prompt_embeds: list[torch.Tensor] = field(default_factory=list)
    negative_prompt_embeds: list[torch.Tensor] | None = None
@@ -136,6 +136,7 @@ class EncodingStage(PipelineStage):
            self.maybe_free_model_hooks()

        self.vae.to("cpu")
+        torch.cuda.empty_cache()

        return batch

@@ -5,6 +5,8 @@ Image encoding stages for I2V diffusion pipelines.
 This module contains implementations of image encoding stages for diffusion pipelines.
 """

+import torch
+
 from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
@@ -66,6 +68,7 @@ class ImageEncodingStage(PipelineStage):

        if fastvideo_args.use_cpu_offload:
            self.image_encoder.to('cpu')
+            torch.cuda.empty_cache()

        return batch

@@ -81,6 +81,7 @@ class TextEncodingStage(PipelineStage):
                    output_hidden_states=True,
                )
            prompt_embeds = postprocess_func(outputs)
+
            batch.prompt_embeds.append(prompt_embeds)
            if batch.prompt_attention_mask is not None:
                batch.prompt_attention_mask.append(attention_mask)
@@ -1,195 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import json
-import os
-
-import pytest
-
-from fastvideo import VideoGenerator
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.tests.utils import compute_video_ssim_torchvision, write_ssim_results
-from diffusers import DiffusionPipeline
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.pipelines import build_pipeline
-from fastvideo.v1.models.loader.utils import hf_to_custom_state_dict, get_param_names_mapping
-from torch.testing import assert_close
-from torch.distributed.tensor import DTensor
-from fastvideo.v1.worker import MultiprocExecutor
-import torch
-logger = init_logger(__name__)
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29500"
-
-# Base parameters for LoRA inference tests
-WAN_LORA_PARAMS = {
-    "num_gpus": 1,
-    "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-    "height": 480,
-    "width": 832,
-    "num_frames": 45,
-    "num_inference_steps": 32,
-    "guidance_scale": 5.0,
-    "flow_shift": 3.0,
-    "seed": 42,
-    "fps": 24,
-    "neg_prompt": "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
-    "text-encoder-precision": ("fp32",),
-    "use_cpu_offload": True,
-}
-
-# LoRA configurations for testing
-LORA_CONFIGS = [
-    {
-        "lora_path": "benjamin-paine/steamboat-willie-1.3b",
-        "lora_nickname": "steamboat",
-        "prompt": "steamboat willie style, golden era animation, close-up of a short fluffy monster kneeling beside a melting red candle. the mood is one of wonder and curiosity, as the monster gazes at the flame with wide eyes and open mouth. Its pose and expression convey a sense of innocence and playfulness, as if it is exploring the world around it for the first time. The use of warm colors and dramatic lighting further enhances the cozy atmosphere of the image.",
-        "negative_prompt": "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
-        "ssim_threshold": 0.79
-    },
-    # {
-    #     "lora_path": "motimalu/wan-flat-color-1.3b-v2",
-    #     "lora_nickname": "flat_color",
-    #     "prompt": "flat color, no lineart, blending, negative space, artist:[john kafka|ponsuke kaikai|hara id 21|yoneyama mai|fuzichoco], 1girl, sakura miko, pink hair, cowboy shot, white shirt, floral print, off shoulder, outdoors, cherry blossom, tree shade, wariza, looking up, falling petals, half-closed eyes, white sky, clouds, live2d animation, upper body, high quality cinematic video of a woman sitting under a sakura tree. Dreamy and lonely, the camera close-ups on the face of the woman as she turns towards the viewer. The Camera is steady, This is a cowboy shot. The animation is smooth and fluid.",
-    #     "negative_prompt": "bad quality video,色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
-    #     "ssim_threshold": 0.79
-    # }
-]
-
-MODEL_TO_PARAMS = {
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WAN_LORA_PARAMS,
-}
-
-@pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys()))
-def test_merge_lora_weights(model_id):
-    lora_config = LORA_CONFIGS[0] # test only one
-    hf_pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
-    hf_pipe.enable_model_cpu_offload()
-
-    lora_nickname = lora_config["lora_nickname"]
-    lora_path = lora_config["lora_path"]
-    args = FastVideoArgs.from_kwargs(
-        model_path=model_id,
-        use_cpu_offload=True,
-        dit_precision="bf16",
-    )
-    pipe = build_pipeline(args)
-    pipe.set_lora_adapter(lora_nickname, lora_path)
-    custom_transformer = pipe.modules["transformer"]
-    custom_state_dict = custom_transformer.state_dict()
-
-    hf_pipe.load_lora_weights(lora_path, adapter_name=lora_nickname)
-    for name, layer in hf_pipe.transformer.named_modules():
-        if hasattr(layer, "unmerge"):
-            layer.unmerge()
-            layer.merge(adapter_names=[lora_nickname])
-
-    hf_transformer = hf_pipe.transformer
-    param_names_mapping = get_param_names_mapping(custom_transformer.param_names_mapping)
-    hf_state_dict, _ = hf_to_custom_state_dict(hf_transformer.state_dict(), param_names_mapping)
-    for key in hf_state_dict.keys():
-        if "base_layer" not in key:
-            continue
-        hf_param = hf_state_dict[key]
-        custom_param = custom_state_dict[key].to_local() if isinstance(custom_state_dict[key], DTensor) else custom_state_dict[key]
-        assert_close(hf_param, custom_param, atol=7e-4, rtol=7e-4)
-
-@pytest.mark.parametrize("ATTENTION_BACKEND", ["TORCH_SDPA"])
-@pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys()))
-def test_lora_inference_similarity(ATTENTION_BACKEND, model_id):
-    """
-    Test that runs LoRA inference with LoRA switching and compares the output
-    to reference videos using SSIM.
-    """
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
-
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-
-    output_dir = os.path.join(script_dir, 'generated_videos', model_id.split('/')[-1], ATTENTION_BACKEND)
-
-    os.makedirs(output_dir, exist_ok=True)
-
-    BASE_PARAMS = MODEL_TO_PARAMS[model_id]
-    num_inference_steps = BASE_PARAMS["num_inference_steps"]
-
-    init_kwargs = {
-        "num_gpus": BASE_PARAMS["num_gpus"],
-        "flow_shift": BASE_PARAMS["flow_shift"],
-        "use_cpu_offload": BASE_PARAMS["use_cpu_offload"],
-    }
-    if "text-encoder-precision" in BASE_PARAMS:
-        init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"]
-
-    generation_kwargs = {
-        "num_inference_steps": num_inference_steps,
-        "output_path": output_dir,
-        "height": BASE_PARAMS["height"],
-        "width": BASE_PARAMS["width"],
-        "num_frames": BASE_PARAMS["num_frames"],
-        "guidance_scale": BASE_PARAMS["guidance_scale"],
-        "seed": BASE_PARAMS["seed"],
-        "fps": BASE_PARAMS["fps"],
-        "save_video": True,
-    }
-    generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs)
-    for lora_config in LORA_CONFIGS:
-        lora_nickname = lora_config["lora_nickname"]
-        lora_path = lora_config["lora_path"]
-        prompt = lora_config["prompt"]
-        generation_kwargs["negative_prompt"] = lora_config["negative_prompt"]
-
-        generator.set_lora_adapter(lora_nickname=lora_nickname, lora_path=lora_path)
-        output_video_name = f"{lora_path.split('/')[-1]}_{prompt[:50]}"
-        generation_kwargs["output_path"] = output_dir
-        generation_kwargs["output_video_name"] = output_video_name
-
-        generator.generate_video(prompt, **generation_kwargs)
-
-        assert os.path.exists(
-            output_dir), f"Output video was not generated at {output_dir}"
-
-        reference_folder = os.path.join(script_dir, 'L40S_reference_videos', model_id.split('/')[-1], ATTENTION_BACKEND)
-        
-        if not os.path.exists(reference_folder):
-            logger.error("Reference folder missing")
-            raise FileNotFoundError(
-                f"Reference video folder does not exist: {reference_folder}")
-
-        # Find the matching reference video for the switched LoRA
-        reference_video_name = None
-
-        for filename in os.listdir(reference_folder):
-            # Check if the filename starts with the expected output_video_name and ends with .mp4
-            if filename.startswith(output_video_name) and filename.endswith('.mp4'):
-                reference_video_name = filename  # Remove .mp4 extension to match the logic below
-                break
-
-        if not reference_video_name:
-            logger.error(f"Reference video not found for adapter: {lora_path} with prompt: {prompt[:50]} and backend: {ATTENTION_BACKEND}")
-            raise FileNotFoundError(f"Reference video missing for adapter {lora_path}")
-
-        reference_video_path = os.path.join(reference_folder, reference_video_name)
-        generated_video_path = os.path.join(output_dir, output_video_name + ".mp4")
-
-        logger.info(
-            f"Computing SSIM between {reference_video_path} and {generated_video_path}"
-        )
-        ssim_values = compute_video_ssim_torchvision(reference_video_path,
-                                                    generated_video_path,
-                                                    use_ms_ssim=True)
-
-        mean_ssim = ssim_values[0]
-        logger.info(f"SSIM mean value: {mean_ssim}")
-        logger.info(f"Writing SSIM results to directory: {output_dir}")
-
-        success = write_ssim_results(output_dir, ssim_values, reference_video_path,
-                                    generated_video_path, num_inference_steps,
-                                    prompt)
-
-        if not success:
-            logger.error("Failed to write SSIM results to file")
-
-        min_acceptable_ssim = lora_config["ssim_threshold"]
-        assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for adapter {lora_config['lora_path']}" 
-
-
-
-
@@ -97,8 +97,3 @@ def run_precision_tests_STA():
@app.function(gpu="H100:1", image=image, timeout=900)
 def run_precision_tests_VSA():
    run_test("python csrc/attn/tests/test_block_sparse.py")
-
-
-@app.function(gpu="L40S:1", image=image, timeout=3600)
-def run_inference_lora_tests():
-    run_test("pytest ./fastvideo/v1/tests/inference/lora/test_lora_inference_similarity.py -vs")
@@ -105,7 +105,7 @@ def run_training():
        "--num_latent_t", "8",
        "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
        "--sp_size", NUM_GPUS_PER_NODE_TRAINING,
-        "--tp_size", 1,
+        "--tp_size", NUM_GPUS_PER_NODE_TRAINING,
        "--hsdp_replicate_dim", "1",
        "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
        "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
@@ -24,7 +24,7 @@ FastHunyuan-diffusers: {
 "flow_shift": 17,
 "seed": 1024,
 "sp_size": 2,
-"tp_size": 1,
+"tp_size": 2,
 "vae_sp": true,
 "fps": 24
 }
@@ -41,7 +41,7 @@ Wan2.1-T2V-1.3B-Diffusers: {
 "flow_shift": 7.0,
 "seed": 1024,
 "sp_size": 2,
-"tp_size": 1,
+"tp_size": 2,
 "vae_sp": True,
 "fps": 24,
 "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
@@ -60,7 +60,7 @@ Wan2.1-I2V-14B-480P-Diffusers: {
 "flow_shift": 7.0,
 "seed": 1024,
 "sp_size": 2,
-"tp_size": 1,
+"tp_size": 2,
 "vae_sp": True,
 "fps": 24,
 "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
@@ -1,31 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 import argparse
 import os
-import json
-from fastvideo.v1.logger import init_logger

 import numpy as np
 import torch
 from pytorch_msssim import ms_ssim, ssim
 from torchvision.io import read_video

-logger = init_logger(__name__)
-

 def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
-    """
-    Compute SSIM between two videos.
-
-    Args:
-        video1_path: Path to the first video.
-        video2_path: Path to the second video.
-        use_ms_ssim: Whether to use Multi-Scale Structural Similarity(MS-SSIM) instead of SSIM.
-    """
    print(f"Computing SSIM between {video1_path} and {video2_path}...")
-    if not os.path.exists(video1_path):
-        raise FileNotFoundError(f"Video1 not found: {video1_path}")
-    if not os.path.exists(video2_path):
-        raise FileNotFoundError(f"Video2 not found: {video2_path}")

    frames1, _, _ = read_video(video1_path,
                               pts_unit='sec',
@@ -81,26 +65,7 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
 def compare_folders(reference_folder, generated_folder, use_ms_ssim=True):
    """
    Compare videos with the same filename between reference_folder and generated_folder
-    
-    Example usage:
-        results = compare_folders(reference_folder, generated_folder,
-                              args.use_ms_ssim)
-        for video_name, ssim_value in results.items():
-            if ssim_value is not None:
-                print(
-                    f"{video_name}: {ssim_value[0]:.4f}, Min SSIM: {ssim_value[1]:.4f}, Max SSIM: {ssim_value[2]:.4f}"
-                )
-            else:
-                print(f"{video_name}: Error during comparison")
-
-        valid_ssims = [v for v in results.values() if v is not None]
-        if valid_ssims:
-            avg_ssim = np.mean([v[0] for v in valid_ssims])
-            print(f"\nAverage SSIM across all videos: {avg_ssim:.4f}")
-        else:
-            print("\nNo valid SSIM values to average")
    """
-    
    reference_videos = [
        f for f in os.listdir(reference_folder) if f.endswith('.mp4')
    ]
@@ -127,40 +92,54 @@ def compare_folders(reference_folder, generated_folder, use_ms_ssim=True):

    return results

-def write_ssim_results(output_dir, ssim_values, reference_path, generated_path,
-                       num_inference_steps, prompt):
-    """
-    Write SSIM results to a JSON file in the same directory as the generated videos.
-    """
-    try:
-        logger.info(
-            f"Attempting to write SSIM results to directory: {output_dir}")

-        if not os.path.exists(output_dir):
-            os.makedirs(output_dir, exist_ok=True)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description='Compare videos using SSIM/MS-SSIM metrics')
+    parser.add_argument('--reference',
+                        '-r',
+                        type=str,
+                        help='Path to reference videos directory')
+    parser.add_argument('--generated',
+                        '-g',
+                        type=str,
+                        help='Path to generated videos directory')
+    parser.add_argument('--use-ms-ssim',
+                        action='store_true',
+                        help='Use MS-SSIM instead of SSIM')
+    args = parser.parse_args()

-        mean_ssim, min_ssim, max_ssim = ssim_values
+    script_dir = os.path.dirname(os.path.abspath(__file__))

-        result = {
-            "mean_ssim": mean_ssim,
-            "min_ssim": min_ssim,
-            "max_ssim": max_ssim,
-            "reference_video": reference_path,
-            "generated_video": generated_path,
-            "parameters": {
-                "num_inference_steps": num_inference_steps,
-                "prompt": prompt
-            }
-        }
+    reference_folder = args.reference if args.reference else os.path.join(
+        script_dir, 'reference_videos')
+    generated_folder = args.generated if args.generated else os.path.join(
+        script_dir, 'generated_videos')

-        test_name = f"steps{num_inference_steps}_{prompt[:100]}"
-        result_file = os.path.join(output_dir, f"{test_name}_ssim.json")
-        logger.info(f"Writing JSON results to: {result_file}")
-        with open(result_file, 'w') as f:
-            json.dump(result, f, indent=2)
+    if not os.path.exists(reference_folder):
+        print(f"ERROR: Reference folder {reference_folder} does not exist!")
+        exit(1)

-        logger.info(f"SSIM results written to {result_file}")
-        return True
-    except Exception as e:
-        logger.error(f"ERROR writing SSIM results: {str(e)}")
-        return False
+    if not os.path.exists(generated_folder):
+        print(f"ERROR: Generated folder {generated_folder} does not exist!")
+        exit(1)
+
+    print(f"Comparing videos between {reference_folder} and {generated_folder}")
+    results = compare_folders(reference_folder, generated_folder,
+                              args.use_ms_ssim)
+
+    print("\n===== SSIM Results Summary =====")
+    for video_name, ssim_value in results.items():
+        if ssim_value is not None:
+            print(
+                f"{video_name}: {ssim_value[0]:.4f}, Min SSIM: {ssim_value[1]:.4f}, Max SSIM: {ssim_value[2]:.4f}"
+            )
+        else:
+            print(f"{video_name}: Error during comparison")
+
+    valid_ssims = [v for v in results.values() if v is not None]
+    if valid_ssims:
+        avg_ssim = np.mean([v[0] for v in valid_ssims])
+        print(f"\nAverage SSIM across all videos: {avg_ssim:.4f}")
+    else:
+        print("\nNo valid SSIM values to average")
@@ -7,7 +7,7 @@ import pytest

 from fastvideo import VideoGenerator
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.tests.utils import compute_video_ssim_torchvision, write_ssim_results
+from fastvideo.v1.tests.ssim.compute_ssim import compute_video_ssim_torchvision
 from fastvideo.v1.worker.multiproc_executor import MultiprocExecutor

 logger = init_logger(__name__)
@@ -33,7 +33,7 @@ HUNYUAN_PARAMS = {
    "flow_shift": 17,
    "seed": 1024,
    "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
    "vae_sp": True,
    "fps": 24,
 }
@@ -50,7 +50,7 @@ WAN_T2V_PARAMS = {
    "flow_shift": 7.0,
    "seed": 1024,
    "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
    "vae_sp": True,
    "fps": 24,
    "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
@@ -69,7 +69,7 @@ WAN_I2V_PARAMS = {
    "flow_shift": 7.0,
    "seed": 1024,
    "sp_size": 2,
-    "tp_size": 1,
+    "tp_size": 2,
    "vae_sp": True,
    "fps": 24,
    "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
@@ -99,6 +99,44 @@ I2V_IMAGE_PATHS = [
 ]


+def write_ssim_results(output_dir, ssim_values, reference_path, generated_path,
+                       num_inference_steps, prompt):
+    """
+    Write SSIM results to a JSON file in the same directory as the generated videos.
+    """
+    try:
+        logger.info(
+            f"Attempting to write SSIM results to directory: {output_dir}")
+
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir, exist_ok=True)
+
+        mean_ssim, min_ssim, max_ssim = ssim_values
+
+        result = {
+            "mean_ssim": mean_ssim,
+            "min_ssim": min_ssim,
+            "max_ssim": max_ssim,
+            "reference_video": reference_path,
+            "generated_video": generated_path,
+            "parameters": {
+                "num_inference_steps": num_inference_steps,
+                "prompt": prompt
+            }
+        }
+
+        test_name = f"steps{num_inference_steps}_{prompt[:100]}"
+        result_file = os.path.join(output_dir, f"{test_name}_ssim.json")
+        logger.info(f"Writing JSON results to: {result_file}")
+        with open(result_file, 'w') as f:
+            json.dump(result, f, indent=2)
+
+        logger.info(f"SSIM results written to {result_file}")
+        return True
+    except Exception as e:
+        logger.error(f"ERROR writing SSIM results: {str(e)}")
+        return False
+

@pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS)
@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"])
@@ -200,7 +238,7 @@ def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
        logger.error("Failed to write SSIM results to file")

    min_acceptable_ssim = 0.97
-    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"
+    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}"

@pytest.mark.parametrize("prompt", TEST_PROMPTS)
@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"])
@@ -299,5 +337,5 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
    if not success:
        logger.error("Failed to write SSIM results to file")

-    min_acceptable_ssim = 0.93
-    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"
+    min_acceptable_ssim = 0.95
+    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}"
@@ -43,7 +43,7 @@ def run_worker():
        "--num_latent_t", "4",
        "--num_gpus", "4",
        "--sp_size", "4",
-        "--tp_size", "1",
+        "--tp_size", "4",
        "--hsdp_replicate_dim", "1",
        "--hsdp_shard_dim", "4",
        "--train_sp_batch_size", "1",
@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import numpy as np
+import pytest
+import torch
+from diffusers import CosmosTransformer3DModel
+
+from fastvideo.v1.configs.pipelines import PipelineConfig
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.loader.component_loader import TransformerLoader
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.configs.models.dits import CosmosConfig
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29504"
+
+BASE_MODEL_PATH = "nvidia/Cosmos-Predict2-2B-Text2Image"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_cosmos2_transformer():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
+                         use_cpu_offload=False,
+                         pipeline_config=PipelineConfig(dit_config=CosmosConfig(), dit_precision=precision_str))
+    args.device = device
+
+    loader = TransformerLoader()
+    model2 = loader.load(TRANSFORMER_PATH, "", args).to(device, dtype=precision)
+
+    model1 = CosmosTransformer3DModel.from_pretrained(
+        TRANSFORMER_PATH, device=device,
+        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+
+    total_params = sum(p.numel() for p in model1.parameters())
+    # Calculate weight sum for model1 (converting to float64 to avoid overflow)
+    weight_sum_model1 = sum(
+        p.to(torch.float64).sum().item() for p in model1.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model1 = weight_sum_model1 / total_params
+    logger.info("Model 1 weight sum: %s", weight_sum_model1)
+    logger.info("Model 1 weight mean: %s", weight_mean_model1)
+
+    # Calculate weight sum for model2 (converting to float64 to avoid overflow)
+    total_params_model2 = sum(p.numel() for p in model2.parameters())
+    weight_sum_model2 = sum(
+        p.to(torch.float64).sum().item() for p in model2.parameters())
+    # Also calculate mean for more stable comparison
+    weight_mean_model2 = weight_sum_model2 / total_params_model2
+    logger.info("Model 2 weight sum: %s", weight_sum_model2)
+    logger.info("Model 2 weight mean: %s", weight_mean_model2)
+
+    weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2)
+    logger.info("Weight sum difference: %s", weight_sum_diff)
+    weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2)
+    logger.info("Weight mean difference: %s", weight_mean_diff)
+
+    # Set both models to eval mode
+    model1 = model1.eval()
+    model2 = model2.eval()
+
+    # Create identical inputs for both models
+    batch_size = 1
+    seq_len = 30
+
+    # Video latents [B, C, T, H, W] - Cosmos2 specific dimensions
+    hidden_states = torch.randn(batch_size,
+                                16,
+                                1,  # Single frame for image generation
+                                32,  # Height patches
+                                32,  # Width patches
+                                device=device,
+                                dtype=precision)
+
+    # Text embeddings [B, L, D] - Cosmos2 uses T5 embeddings with 1024 dim
+    encoder_hidden_states = torch.randn(batch_size,
+                                        seq_len,
+                                        1024,  # T5 embedding dimension
+                                        device=device,
+                                        dtype=precision)
+
+    # Timestep
+    timestep = torch.tensor([500], device=device, dtype=precision)
+
+    # padding mask
+    padding_mask = hidden_states.new_zeros(1, 1, 32, 32, device=device, dtype=precision)
+    # print(padding_mask.shape)
+
+    forward_batch = ForwardBatch(
+        data_type="dummy",
+    )
+
+    with torch.amp.autocast('cuda', dtype=precision):
+        output1 = model1(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep=timestep,
+            padding_mask=padding_mask,
+            return_dict=False,
+        )[0]
+        with set_forward_context(
+                current_timestep=0,
+                attn_metadata=None,
+                forward_batch=forward_batch,
+        ):
+            output2 = model2(hidden_states=hidden_states,
+                             encoder_hidden_states=encoder_hidden_states,
+                             timestep=timestep,
+                             padding_mask=padding_mask)
+
+    # Check if outputs have the same shape
+    assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+    assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
+
+    # Check if outputs are similar (allowing for small numerical differences)
+    max_diff = torch.max(torch.abs(output1 - output2))
+    mean_diff = torch.mean(torch.abs(output1 - output2))
+    logger.info("Max Diff: %s", max_diff.item())
+    logger.info("Mean Diff: %s", mean_diff.item())
+    assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}"
+    # mean diff
+    assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}"
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_cosmos2_transformer_video2world():
+    """Test Cosmos2 Video2World variant"""
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    
+    # Use Video2World model path
+    base_model_path = "nvidia/Cosmos-Predict2-2B-Video2World"
+    model_path = maybe_download_model(base_model_path,
+                                      local_dir=os.path.join(
+                                          'data', base_model_path))
+    transformer_path = os.path.join(model_path, "transformer")
+    
+    args = FastVideoArgs(model_path=transformer_path,
+                         use_cpu_offload=False,
+                         pipeline_config=PipelineConfig(dit_config=CosmosConfig(), dit_precision=precision_str))
+    args.device = device
+
+    loader = TransformerLoader()
+    model2 = loader.load(transformer_path, "", args).to(device, dtype=precision)
+
+    model1 = CosmosTransformer3DModel.from_pretrained(
+        transformer_path, device=device,
+        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+
+    # Set both models to eval mode
+    model1 = model1.eval()
+    model2 = model2.eval()
+
+    # Create identical inputs for both models
+    batch_size = 1
+    seq_len = 30
+
+    # Video latents [B, C, T, H, W] - Video2World has additional condition channel
+    hidden_states = torch.randn(batch_size,
+                                17,  # 16 + 1 for condition channel
+                                8,   # Multiple frames for video
+                                32,  # Height patches
+                                32,  # Width patches
+                                device=device,
+                                dtype=precision)
+
+    # Text embeddings [B, L, D]
+    encoder_hidden_states = torch.randn(batch_size,
+                                        seq_len,
+                                        1024,  # T5 embedding dimension
+                                        device=device,
+                                        dtype=precision)
+
+    # Timestep
+    timestep = torch.tensor([500], device=device, dtype=precision)
+
+    forward_batch = ForwardBatch(
+        data_type="dummy",
+    )
+
+    with torch.amp.autocast('cuda', dtype=precision):
+        output1 = model1(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep=timestep,
+            return_dict=False,
+        )[0]
+        with set_forward_context(
+                current_timestep=0,
+                attn_metadata=None,
+                forward_batch=forward_batch,
+        ):
+            output2 = model2(hidden_states=hidden_states,
+                             encoder_hidden_states=encoder_hidden_states,
+                             timestep=timestep)
+
+    # Check if outputs have the same shape
+    assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+    assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
+
+    # Check if outputs are similar (allowing for small numerical differences)
+    max_diff = torch.max(torch.abs(output1 - output2))
+    mean_diff = torch.mean(torch.abs(output1 - output2))
+    logger.info("Max Diff: %s", max_diff.item())
+    logger.info("Mean Diff: %s", mean_diff.item())
+    assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}"
+    # mean diff
+    assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}"
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+import gc
 import math
 import os
 import time
@@ -85,14 +86,6 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
        assert self.transformer is not None
        self.set_schemas()

-        # Set random seeds for deterministic training
-        set_random_seed(self.seed)
-        self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
-            self.seed)
-        self.validation_random_generator = torch.Generator(
-            device="cpu").manual_seed(self.seed)
-        logger.info("Initialized random seeds with seed: %s", self.seed)
-
        self.transformer.requires_grad_(True)
        self.transformer.train()
        if training_args.enable_gradient_checkpointing_type is not None:
@@ -431,6 +424,12 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
                    local_main_process_only=False)
        assert self.training_args is not None

+        # Set random seeds for deterministic training
+        set_random_seed(self.seed)
+        self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
+            self.seed)
+        logger.info("Initialized random seeds with seed: %s", self.seed)
+
        self.noise_scheduler = FlowMatchEulerDiscreteScheduler()

        if self.training_args.resume_from_checkpoint:
@@ -576,7 +575,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
        batch = ForwardBatch(
            **shallow_asdict(sampling_param),
            latents=None,
-            generator=self.validation_random_generator,
+            generator=torch.Generator(device="cpu").manual_seed(self.seed),
            n_tokens=n_tokens,
            eta=0.0,
            VSA_sparsity=training_args.VSA_sparsity,
@@ -599,6 +598,10 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
        # Create sampling parameters if not provided
        sampling_param = SamplingParam.from_pretrained(training_args.model_path)

+        # Set deterministic seed for validation
+        set_random_seed(self.seed)
+        logger.info("Using validation seed: %s", self.seed)
+
        # Prepare validation prompts
        logger.info('rank: %s: fastvideo_args.validation_dataset_file: %s',
                    self.global_rank,
@@ -703,3 +706,5 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
        # Re-enable gradients for training
        training_args.inference_mode = False
        transformer.train()
+        gc.collect()
+        torch.cuda.empty_cache()
@@ -3,7 +3,6 @@ import json
 import math
 import os
 import time
-from collections.abc import Iterator
 from typing import Any

 import torch
@@ -163,9 +162,9 @@ def save_checkpoint(transformer,
                    weight_path,
                    local_main_process_only=False)

-        # Convert training format to diffusers format and save
-        diffusers_state_dict = custom_to_hf_state_dict(
-            cpu_state, transformer.reverse_param_names_mapping)
+        # Convert fastvideo custom format to diffusers format and save
+        diffusers_state_dict = convert_custom_format_to_diffusers_format(
+            cpu_state, transformer)
        save_file(diffusers_state_dict, weight_path)

        logger.info("rank: %s, consolidated checkpoint saved to %s",
@@ -488,25 +487,24 @@ def _has_foreach_support(tensors: list[torch.Tensor],
        t is None or type(t) in [torch.Tensor] for t in tensors)


-def custom_to_hf_state_dict(
-    state_dict: dict[str, Any] | Iterator[tuple[str, torch.Tensor]],
-    reverse_param_names_mapping: dict[str, tuple[str, int,
-                                                 int]]) -> dict[str, Any]:
+def convert_custom_format_to_diffusers_format(state_dict: dict[str, Any],
+                                              transformer) -> dict[str, Any]:
    """
-    Convert fastvideo's custom model format to diffusers format using reverse_param_names_mapping.
+    Convert fastvideo custom format state dict to diffusers format using reverse_param_names_mapping.
    
    Args:
-        state_dict: State dict in fastvideo's custom format
-        reverse_param_names_mapping: Reverse mapping from fastvideo's custom format to diffusers format
+        state_dict: State dict in training format
+        transformer: Transformer model object with _reverse_param_names_mapping
        
    Returns:
        State dict in diffusers format
    """
-    assert len(
-        reverse_param_names_mapping) > 0, "reverse_param_names_mapping is empty"
-    if isinstance(state_dict, Iterator):
-        state_dict = dict(state_dict)
    new_state_dict = {}
+
+    # Get the reverse mapping from the transformer
+    reverse_param_names_mapping = transformer._reverse_param_names_mapping
+    assert reverse_param_names_mapping != {}, "reverse_param_names_mapping is empty"
+
    # Group parameters that need to be split (merged parameters)
    merge_groups: dict[str, list[tuple[str, int, int]]] = {}

@@ -29,7 +29,6 @@ from diffusers.loaders.lora_base import (
    _best_guess_weight_name)  # watch out for potetential removal from diffusers
 from huggingface_hub import snapshot_download
 from remote_pdb import RemotePdb
-from torch.distributed.fsdp import MixedPrecisionPolicy

 import fastvideo.v1.envs as envs
 from fastvideo.v1.logger import init_logger
@@ -685,11 +684,11 @@ def remote_breakpoint() -> None:

@dataclass
 class MixedPrecisionState:
+    master_dtype: torch.dtype | None = None
    param_dtype: torch.dtype | None = None
    reduce_dtype: torch.dtype | None = None
    output_dtype: torch.dtype | None = None
    compute_dtype: torch.dtype | None = None
-    mp_policy: MixedPrecisionPolicy | None = None


 # Thread-local storage for mixed precision state
@@ -703,12 +702,10 @@ def get_mixed_precision_state() -> MixedPrecisionState:
    return cast(MixedPrecisionState, _mixed_precision_state.state)


-def set_mixed_precision_policy(
-    param_dtype: torch.dtype,
-    reduce_dtype: torch.dtype,
-    output_dtype: torch.dtype | None = None,
-    mp_policy: MixedPrecisionPolicy | None = None,
-):
+def set_mixed_precision_policy(master_dtype: torch.dtype,
+                               param_dtype: torch.dtype,
+                               reduce_dtype: torch.dtype,
+                               output_dtype: torch.dtype | None = None):
    """Set mixed precision policy globally.
    
    Args:
@@ -717,10 +714,10 @@ def set_mixed_precision_policy(
        output_dtype: Optional output dtype
    """
    state = MixedPrecisionState(
+        master_dtype=master_dtype,
        param_dtype=param_dtype,
        reduce_dtype=reduce_dtype,
        output_dtype=output_dtype,
-        mp_policy=mp_policy,
    )
    _mixed_precision_state.state = state

@@ -1,5 +0,0 @@
-from .executor import Executor
-from .gpu_worker import run_worker_process
-from .multiproc_executor import MultiprocExecutor
-
-__all__ = ["Executor", "run_worker_process", "MultiprocExecutor"]
@@ -49,9 +49,7 @@ class Executor(ABC):
        return cast(ForwardBatch, outputs[0]["output_batch"])

    @abstractmethod
-    def set_lora_adapter(self,
-                         lora_nickname: str,
-                         lora_path: str | None = None) -> None:
+    def set_lora_adapter(self, lora_nickname: str, lora_path: str) -> None:
        """
        Set the LoRA adapter for the workers.
        """
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import contextlib
 import faulthandler
+import gc
 import multiprocessing as mp
 import os
 import signal
@@ -68,6 +69,8 @@ class Worker:
        torch.cuda.set_device(self.device)

        # _check_if_gpu_supports_dtype(self.model_config.dtype)
+        gc.collect()
+        torch.cuda.empty_cache()
        self.init_gpu_memory = torch.cuda.mem_get_info()[0]

        os.environ["MASTER_ADDR"] = "localhost"
@@ -87,9 +90,7 @@ class Worker:
        output_batch = self.pipeline.forward(forward_batch, self.fastvideo_args)
        return cast(ForwardBatch, output_batch)

-    def set_lora_adapter(self,
-                         lora_nickname: str,
-                         lora_path: str | None = None) -> None:
+    def set_lora_adapter(self, lora_nickname: str, lora_path: str) -> None:
        self.pipeline.set_lora_adapter(lora_nickname, lora_path)

    def shutdown(self) -> dict[str, Any]:
@@ -101,6 +102,9 @@ class Worker:
        if hasattr(self, 'pipeline') and self.pipeline is not None:
            # Clean up pipeline resources if needed
            pass
+        # Release CUDA resources
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()

        # Destroy the distributed environment
        cleanup_dist_env_and_memory(shutdown_ray=False)
@@ -129,18 +133,13 @@ class Worker:

                # Handle regular RPC calls
                if method_name == 'execute_forward':
+                    gc.collect()
+                    torch.cuda.empty_cache()
                    forward_batch = recv_rpc['kwargs']['forward_batch']
                    fastvideo_args = recv_rpc['kwargs']['fastvideo_args']
                    output_batch = self.execute_forward(forward_batch,
                                                        fastvideo_args)
                    self.pipe.send({"output_batch": output_batch.output.cpu()})
-                elif method_name == 'set_lora_adapter':
-                    lora_nickname = recv_rpc['kwargs']['lora_nickname']
-                    lora_path = recv_rpc['kwargs']['lora_path']
-                    self.set_lora_adapter(lora_nickname, lora_path)
-                    logger.info("Worker %d set LoRA adapter %s with path %s",
-                                self.rank, lora_nickname, lora_path)
-                    self.pipe.send({"status": "lora_adapter_set"})
                else:
                    # Handle other methods dynamically if needed
                    args = recv_rpc.get('args', ())
@@ -75,18 +75,12 @@ class MultiprocExecutor(Executor):
                                        })
        return cast(ForwardBatch, responses[0]["output_batch"])

-    def set_lora_adapter(self,
-                         lora_nickname: str,
-                         lora_path: str | None = None) -> None:
-        responses = self.collective_rpc("set_lora_adapter",
-                                        kwargs={
-                                            "lora_nickname": lora_nickname,
-                                            "lora_path": lora_path
-                                        })
-        for i, response in enumerate(responses):
-            if response["status"] != "lora_adapter_set":
-                raise RuntimeError(
-                    f"Worker {i} failed to set LoRA adapter to {lora_path}")
+    def set_lora_adapter(self, lora_nickname: str, lora_path: str) -> None:
+        self.collective_rpc("set_lora_adapter",
+                            kwargs={
+                                "lora_nickname": lora_nickname,
+                                "lora_path": lora_path
+                            })

    def collective_rpc(self,
                       method: str | Callable,
@@ -1,4 +1,3 @@
-# test pr
 [build-system]
 requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"
@@ -79,7 +78,6 @@ exclude = ["assets*", "docker*", "docs", "scripts*"]
 [tool.wheel]
 exclude = ["assets*", "docker*", "docs", "scripts*"]

-
 [tool.mypy]
 warn_unused_configs = true
 ignore_missing_imports = true
@@ -20,7 +20,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
    --train_batch_size=4 \
    --num_latent_t 20 \
    --sp_size 4 \
-    --tp_size 1 \
+    --tp_size 4 \
    --hsdp_replicate_dim 1 \
    --hsdp_shard_dim 4 \
    --num_gpus $NUM_GPUS \
@@ -3,10 +3,13 @@
 num_gpus=4
 export MODEL_BASE=FastVideo/FastHunyuan-Diffusers
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
    --model-path $MODEL_BASE \
    --sp-size $num_gpus \
-    --tp-size 1 \
+    --tp-size $num_gpus \
    --num-gpus $num_gpus \
    --height 720 \
    --width 1280 \
@@ -4,10 +4,13 @@ num_gpus=4
 export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
    --model-path $MODEL_BASE \
    --sp-size $num_gpus \
-    --tp-size 1 \
+    --tp-size $num_gpus \
    --num-gpus $num_gpus \
    --height 720 \
    --width 1280 \
@@ -5,10 +5,13 @@ export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_hunyuan.json
 export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
 export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
 # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
    --model-path $MODEL_BASE \
    --sp-size ${num_gpus} \
-    --tp-size 1 \
+    --tp-size ${num_gpus} \
    --height 768 \
    --width 1280 \
    --num-frames 117 \
@@ -4,10 +4,13 @@ num_gpus=2
 export FASTVIDEO_ATTENTION_BACKEND=
 export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
 # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
    --model-path $MODEL_BASE \
    --sp-size $num_gpus \
-    --tp-size 1 \
+    --tp-size $num_gpus \
    --num-gpus $num_gpus \
    --height 480 \
    --width 832 \
@@ -4,10 +4,13 @@ num_gpus=2
 export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_wan.json
 export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
 export MODEL_BASE=Wan-AI/Wan2.1-T2V-14B-Diffusers
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
    --model-path $MODEL_BASE \
    --sp-size $num_gpus \
-    --tp-size 1 \
+    --tp-size $num_gpus \
    --num-gpus $num_gpus \
    --height 768 \
    --width 1280 \
@@ -4,11 +4,14 @@ num_gpus=1
 export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
 # change model path to local dir if you want to inference using your checkpoint
 export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo 
+# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
    --model-path $MODEL_BASE \
    --sp-size $num_gpus \
-    --tp-size 1 \
+    --tp-size $num_gpus \
    --num-gpus $num_gpus \
    --height 448 \
    --width 832 \
@@ -4,6 +4,9 @@ num_gpus=2
 export FASTVIDEO_ATTENTION_BACKEND=
 export MODEL_BASE=Wan-AI/Wan2.1-I2V-14B-480P-Diffusers
 # export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
 fastvideo generate \
    --model-path $MODEL_BASE \
    --sp-size $num_gpus \
@@ -0,0 +1,113 @@
+import os
+import unittest
+
+import torch
+from transformers import AutoTokenizer, T5EncoderModel
+
+from fastvideo.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D
+
+
+class TestAutoencoderKLCausal3D(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        setUpClass is called once, before any test is run.
+        We can set environment variables or load heavy resources here.
+        """
+        os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
+
+        # Load tokenizer/model that can be reused across all tests
+        cls.tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
+        cls.text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
+
+    def setUp(self):
+        """
+        setUp is called before each test method to prepare fresh state.
+        """
+        self.batch_size = 1
+        self.init_time_len = 9
+        self.init_height = 16
+        self.init_width = 16
+        self.latent_channels = 4
+        self.spatial_compression_ratio = 8
+        self.time_compression_ratio = 4
+
+        # Model initialization config
+        self.init_dict = {
+            "in_channels":
+            3,
+            "out_channels":
+            3,
+            "latent_channels":
+            self.latent_channels,
+            "down_block_types": (
+                "DownEncoderBlockCausal3D",
+                "DownEncoderBlockCausal3D",
+                "DownEncoderBlockCausal3D",
+                "DownEncoderBlockCausal3D",
+            ),
+            "up_block_types": (
+                "UpDecoderBlockCausal3D",
+                "UpDecoderBlockCausal3D",
+                "UpDecoderBlockCausal3D",
+                "UpDecoderBlockCausal3D",
+            ),
+            "block_out_channels": (8, 8, 8, 8),
+            "layers_per_block":
+            1,
+            "act_fn":
+            "silu",
+            "norm_num_groups":
+            4,
+            "scaling_factor":
+            0.476986,
+            "spatial_compression_ratio":
+            self.spatial_compression_ratio,
+            "time_compression_ratio":
+            self.time_compression_ratio,
+            "mid_block_add_attention":
+            True,
+        }
+
+        # Instantiate the model
+        self.model = AutoencoderKLCausal3D(**self.init_dict)
+
+        # Create a random input tensor
+        self.input_tensor = torch.rand(self.batch_size, 3, self.init_time_len, self.init_height, self.init_width)
+
+    def test_encode_shape(self):
+        """
+        Check that the shape of the encoded output matches expectations.
+        """
+        vae_encoder_output = self.model.encode(self.input_tensor)
+
+        # The distribution from the VAE has a .sample() method
+        # so we verify the shape of that sample.
+        sample_shape = vae_encoder_output["latent_dist"].sample().shape
+
+        # We expect shape: [batch_size, latent_channels,
+        #                   (init_time_len // time_compression_ratio) + 1,
+        #                   init_height // spatial_compression_ratio,
+        #                   init_width // spatial_compression_ratio]
+        expected_shape = (
+            self.batch_size,
+            self.latent_channels,
+            (self.init_time_len // self.time_compression_ratio) + 1,
+            self.init_height // self.spatial_compression_ratio,
+            self.init_width // self.spatial_compression_ratio,
+        )
+
+        # (Optional) Print them if you like, or just rely on assertions:
+        print(f"sample_shape: {sample_shape}")
+        print(f"expected_shape: {expected_shape}")
+
+        self.assertEqual(
+            sample_shape,
+            expected_shape,
+            f"Encoded sample shape {sample_shape} does not match {expected_shape}.",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,39 @@
+import os
+import shutil
+
+import pytest
+import torch
+import torch.distributed as dist
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+
+@pytest.fixture(scope="module", autouse=True)
+def setup_distributed():
+    os.environ["RANK"] = "0"
+    os.environ["WORLD_SIZE"] = "1"
+    os.environ["LOCAL_RANK"] = "0"
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = "12345"
+
+    dist.init_process_group("nccl")
+    yield
+    dist.destroy_process_group()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires at least 2 GPUs to run NCCL tests")
+def test_save_and_remove_checkpoint():
+    from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
+    from fastvideo.utils.checkpoint import save_checkpoint
+    from fastvideo.utils.fsdp_util import get_dit_fsdp_kwargs
+
+    transformer = MochiTransformer3DModel(num_layers=0)
+    fsdp_kwargs, _ = get_dit_fsdp_kwargs(transformer, "none")
+    transformer = FSDP(transformer, **fsdp_kwargs)
+
+    test_folder = "./test_checkpoint"
+    save_checkpoint(transformer, 0, test_folder, 0)
+
+    assert os.path.exists(test_folder), "Checkpoint folder was not created."
+
+    shutil.rmtree(test_folder)
+    assert not os.path.exists(test_folder), "Checkpoint folder still exists."
@@ -0,0 +1,111 @@
+from functools import partial
+from multiprocessing import Manager
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from fastvideo.utils.communications import nccl_info, prepare_sequence_parallel_data
+
+
+def _init_distributed_test_gpu(rank, world_size, backend, port, data, results):
+    dist.init_process_group(
+        backend=backend,
+        init_method=f"tcp://127.0.0.1:{port}",
+        world_size=world_size,
+        rank=rank,
+    )
+
+    device = torch.device(f"cuda:{rank}")
+
+    nccl_info.sp_size = world_size
+    nccl_info.rank_within_group = rank
+    nccl_info.group_id = 0
+
+    seq_group = dist.new_group(ranks=list(range(world_size)))
+    nccl_info.group = seq_group
+
+    hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask = data
+    hidden_states = hidden_states[rank].unsqueeze(dim=0).to(device)
+    encoder_hidden_states = encoder_hidden_states.to(device)
+    attention_mask = attention_mask.to(device)
+    encoder_attention_mask = encoder_attention_mask.to(device)
+    print(f"Rank {rank} input hidden_states:\n", hidden_states)
+    print(f"Rank {rank} input hidden_states shape:\n", hidden_states.shape)
+    out_hidden, out_encoder, out_attn_mask, out_encoder_mask = prepare_sequence_parallel_data(
+        hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask)
+    print(f"Rank {rank} output out_hidden:\n", out_hidden)
+
+    shapes = (
+        out_hidden.shape,
+        out_encoder.shape,
+        out_attn_mask.shape,
+        out_encoder_mask.shape,
+    )
+    shape_tensor = torch.tensor([*shapes[0], *shapes[1], *shapes[2], *shapes[3]], dtype=torch.int32, device=device)
+    shape_list = [torch.zeros_like(shape_tensor) for _ in range(world_size)]
+    dist.all_gather(shape_list, shape_tensor, group=seq_group)
+    gathered_shapes = [tuple(s.tolist()) for s in shape_list]
+    out_hidden_cpu = out_hidden.to("cpu")
+
+    results[rank] = {
+        "shapes": gathered_shapes,
+        "out_hidden": out_hidden_cpu,
+    }
+
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 2,
+                    reason="Requires at least 2 GPUs to run NCCL tests")
+def test_prepare_sequence_parallel_data_gpu():
+    world_size = 2
+    backend = "nccl"
+    port = 12355  # or use a random free port if collisions occur
+
+    # Create test tensors on CPU; the dimension at index=2 should be divisible by world_size=2 (if applicable).
+    hidden_states = torch.randn(2, 1, 2, 1, 1)
+    encoder_hidden_states = torch.randn(2, 2)
+    attention_mask = torch.randn(2, 2)
+    encoder_attention_mask = torch.randn(2, 2)
+
+    print("init hidden states", hidden_states)
+
+    manager = Manager()
+    results_dict = manager.dict()
+
+    # Wrap our helper function with partial
+    mp_func = partial(_init_distributed_test_gpu,
+                      world_size=world_size,
+                      backend=backend,
+                      port=port,
+                      data=(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask),
+                      results=results_dict)
+
+    # Spawn two GPU processes (rank=0, rank=1)
+    mp.spawn(mp_func, nprocs=world_size)
+
+    first_rank_shapes = None
+
+    overall_hidden_out = []
+
+    for rank in sorted(results_dict.keys()):
+        rank_data = results_dict[rank]
+        rank_shapes = rank_data["shapes"]
+        if first_rank_shapes is None:
+            first_rank_shapes = rank_shapes
+        assert rank_shapes == first_rank_shapes, (
+            f"Mismatch in shapes across ranks: {rank_shapes} != {first_rank_shapes}")
+        overall_hidden_out.append(rank_data["out_hidden"])
+
+    overall_hidden_out = torch.cat(overall_hidden_out, dim=2)
+    print("overall_hidden_out", overall_hidden_out)
+    print("overall_hidden_out_shape", overall_hidden_out.shape)
+
+    assert torch.allclose(hidden_states, torch.tensor(overall_hidden_out), rtol=1e-7, atol=1e-6)
+
+
+if __name__ == "__main__":
+    test_prepare_sequence_parallel_data_gpu()
Author	SHA1	Message	Date
SolitaryThinker	c28b0c9d89	run lint	2025-07-06 01:16:42 -07:00
SolitaryThinker	feed03b456	cosmos2 dit	2025-07-06 01:12:05 -07:00