fix

trigger ci
2025-06-20 22:47:06 -07:00 · 2025-06-20 22:46:56 -07:00
133 changed files with 1060 additions and 3818 deletions
@@ -2,26 +2,25 @@ env:
  IMAGE_VERSION: "py3.12-latest"

 steps:
-  - label: "pre-commit"
-    command: ".buildkite/scripts/pre_commit.sh"
-    agents:
-      queue: "default"
-    env:
-      - BUILDKITE_CLEAN_CHECKOUT=true
-
-  - wait
+  - block: "Start Build"
+    blocked_state: "running"
+    prompt: "Approve build?"

  - label: "Trigger Tests"
+    command: |
+      echo "Current working directory: $(pwd)"
+      echo "Current branch:"
+      git branch --show-current
+      echo "Full diff:"
+      git diff --name-only $BUILDKITE_PULL_REQUEST_BASE_BRANCH...HEAD
    plugins:
      - monorepo-diff#v1.4.0:
          diff: "git diff --name-only $BUILDKITE_PULL_REQUEST_BASE_BRANCH...HEAD"
          watch:
            - path: 
                - "fastvideo/v1/models/encoders/**"
-                - "fastvideo/v1/models/loader/**"
+                - "fastvideo/v1/models/loaders/**"
                - "fastvideo/v1/tests/encoders/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 30m .buildkite/scripts/pr_test.sh"
                label: "Encoder Tests"
@@ -32,10 +31,8 @@ steps:
                  queue: "default"
            - path:
                - "fastvideo/v1/models/vaes/**"
-                - "fastvideo/v1/models/loader/**"
+                - "fastvideo/v1/models/loaders/**"
                - "fastvideo/v1/tests/vaes/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 30m .buildkite/scripts/pr_test.sh"
                label: "VAE Tests"
@@ -46,12 +43,10 @@ steps:
                  queue: "default"
            - path:
                - "fastvideo/v1/models/dits/**"
-                - "fastvideo/v1/models/loader/**"
+                - "fastvideo/v1/models/loaders/**"
                - "fastvideo/v1/tests/transformers/**"
                - "fastvideo/v1/layers/**"
                - "fastvideo/v1/attention/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 30m .buildkite/scripts/pr_test.sh"
                label: "Transformer Tests"
@@ -60,8 +55,7 @@ steps:
                  - TEST_TYPE=transformer
                agents:
                  queue: "default"
-            - path: 
-                - "fastvideo/v1/**/*.py"
+            - path: "fastvideo/v1/**/*.py"
              config:
                command: "timeout 60m .buildkite/scripts/pr_test.sh"
                label: "SSIM Tests"
@@ -70,79 +64,3 @@ steps:
                  - TEST_TYPE=ssim
                agents:
                  queue: "default"
-            - path:
-                - "fastvideo/v1/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: "Training Tests"
-                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
-                  - TEST_TYPE=training
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/v1/**"
-                - "csrc/attn/vsa/**"
-                - "csrc/attn/tk/**"
-                - "csrc/attn/setup_vsa.py"
-                - "csrc/attn/config_vsa.py"
-                - "csrc/attn/vsa.cpp"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: "Training Tests VSA"
-                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
-                  - TEST_TYPE=training_vsa
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/v1/**"
-                - "csrc/attn/st_attn/**"
-                - "csrc/attn/setup_sta.py"
-                - "csrc/attn/config_sta.py"
-                - "csrc/attn/st_attn.cpp"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: "Inference Tests STA"
-                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
-                  - TEST_TYPE=inference_sta
-                agents:
-                  queue: "default"
-            - path:
-                - "csrc/attn/st_attn/**"
-                - "csrc/attn/setup_sta.py"
-                - "csrc/attn/config_sta.py"
-                - "csrc/attn/st_attn.cpp"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: "Precision Tests STA"
-                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
-                  - TEST_TYPE=precision_sta
-                agents:
-                  queue: "default"
-            - path:
-                - "csrc/attn/vsa/**"
-                - "csrc/attn/tk/**"
-                - "csrc/attn/setup_vsa.py"
-                - "csrc/attn/config_vsa.py"
-                - "csrc/attn/vsa.cpp"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: "Precision Tests VSA"
-                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
-                  - TEST_TYPE=precision_vsa
-                agents:
-                  queue: "default"
@@ -31,10 +31,6 @@ log "Setting up Modal authentication from Buildkite secrets..."
 MODAL_TOKEN_ID=$(buildkite-agent secret get modal_token_id)
 MODAL_TOKEN_SECRET=$(buildkite-agent secret get modal_token_secret)

-WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
-
-WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
-
 if [ -n "$MODAL_TOKEN_ID" ] && [ -n "$MODAL_TOKEN_SECRET" ]; then
    log "Retrieved Modal credentials from Buildkite secrets"
    python3 -m modal token set --token-id "$MODAL_TOKEN_ID" --token-secret "$MODAL_TOKEN_SECRET" --profile buildkite-ci --activate --verify
@@ -58,44 +54,22 @@ if [ -z "${TEST_TYPE:-}" ]; then
 fi
 log "Test type: $TEST_TYPE"

-MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT IMAGE_VERSION=$IMAGE_VERSION"
-
 case "$TEST_TYPE" in
    "encoder")
        log "Running encoder tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
+        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
        ;;
    "vae")
        log "Running VAE tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
+        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
        ;;
    "transformer")
        log "Running transformer tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
+        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
        ;;
    "ssim")
        log "Running SSIM tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_ssim_tests"
-        ;;
-    "training")
-        log "Running training tests..."
-        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests"
-        ;;
-    "training_vsa")
-        log "Running training VSA tests..."
-        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests_VSA"
-        ;;
-    "inference_sta")
-        log "Running inference STA tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_STA"
-        ;;
-    "precision_sta")
-        log "Running precision STA tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_STA"
-        ;;
-    "precision_vsa")
-        log "Running precision VSA tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_VSA"
+        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_ssim_tests"
        ;;
    *)
        log "Error: Unknown test type: $TEST_TYPE"
@@ -1,40 +0,0 @@
-#!/bin/bash
-set -uo pipefail
-
-log() {
-  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
-}
-
-log "=== Starting pre-commit checks ==="
-
-cd "$(dirname "$0")/../.."
-PROJECT_ROOT=$(pwd)
-log "Project root: $PROJECT_ROOT"
-
-if ! python3 -m pre_commit --version &> /dev/null; then
-    log "pre-commit not found, installing..."
-    python3 -m pip install --user pre-commit==4.0.1
-    
-    if ! python3 -m pre_commit --version &> /dev/null; then
-        log "Error: Failed to install pre-commit."
-        exit 1
-    fi
-fi
-
-log "Pre-commit version: $(python3 -m pre_commit --version)"
-
-log "Installing/updating pre-commit hooks..."
-python3 -m pre_commit install --install-hooks
-
-log "Running pre-commit checks on all files..."
-python3 -m pre_commit run --all-files
-PRE_COMMIT_EXIT_CODE=$?
-
-if [ $PRE_COMMIT_EXIT_CODE -eq 0 ]; then
-    log "Pre-commit checks completed successfully"
-else
-    log "Error: Pre-commit checks failed with exit code: $PRE_COMMIT_EXIT_CODE"
-fi
-
-log "=== Pre-commit checks completed with exit code: $PRE_COMMIT_EXIT_CODE ==="
-exit $PRE_COMMIT_EXIT_CODE 
@@ -14,9 +14,13 @@ on:
      - ".github/workflows/pr-test.yml"
      - "pyproject.toml"
      - "docker/Dockerfile.python3.12"
-      - "csrc/**"
  workflow_dispatch:
    inputs:
+      custom_image:
+        description: "Custom image from this repository (default: fastvideo-dev:py3.12-latest)"
+        required: false
+        default: "fastvideo-dev:py3.12-latest"
+        type: string
      run_encoder_test:
        description: "Run encoder-test"
        required: false
@@ -52,16 +56,6 @@ on:
        required: false
        default: false
        type: boolean
-      run_precision_test_STA:
-        description: "Run precision-test-STA"
-        required: false
-        default: false
-        type: boolean
-      run_precision_test_VSA:
-        description: "Run precision-test-VSA"
-        required: false
-        default: false
-        type: boolean
      run_nightly_test:
        description: "Run nightly-test"
        required: false
@@ -71,7 +65,6 @@ on:
 env:
  PYTHONUNBUFFERED: "1"

-
 concurrency:
  group: pr-test-${{ github.ref }}
  cancel-in-progress: true
@@ -91,69 +84,44 @@ jobs:
      training-test: ${{ steps.filter.outputs.training-test }}
      training-test-VSA: ${{ steps.filter.outputs.training-test-VSA }}
      inference-test-STA: ${{ steps.filter.outputs.inference-test-STA }}
-      precision-test-STA: ${{ steps.filter.outputs.precision-test-STA }}
-      precision-test-VSA: ${{ steps.filter.outputs.precision-test-VSA }}
    steps:
      - uses: actions/checkout@v4
      - uses: dorny/paths-filter@v3
        id: filter
        with:
          filters: |
-            # Define reusable path patterns
-            common-paths: &common-paths
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
-            sta-kernel-paths: &sta-kernel-paths
-              - 'csrc/attn/st_attn/**'
-              - 'csrc/attn/setup_sta.py'
-              - 'csrc/attn/config_sta.py'
-              - 'csrc/attn/st_attn.cpp'
-            vsa-kernel-paths: &vsa-kernel-paths
-              - 'csrc/attn/vsa/**'
-              - 'csrc/attn/tk/**'
-              - 'csrc/attn/setup_vsa.py'
-              - 'csrc/attn/config_vsa.py'
-              - 'csrc/attn/vsa.cpp'
-            vsa-paths: &vsa-paths
-              - 'fastvideo/v1/**'
-              - *common-paths
-              - *vsa-kernel-paths
-            
-            # Actual tests
            encoder-test:
              - 'fastvideo/v1/models/encoders/**'
-              - 'fastvideo/v1/models/loader/**'
+              - 'fastvideo/v1/models/loaders/**'
              - 'fastvideo/v1/tests/encoders/**'
-              - *common-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
            vae-test:
              - 'fastvideo/v1/models/vaes/**'
-              - 'fastvideo/v1/models/loader/**'
+              - 'fastvideo/v1/models/loaders/**'
              - 'fastvideo/v1/tests/vaes/**'
-              - *common-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
            transformer-test:
              - 'fastvideo/v1/models/dits/**'
-              - 'fastvideo/v1/models/loader/**'
+              - 'fastvideo/v1/models/loaders/**'
              - 'fastvideo/v1/tests/transformers/**'
              - 'fastvideo/v1/layers/**'
              - 'fastvideo/v1/attention/**'
-              - *common-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
            training-test:
              - 'fastvideo/v1/**'
-              - *common-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
            training-test-VSA:
              - 'fastvideo/v1/**'
-              - *common-paths
-              - *vsa-kernel-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
            inference-test-STA:
              - 'fastvideo/v1/**'
-              - *common-paths
-              - *sta-kernel-paths
-            precision-test-STA:
-              - *common-paths
-              - *sta-kernel-paths
-            precision-test-VSA:
-              - *common-paths
-              - *vsa-kernel-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'

  encoder-test:
    needs: change-filter
@@ -166,7 +134,7 @@ jobs:
      gpu_type: "NVIDIA A40"
      gpu_count: 1
      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
      timeout_minutes: 30
    secrets:
@@ -184,7 +152,7 @@ jobs:
      gpu_type: "NVIDIA A40"
      gpu_count: 1
      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
      timeout_minutes: 30
    secrets:
@@ -202,7 +170,7 @@ jobs:
      gpu_type: "NVIDIA L40S"
      gpu_count: 1
      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
      timeout_minutes: 30
    secrets:
@@ -212,7 +180,8 @@ jobs:
  ssim-test:
    needs: change-filter
    if: >-
-      github.event_name != 'workflow_dispatch' || (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
    strategy:
      fail-fast: false
      matrix:
@@ -238,7 +207,7 @@ jobs:
  training-test:
    needs: change-filter
    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.training-test == 'true') || 
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
@@ -247,7 +216,7 @@ jobs:
      gpu_count: 4
      volume_size: 100
      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/Vanilla -srP"
      timeout_minutes: 30
    secrets:
@@ -258,7 +227,7 @@ jobs:
  training-test-VSA:
    needs: change-filter
    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.training-test-VSA == 'true') || 
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test_VSA == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
@@ -267,7 +236,7 @@ jobs:
      gpu_count: 1
      volume_size: 100
      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/VSA -srP"
      timeout_minutes: 30
    secrets:
@@ -278,7 +247,7 @@ jobs:
  inference-test-STA:
    needs: change-filter
    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.inference-test-STA == 'true') || 
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_inference_test_STA == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
@@ -287,51 +256,13 @@ jobs:
      gpu_count: 1
      volume_size: 100
      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/inference/STA -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

-  precision-test-STA:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-STA == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_STA == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "precision-test-STA"
-      gpu_type: "NVIDIA H100 NVL"
-      gpu_count: 1
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_sta.py"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  precision-test-VSA:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-VSA == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_VSA == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "precision-test-VSA"
-      gpu_type: "NVIDIA H100 NVL"
-      gpu_count: 1
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_block_sparse.py"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
  nightly-test:
    if: >-
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
@@ -342,7 +273,7 @@ jobs:
      gpu_count: 4
      volume_size: 100
      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs"
      timeout_minutes: 30
    secrets:
@@ -351,8 +282,7 @@ jobs:
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}

  runpod-cleanup:
-    # Add other jobs to this list as you create them
-    needs: [encoder-test, vae-test, transformer-test, ssim-test, training-test, training-test-VSA, inference-test-STA, precision-test-STA, precision-test-VSA] 
+    needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
    if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }}
    runs-on: ubuntu-latest
    steps:
@@ -369,7 +299,7 @@ jobs:

      - name: Cleanup all RunPod instances
        env:
-          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"]'
+          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12"]'
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        run: python .github/scripts/runpod_cleanup.py
@@ -4,7 +4,7 @@


 ## Installation
-We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only support H100/H200, because ThunderKittens uses TMA but doesn't support Blackwell yet.
+We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
 First, install C++20 for ThunderKittens:
 ```bash
 sudo apt update
@@ -53,14 +53,8 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False)

 ## Test
 ```bash
-python tests/test_sta.py # test STA
-python tests/test_block_sparse.py # test VSA
+python test/test_sta.py
 ```
-## Benchmark
-```bash
-python benchmarks/bench_sta.py
-```
-

 ## How Does STA Work?
 We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 
@@ -5,7 +5,6 @@ import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from st_attn import sliding_tile_attention
-from triton.testing import do_bench


 def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"):
@@ -14,16 +13,16 @@ def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"):
    return f if mode == "fwd" else (2.5 * f if mode == "bwd" else 3.5 * f)


-def compute_TFLOPS(flops, ms):
-    flops = flops / 1e12
-    ms = ms / 1e3
-    return flops / ms
+def efficiency(flop, time):
+    flop = flop / 1e12
+    time = time / 1e6
+    return flop / time


 def benchmark_attention(configurations):
    results = {'fwd': defaultdict(list), 'bwd': defaultdict(list)}

-    for B, H, N, D, causal, dit_seq_shape, window_size in configurations:
+    for B, H, N, D, causal in configurations:
        print("=" * 60)
        print(f"Timing forward and backward pass for B={B}, H={H}, N={N}, D={D}, causal={causal}")

@@ -31,31 +30,38 @@ def benchmark_attention(configurations):
        k = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
        v = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()

-        # grad_output = torch.randn_like(q, requires_grad=False).contiguous()
-        # qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous()
-        # kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous()
-        # vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous()
+        grad_output = torch.randn_like(q, requires_grad=False).contiguous()

+        qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous()
+        kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous()
+        vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous()

-        # # Warmup for forward pass
-        # for _ in range(10):
-        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
+        # Prepare for timing forward pass
+        start_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+        end_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]

-        # # Time the forward pass
-        # for i in range(10):
-        #     start_events_fwd[i].record()
-        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
-        #     end_events_fwd[i].record()
-        ms = do_bench(lambda: sliding_tile_attention(q, k, v, [window_size] * 24, 0, False, dit_seq_shape))
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()

-        # times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)]
-        # time_us_fwd = np.mean(times_fwd) * 1000
+        # Warmup for forward pass
+        for _ in range(10):
+            o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, '18x48x80')

-        tflops_fwd = compute_TFLOPS(flops(B, N, H, D, causal, 'fwd'), ms)
+        # Time the forward pass
+        for i in range(10):
+            start_events_fwd[i].record()
+            o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, '18x48x80')
+            end_events_fwd[i].record()
+
+        torch.cuda.synchronize()
+        times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)]
+        time_us_fwd = np.mean(times_fwd) * 1000
+
+        tflops_fwd = efficiency(flops(B, N, H, D, causal, 'fwd'), time_us_fwd)
        results['fwd'][(D, causal)].append((N, tflops_fwd))

-        print(f"Average time for forward pass (ms): {ms:.2f}")
-        print(f"Average TFLOPS: {tflops_fwd}")
+        print(f"Average time for forward pass in us: {time_us_fwd:.2f}")
+        print(f"Average efficiency for forward pass in TFLOPS: {tflops_fwd}")
        print("-" * 60)

        # torch.cuda.empty_cache()
@@ -79,14 +85,15 @@ def benchmark_attention(configurations):
        # times_bwd = [s.elapsed_time(e) for s, e in zip(start_events_bwd, end_events_bwd)]
        # time_us_bwd = np.mean(times_bwd) * 1000

-        # tflops_bwd = compute_TFLOPS(flops(B, N, H, D, causal, 'bwd'), ms)
+        # tflops_bwd = efficiency(flops(B, N, H, D, causal, 'bwd'), time_us_bwd)
        # results['bwd'][(D, causal)].append((N, tflops_bwd))

-        # print(f"Average time for backward pass(ms): {ms:.2f}")
-        # print(f"Average TFLOPS: {tflops_bwd}")
-        # print("=" * 60)
+        # print(f"Average time for backward pass in us: {time_us_bwd:.2f}")
+        # print(f"Average efficiency for backward pass in TFLOPS: {tflops_bwd}")
+        print("=" * 60)

        torch.cuda.empty_cache()
+        torch.cuda.synchronize()

    return results

@@ -117,10 +124,7 @@ def plot_results(results):

 # Example list of configurations to test
 configurations = [
-    (2, 24, 69120, 128, False, '18x48x80', [3, 6, 10]),
-    (2, 24, 69120, 128, True, '18x48x80', [3, 6, 10]),
-    (2, 24, 82944, 128, False, '36x48x48', [3, 3, 6]), # Stepvideo
-    (2, 24, 82944, 128, True, '36x48x48', [3, 3, 6]),
+    (2, 24, 69120, 128, False),
    # (16, 16, 768*16,  128, False),
    # (16, 16, 768*2,  128, False),
    # (16, 16, 768*4,  128, False),
@@ -4,17 +4,9 @@
 #include <cooperative_groups.h>
 #include <iostream>
 #include <stdio.h>
-#include <c10/cuda/CUDAGuard.h>
-
-// #define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value)))
-__device__ __forceinline__ int clamp_int(int value, int min, int max) {
-    return (value < min) ? min : ((value > max) ? max : value);
-}
-// #define ABS(x) ((x) < 0 ? -(x) : (x))
-__device__ __forceinline__ int abs_int(int value) {
-    return (value < 0) ? -value : value;
-}

+#define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value)))
+#define ABS(x) ((x) < 0 ? -(x) : (x))

 constexpr int CONSUMER_WARPGROUPS = (3); 
 constexpr int PRODUCER_WARPGROUPS = (1); 
@@ -125,16 +117,16 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
            int qt = seq_idx / 6 / (CH * CW);
            int qh = (seq_idx / 6) % (CH * CW) / CW;
            int qw = (seq_idx / 6) % CW;
-            qt = clamp_int(qt, DT, CT-DT-1);
-            qh = clamp_int(qh, DH, CH-DH-1);
-            qw = clamp_int(qw, DW, CW-DW-1);
+            qt = CLAMP(qt, DT, CT-DT-1);
+            qh = CLAMP(qh, DH, CH-DH-1);
+            qw = CLAMP(qw, DW, CW-DW-1);
            int count = 0;
            int j = 0;
            while (count < K::stages - 1) {
                int kt = j / 3 / (CH * CW);
                int kh = (j / 3) % (CH * CW) / CW;
                int kw = (j / 3) % CW;
-                bool mask = (abs_int(qt - kt) <= DT) && (abs_int(qh - kh) <= DH) && (abs_int(qw - kw) <= DW);
+                bool mask = (ABS(qt - kt) <= DT) && (ABS(qh - kh) <= DH) && (ABS(qw - kw) <= DW);
                if (mask){
                    coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0};
                    tma::expect_bytes(k_smem_arrived[count], sizeof(k_tile));
@@ -175,15 +167,15 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
                int qt = seq_idx / 6 / (CH * CW);
                int qh = (seq_idx / 6) % (CH * CW) / CW;
                int qw = (seq_idx / 6) % CW;
-                qt = clamp_int(qt, DT, CT-DT-1);
-                qh = clamp_int(qh, DH, CH-DH-1);
-                qw = clamp_int(qw, DW, CW-DW-1);
-                int k_t_min = clamp_int(qt-DT, 0, CT-1);
-                int k_t_max = clamp_int(qt+DT, 0, CT-1);
-                int k_h_min = clamp_int(qh-DH, 0, CH-1);
-                int k_h_max = clamp_int(qh+DH, 0, CH-1);
-                int k_w_min = clamp_int(qw-DW, 0, CW-1);
-                int k_w_max = clamp_int(qw+DW, 0, CW-1);
+                qt = CLAMP(qt, DT, CT-DT-1);
+                qh = CLAMP(qh, DH, CH-DH-1);
+                qw = CLAMP(qw, DW, CW-DW-1);
+                int k_t_min = CLAMP(qt-DT, 0, CT-1);
+                int k_t_max = CLAMP(qt+DT, 0, CT-1);
+                int k_h_min = CLAMP(qh-DH, 0, CH-1);
+                int k_h_max = CLAMP(qh+DH, 0, CH-1);
+                int k_w_min = CLAMP(qw-DW, 0, CW-1);
+                int k_w_max = CLAMP(qw+DW, 0, CW-1);
                int count = 0;
                for (int kt = k_t_min; kt <= k_t_max; kt++) {
                    for (int kh = k_h_min; kh <= k_h_max; kh++) {
@@ -242,7 +234,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
            // the last three kv blocks are for text, we process them separately
            kv_iters = img_kv_blocks - 1;
        } else {
-            kv_iters = clamp_int(DT*2+1, 1, CT) * clamp_int(DH*2+1, 1, CH) * clamp_int(DW*2+1, 1, CW) * 3 - 1 ; 
+            kv_iters = CLAMP(DT*2+1, 1, CT) * CLAMP(DH*2+1, 1, CH) * CLAMP(DW*2+1, 1, CW) * 3 - 1 ; 
        }

        kittens::wait(qsmem_semaphore, 0);
@@ -423,9 +415,8 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
    float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
    float* d_l   = reinterpret_cast<float*>(l_ptr);

-    //cudadevicesynchronize();
-    const c10::cuda::OptionalCUDAGuard device_guard(q.device());
-    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
+    cudaDeviceSynchronize();
+    auto stream = at::cuda::getCurrentCUDAStream().stream(); 


    if (head_dim == 128) {
@@ -451,8 +442,8 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,

        globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len),  static_cast<int>(text_length), static_cast<int>(hr)};

-        constexpr int mem_size = kittens::MAX_SHARED_MEMORY;
-        int threads = NUM_WORKERS * kittens::WARP_THREADS;
+        auto mem_size = kittens::MAX_SHARED_MEMORY;
+        auto threads  = NUM_WORKERS * kittens::WARP_THREADS;
        if (has_text) {
            // TORCH_CHECK(seq_len % (CONSUMER_WARPGROUPS*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 192");
            dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4)-2, qo_heads, batch);
@@ -832,10 +823,10 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,

        }
        CHECK_CUDA_ERROR(cudaGetLastError());
-        // cudaStreamSynchronize(stream);
+        cudaStreamSynchronize(stream);
    }

    return o;
-    //cudadevicesynchronize();
+    cudaDeviceSynchronize();
 }

@@ -7,7 +7,6 @@ from vsa import BLOCK_M, BLOCK_N

 import numpy as np
 import random
-import gc

 def set_seed(seed: int = 42):
    # Python random module
@@ -21,6 +20,15 @@ def set_seed(seed: int = 42):
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if using multi-GPU

+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
+    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
+    parser.add_argument('--num_heads', type=int, default=6, help='Number of heads')
+    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
+    parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to')
+    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark')
+    parser.add_argument('--num_iterations', type=int, default=100, help='Number of test iterations to run')
+    return parser.parse_args()

@torch.no_grad
 def precision_metric(quant_o, fa2_o): 
@@ -127,7 +135,9 @@ def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device=
                
    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask

-def main(args):
+def main():
+    args = parse_arguments()
+
    set_seed(42)
    
    # Extract parameters
@@ -181,36 +191,23 @@ def main(args):
            block_mask_expanded = block_sparse_mask.unsqueeze(-1).unsqueeze(-2)  # [b, h, num_q_blocks, num_kv_blocks, 1, 1]
            block_mask_expanded = block_mask_expanded.expand(-1, -1, -1, -1, BLOCK_M, BLOCK_N)  # [b, h, num_q_blocks, num_kv_blocks, BLOCK_M, BLOCK_N]
            full_mask = block_mask_expanded.permute(0, 1, 2, 4, 3, 5).reshape(batch, head, seq_len, seq_len)
+            
+            q_sdpa = q.clone()
+            k_sdpa = k.clone()
+            v_sdpa = v.clone()

            q.requires_grad = True
            k.requires_grad = True
            v.requires_grad = True
-
-
-            # testing forward
-            o = BlockSparseAttentionFunction.apply(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
-            del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded
-            grad_o = torch.randn_like(o)
-            o.backward(grad_o)
-            # clear memory
-            q_sdpa = q.detach().clone()
-            k_sdpa = k.detach().clone()
-            v_sdpa = v.detach().clone()
            q_sdpa.requires_grad = True
            k_sdpa.requires_grad = True
            v_sdpa.requires_grad = True
-            q.data = torch.empty(0, device=q.device)
-            k.data = torch.empty(0, device=k.device)
-            v.data = torch.empty(0, device=v.device)
-            torch.cuda.empty_cache()

+            # testing forward
+            o = BlockSparseAttentionFunction.apply(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
            o_sdpa = torch.nn.functional.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, attn_mask=full_mask)

-
            sim, l1, rmse = precision_metric(o, o_sdpa)
-            assert sim > 0.9999, f"SSIM too low: {sim}"
-            assert l1 < 8e-5, f"l1 too large: {l1}"
-            assert rmse < 2e-5, f"RMSE too large: {rmse}"
            forward_metrics['sim'].append(sim)
            forward_metrics['l1'].append(l1)
            forward_metrics['rmse'].append(rmse)
@@ -218,72 +215,52 @@ def main(args):
            print(f"block_sparse_attention_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")

            # test backward
+            grad_o = torch.randn_like(o)
+            o.backward(grad_o)
            o_sdpa.backward(grad_o)

            sim, l1, rmse = precision_metric(q.grad, q_sdpa.grad)
-            # Error bounds collected on H100
-            assert sim > 0.9999, f"SSIM too low: {sim}"
-            assert l1 < 4e-3, f"l1 too large: {l1}"
-            assert rmse < 3e-4, f"RMSE too large: {rmse}"
            grad_q_metrics['sim'].append(sim)
            grad_q_metrics['l1'].append(l1)
            grad_q_metrics['rmse'].append(rmse)
            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}")        
            
            sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad)
-            assert sim > 0.9999, f"SSIM too low: {sim}"
-            assert l1 < 4e-3, f"l1 too large: {l1}"
-            assert rmse < 2e-4, f"RMSE too large: {rmse}"
            grad_k_metrics['sim'].append(sim)
            grad_k_metrics['l1'].append(l1)
            grad_k_metrics['rmse'].append(rmse)
            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
            
            sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad)
-            assert sim > 0.9999, f"SSIM too low: {sim}"
-            assert l1 < 1e-4, f"l1 too large: {l1}"
-            assert rmse < 2e-5, f"RMSE too large: {rmse}"
            grad_v_metrics['sim'].append(sim)
            grad_v_metrics['l1'].append(l1)
            grad_v_metrics['rmse'].append(rmse)
            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
-            
-            del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa
-            gc.collect()
-            torch.cuda.empty_cache()
-
+        
        # Print summary statistics if multiple iterations were run
        if num_iterations > 1:
            print("\n" + "="*50)
            print(f"Summary Statistics (over {num_iterations} iterations):")
            
            print("\nForward metrics:")
-            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}, min={np.min(forward_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}, max={np.max(forward_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}, max={np.max(forward_metrics['rmse']):.6f}")
+            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}")
            
            print("\nGradient Q metrics:")
-            print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}, min={np.min(grad_q_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}, max={np.max(grad_q_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}, max={np.max(grad_q_metrics['rmse']):.6f}")
+            print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}")
            
            print("\nGradient K metrics:")
-            print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}, min={np.min(grad_k_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}, max={np.max(grad_k_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}, max={np.max(grad_k_metrics['rmse']):.6f}")
+            print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}")
            
            print("\nGradient V metrics:")
-            print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}, min={np.min(grad_v_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}, max={np.max(grad_v_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}, max={np.max(grad_v_metrics['rmse']):.6f}")
+            print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}")

 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
-    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
-    parser.add_argument('--num_heads', type=int, default=6, help='Number of heads')
-    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
-    parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to')
-    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark')
-    parser.add_argument('--num_iterations', type=int, default=50, help='Number of test iterations to run')
-    args = parser.parse_args()
-    main(args)
+    main()
@@ -81,7 +81,5 @@ std = 10

 # Run correctness check directly
 results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output')
-assert results['TK vs FLEX']['avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large"
-assert results['TK vs FLEX']['max_diff'] < 4e-2, f"Maximum difference: {results['TK vs FLEX']['max_diff']} is too large"
 print(f"Average difference: {results['TK vs FLEX']['avg_diff']}")
 print(f"Maximum difference: {results['TK vs FLEX']['max_diff']}")
@@ -3,8 +3,6 @@
 #include "kittens.cuh"
 #include <cooperative_groups.h>
 #include <iostream>
-#include <c10/cuda/CUDAGuard.h>
-

 using namespace kittens;
 namespace cg = cooperative_groups;
@@ -942,9 +940,8 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
    float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
    float* d_l   = reinterpret_cast<float*>(l_ptr);

-    //cudadevicesynchronize();
-    const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
-    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
+    cudaDeviceSynchronize();
+    auto stream = at::cuda::getCurrentCUDAStream().stream(); 

    if (head_dim == 64) {
        using q_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>;
@@ -969,7 +966,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v

        globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_kv_blocks_per_q), reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr())};

-        constexpr int mem_size = 54000;
+        auto mem_size = 54000;

        dim3 grid(seq_len/(64), qo_heads, batch);

@@ -982,7 +979,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
        fwd_attend_ker<64><<<grid, (128), mem_size, stream>>>(g);

        CHECK_CUDA_ERROR(cudaGetLastError());
-        // cudaStreamSynchronize(stream);
+        cudaStreamSynchronize(stream);
    }

    if (head_dim == 128) {
@@ -1008,7 +1005,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v

        globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_kv_blocks_per_q), reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr())};

-        constexpr int mem_size = 54000;
+        auto mem_size = 54000;

        dim3 grid(seq_len/(64), qo_heads, batch);

@@ -1021,11 +1018,11 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
        fwd_attend_ker<128><<<grid, (128), mem_size, stream>>>(g);

        CHECK_CUDA_ERROR(cudaGetLastError());
-        // cudaStreamSynchronize(stream);
+        cudaStreamSynchronize(stream);
    }

    return {o, l_vec};
-    //cudadevicesynchronize();
+    cudaDeviceSynchronize();
 }

 std::vector<torch::Tensor> 
@@ -1135,14 +1132,13 @@ block_sparse_attention_backward(torch::Tensor q,
    float* d_kg = reinterpret_cast<float*>(kg_ptr);
    float* d_vg = reinterpret_cast<float*>(vg_ptr);

-    constexpr int mem_size = kittens::MAX_SHARED_MEMORY; 
-    int threads  = 4 * kittens::WARP_THREADS;
+    auto mem_size = kittens::MAX_SHARED_MEMORY; 
+    auto threads  = 4 * kittens::WARP_THREADS;

-    //cudadevicesynchronize();
-    const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
-    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
+    cudaDeviceSynchronize();
+    auto stream = at::cuda::getCurrentCUDAStream().stream();

-    //  cudaStreamSynchronize(stream);
+    cudaStreamSynchronize(stream);

    // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256");
    dim3 grid_bwd(seq_len/(4*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
@@ -1226,7 +1222,7 @@ block_sparse_attention_backward(torch::Tensor q,
        dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
        threads = 128;

-        //cudadevicesynchronize();
+        cudaDeviceSynchronize();

        {
            cudaFuncSetAttribute(
@@ -1244,8 +1240,8 @@ block_sparse_attention_backward(torch::Tensor q,
        }

        // CHECK_CUDA_ERROR(cudaGetLastError());
-        // cudaStreamSynchronize(stream);
-        //cudadevicesynchronize();
+        cudaStreamSynchronize(stream);
+        cudaDeviceSynchronize();
        // const auto kernel_end = std::chrono::high_resolution_clock::now();
        // std::cout << "Kernel Time: " << std::chrono::duration_cast<std::chrono::microseconds>(kernel_end - start).count() << "us" << std::endl;
        // std::cout << "---" << std::endl;
@@ -1330,7 +1326,7 @@ block_sparse_attention_backward(torch::Tensor q,
        dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
        threads = 128;

-        //cudadevicesynchronize();
+        cudaDeviceSynchronize();

        {
            cudaFuncSetAttribute(
@@ -1342,10 +1338,10 @@ block_sparse_attention_backward(torch::Tensor q,
            bwd_attend_ker<128><<<grid_bwd_2, threads, 113000, stream>>>(bwd_global); 
        }

-        // cudaStreamSynchronize(stream);
-        //cudadevicesynchronize();
+        cudaStreamSynchronize(stream);
+        cudaDeviceSynchronize();
    }

    return {qg, kg, vg};
-    //cudadevicesynchronize();
+    cudaDeviceSynchronize();
 }
@@ -10,7 +10,7 @@ def main():
    # attempt to identify the optimal arguments.
    generator = VideoGenerator.from_pretrained(
        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        # FastVideo will automatically handle distributed setup
+        # if num_gpus > 1, FastVideo will automatically handle distributed setup
        num_gpus=2,
        use_fsdp_inference=True,
        use_cpu_offload=False
@@ -1,10 +0,0 @@
-This directory contain e2e examples scripts for finetuning Wan2.1 I2V. 
-
-Execute the following commands from `FastVideo/` to run training:
-
- Download crush-smol dataset:
-`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
- Preprocess the videos and captions into latents:
-`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
- Edit the following file and run finetuning:
-`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`
@@ -1,3 +0,0 @@
-#!/bin/bash
-
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -1,91 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
-VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
-NUM_GPUS=8
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_i2v_finetune"
-  --output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
-  --max_train_steps 2000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 8
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 8
-  --tp_size 8
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 8
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_preprocessed_path "$VALIDATION_DIR"
-  --validation_steps 100
-  --validation_sampling_steps "40"
-  --validation_guidance_scale "1.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --checkpointing_steps 1000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --allow_tf32
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-)
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -1,130 +0,0 @@
-#!/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
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate will-fv
-
-# Basic Info
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=8
-
-MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
-VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-
-# Training arguments
-training_args=(
-  --tracker_project_name wan_i2v_finetune
-  --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"
-  --max_train_steps=2000
-  --train_batch_size=2
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps=1
-  --num_latent_t 8
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size $NUM_GPUS
-  --tp_size $NUM_GPUS
-  --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
-  --hsdp_shard_dim $NUM_GPUS
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 10
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_preprocessed_path "$VALIDATION_DIR"
-  --validation_steps 100
-  --validation_sampling_steps "40"
-  --validation_guidance_scale "1.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate=1e-5
-  --mixed_precision="bf16"
-  --checkpointing_steps=1000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --allow_tf32
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-)
-
-srun torchrun \
--nnodes $SLURM_JOB_NUM_NODES \
--nproc_per_node $NUM_GPUS \
--node_rank $SLURM_PROCID \
--rdzv_backend=c10d \
--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -1,25 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=1 # 2,4,8
-MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-MODEL_TYPE="wan"
-DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_i2v/"
-VALIDATION_PATH="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
-
-torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
-    --model_path $MODEL_PATH \
-    --data_merge_path $DATA_MERGE_PATH \
-    --preprocess_video_batch_size 8 \
-    --max_height 480 \
-    --max_width 832 \
-    --num_frames 77 \
-    --dataloader_num_workers 0 \
-    --output_dir=$OUTPUT_DIR \
-    --model_type $MODEL_TYPE \
-    --train_fps 16 \
-    --validation_dataset_file $VALIDATION_PATH \
-    --samples_per_file 8 \
-    --flush_frequency 8 \
-    --preprocess_task "i2v" 
@@ -1,31 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
-      "image_path": null,
-      "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
-      "image_path": null,
-      "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}
@@ -1,10 +0,0 @@
-This directory contain e2e examples scripts for finetuning Wan2.1 T2v. 
-
-Execute the following commands from `FastVideo/` to run training:
-
- Download crush-smol dataset:
-`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/download_dataset.sh`
- Preprocess the videos and captions into latents:
-`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/preprocess_wan_data_t2v.sh`
- Edit the following file and run finetuning:
-`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh`
@@ -1,3 +0,0 @@
-#!/bin/bash
-
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -1,90 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
-VALIDATION_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/"
-NUM_GPUS=4
-# export CUDA_VISIBLE_DEVICES=4,5
-
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_t2v_finetune"
-  --output_dir "outputs/wan_t2v_finetune"
-  --max_train_steps 5000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 8
-  --num_latent_t 8 
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS 
-  --sp_size $NUM_GPUS 
-  --tp_size $NUM_GPUS
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim $NUM_GPUS
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path $DATA_DIR
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation 
-  --validation_preprocessed_path $VALIDATION_DIR
-  --validation_steps 50
-  --validation_sampling_steps "50" 
-  --validation_guidance_scale "1.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 5e-5
-  --mixed_precision "bf16"
-  --checkpointing_steps 6000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --allow_tf32
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-)
-
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/v1/training/wan_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -1,127 +0,0 @@
-#!/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
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate will-fv
-
-# Basic Info
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=8
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
-VALIDATION_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/"
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name wan_t2v_finetune
-  --output_dir="outputs/wan_t2v_finetune"
-  --max_train_steps=1000
-  --train_batch_size=4
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps=1
-  --num_latent_t 8
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 4
-  --tp_size 4
-  --hsdp_replicate_dim 2
-  --hsdp_shard_dim 4
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 10
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_preprocessed_path "$VALIDATION_DIR"
-  --validation_steps 100
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "1.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate=5e-5
-  --mixed_precision="bf16"
-  --checkpointing_steps=500
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --allow_tf32
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-)
-
-srun torchrun \
--nnodes $SLURM_JOB_NUM_NODES \
--nproc_per_node $NUM_GPUS \
--node_rank $SLURM_PROCID \
--rdzv_backend=c10d \
--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -1,25 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=1 # 2,4,8
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_TYPE="wan"
-DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_t2v/"
-VALIDATION_PATH="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
-
-torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
-    --model_path $MODEL_PATH \
-    --data_merge_path $DATA_MERGE_PATH \
-    --preprocess_video_batch_size 8 \
-    --max_height 480 \
-    --max_width 832 \
-    --num_frames 77 \
-    --dataloader_num_workers 0 \
-    --output_dir=$OUTPUT_DIR \
-    --model_type $MODEL_TYPE \
-    --train_fps 16 \
-    --validation_dataset_file $VALIDATION_PATH \
-    --samples_per_file 8 \
-    --flush_frequency 8 \
-    --preprocess_task "t2v" 
@@ -1,31 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 50,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 50,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 50,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field, fields
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict

 from fastvideo.v1.logger import init_logger

@@ -12,9 +12,7 @@ logger = init_logger(__name__)
 # 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
@dataclass
 class ArchConfig:
-    stacked_params_mapping: List[Tuple[str, str, str]] = field(
-        default_factory=list
-    )  # mapping from huggingface weight names to custom names
+    pass


@dataclass
@@ -12,7 +12,6 @@ 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)
    _supported_attention_backends: Tuple[AttentionBackendEnum, ...] = (
        AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,
@@ -147,9 +147,6 @@ class HunyuanVideoArchConfig(DiTArchConfig):
            r"final_layer.linear.\1",
        })

-    # 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
    in_channels: int = 16
@@ -5,11 +5,13 @@ from typing import List, Optional, Tuple, Union
 from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig


+def is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
@dataclass
 class StepVideoArchConfig(DiTArchConfig):
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda:
-        [lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit()])
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])

    _param_names_mapping: dict = field(
        default_factory=lambda: {
@@ -49,13 +49,9 @@ class WanVideoArchConfig(DiTArchConfig):
            r"blocks.\1.ffn.fc_in.\2",
            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
            r"blocks.\1.ffn.fc_out.\2",
-            r"^blocks\.(\d+)\.norm2\.(.*)$":
+            r"blocks\.(\d+)\.norm2\.(.*)$":
            r"blocks.\1.self_attn_residual_norm.norm.\2",
        })
-
-    # 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(
@@ -32,11 +32,8 @@ class TextEncoderArchConfig(EncoderArchConfig):
    output_past: bool = True
    scalable_attention: bool = True
    tie_word_embeddings: bool = False
-    stacked_params_mapping: List[Tuple[str, str, str]] = field(
-        default_factory=list
-    )  # mapping from huggingface weight names to custom names
+
    tokenizer_kwargs: Dict[str, Any] = field(default_factory=dict)
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [])

    def __post_init__(self) -> None:
        self.tokenizer_kwargs = {
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple
+from typing import Optional

 from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                       ImageEncoderConfig,
@@ -8,14 +8,6 @@ from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                       TextEncoderConfig)


-def _is_transformer_layer(n: str, m) -> bool:
-    return "layers" in n and str.isdigit(n.split(".")[-1])
-
-
-def _is_embeddings(n: str, m) -> bool:
-    return n.endswith("embeddings")
-
-
@dataclass
 class CLIPTextArchConfig(TextEncoderArchConfig):
    vocab_size: int = 49408
@@ -35,15 +27,6 @@ class CLIPTextArchConfig(TextEncoderArchConfig):
    bos_token_id: int = 49406
    eos_token_id: int = 49407
    text_len: int = 77
-    stacked_params_mapping: List[Tuple[str, str,
-                                       str]] = field(default_factory=lambda: [
-                                           # (param_name, shard_name, shard_id)
-                                           ("qkv_proj", "q_proj", "q"),
-                                           ("qkv_proj", "k_proj", "k"),
-                                           ("qkv_proj", "v_proj", "v"),
-                                       ])
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [_is_transformer_layer, _is_embeddings])


@dataclass
@@ -1,23 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple
+from typing import Optional

 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                       TextEncoderConfig)


-def _is_transformer_layer(n: str, m) -> bool:
-    return "layers" in n and str.isdigit(n.split(".")[-1])
-
-
-def _is_embeddings(n: str, m) -> bool:
-    return n.endswith("embed_tokens")
-
-
-def _is_final_norm(n: str, m) -> bool:
-    return n.endswith("norm")
-
-
@dataclass
 class LlamaArchConfig(TextEncoderArchConfig):
    vocab_size: int = 32000
@@ -44,18 +32,6 @@ class LlamaArchConfig(TextEncoderArchConfig):
    head_dim: Optional[int] = None
    hidden_state_skip_layer: int = 2
    text_len: int = 256
-    stacked_params_mapping: List[Tuple[str, str, str]] = field(
-        default_factory=lambda: [
-            # (param_name, shard_name, shard_id)
-            (".qkv_proj", ".q_proj", "q"),
-            (".qkv_proj", ".k_proj", "k"),
-            (".qkv_proj", ".v_proj", "v"),
-            (".gate_up_proj", ".gate_proj", 0),  # type: ignore
-            (".gate_up_proj", ".up_proj", 1),  # type: ignore
-        ])
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda:
-        [_is_transformer_layer, _is_embeddings, _is_final_norm])


@dataclass
@@ -1,23 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple
+from typing import Optional

 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                       TextEncoderConfig)


-def _is_transformer_layer(n: str, m) -> bool:
-    return "block" in n and str.isdigit(n.split(".")[-1])
-
-
-def _is_embeddings(n: str, m) -> bool:
-    return n.endswith("shared")
-
-
-def _is_final_layernorm(n: str, m) -> bool:
-    return n.endswith("final_layer_norm")
-
-
@dataclass
 class T5ArchConfig(TextEncoderArchConfig):
    vocab_size: int = 32128
@@ -41,16 +29,6 @@ class T5ArchConfig(TextEncoderArchConfig):
    eos_token_id: int = 1
    classifier_dropout: float = 0.0
    text_len: int = 512
-    stacked_params_mapping: List[Tuple[str, str,
-                                       str]] = field(default_factory=lambda: [
-                                           # (param_name, shard_name, shard_id)
-                                           (".qkv_proj", ".q", "q"),
-                                           (".qkv_proj", ".k", "k"),
-                                           (".qkv_proj", ".v", "v"),
-                                       ])
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda:
-        [_is_transformer_layer, _is_embeddings, _is_final_layernorm])

    # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py
    def __post_init__(self):
@@ -1,17 +1,19 @@
-# SPDX-License-Identifier: Apache-2.0
+import os
+
 from torchvision import transforms
 from torchvision.transforms import Lambda
+from transformers import AutoTokenizer

-from fastvideo.v1.dataset.parquet_dataset_map_style import (
-    build_parquet_map_style_dataloader)
-from fastvideo.v1.dataset.preprocessing_datasets import (
-    VideoCaptionMergedDataset)
+from fastvideo.v1.dataset.t2v_datasets import T2V_dataset
 from fastvideo.v1.dataset.transform import (CenterCropResizeVideo, Normalize255,
                                            TemporalRandomCrop)
-from fastvideo.v1.dataset.validation_dataset import ValidationDataset
+
+from .parquet_dataset_map_style import build_parquet_map_style_dataloader
+
+__all__ = ["build_parquet_map_style_dataloader"]


-def getdataset(args) -> VideoCaptionMergedDataset:
+def getdataset(args, start_idx=0) -> T2V_dataset:
    temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
    norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
    resize_topcrop = [
@@ -29,14 +31,15 @@ def getdataset(args) -> VideoCaptionMergedDataset:
        *resize_topcrop,
        norm_fun,
    ])
-    return VideoCaptionMergedDataset(data_merge_path=args.data_merge_path,
-                                     args=args,
-                                     transform=transform,
-                                     temporal_sample=temporal_sample,
-                                     transform_topcrop=transform_topcrop)
+    tokenizer_path = os.path.join(args.model_path, "tokenizer")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
+                                              cache_dir=args.cache_dir)
+    if args.dataset == "t2v":
+        return T2V_dataset(args,
+                           transform=transform,
+                           temporal_sample=temporal_sample,
+                           tokenizer=tokenizer,
+                           transform_topcrop=transform_topcrop,
+                           start_idx=start_idx)

-
-__all__ = [
-    "build_parquet_map_style_dataloader", "ValidationDataset",
-    "VideoCaptionMergedDataset"
-]
+    raise NotImplementedError(args.dataset)
@@ -11,7 +11,7 @@ from fastvideo.v1.dataset.parquet_dataset_iterable_style import (
    build_parquet_iterable_style_dataloader)
 from fastvideo.v1.distributed import get_world_rank
 from fastvideo.v1.distributed.parallel_state import (
-    cleanup_dist_env_and_memory, get_local_torch_device,
+    cleanup_dist_env_and_memory, get_torch_device,
    maybe_init_distributed_environment_and_model_parallel)
 from fastvideo.v1.logger import init_logger

@@ -148,8 +148,8 @@ def main() -> None:
                break

            # Move data to device
-            latents = latents.to(get_local_torch_device())
-            embeddings = embeddings.to(get_local_torch_device())
+            latents = latents.to(get_torch_device())
+            embeddings = embeddings.to(get_torch_device())

            # Calculate actual batch size
            batch_size = latents.size(0)
@@ -8,12 +8,11 @@ import torch
 import torch.distributed as dist
 import torch.distributed.checkpoint as dist_cp

-from fastvideo.v1.dataset.dataloader.schema import pyarrow_schema_t2v
 from fastvideo.v1.dataset.parquet_dataset_map_style import (
    build_parquet_map_style_dataloader)
 from fastvideo.v1.distributed import get_world_rank
 from fastvideo.v1.distributed.parallel_state import (
-    cleanup_dist_env_and_memory, get_local_torch_device,
+    cleanup_dist_env_and_memory, get_torch_device,
    maybe_init_distributed_environment_and_model_parallel)
 from fastvideo.v1.logger import init_logger

@@ -68,18 +67,14 @@ def main() -> None:

    # Create DataLoader with proper settings
    dataset, dataloader = build_parquet_map_style_dataloader(
-        args.path,
-        args.batch_size,
-        parquet_schema=pyarrow_schema_t2v,
-        num_data_workers=args.num_data_workers)
+        args.path, args.batch_size, args.num_data_workers)
    logger.info("Initialized dataloader with %d batches", len(dataloader))

    if args.verify_resume:
        # First pass - record latent sums
        first_pass_sums = []
-        for i, batch in enumerate(dataloader):
-            latents = batch['vae_latent']
-            embeddings = batch['text_embedding']
+        for i, (latents, embeddings, masks,
+                caption_text) in enumerate(dataloader):
            latent_sum = latents.sum().item()
            first_pass_sums.append(latent_sum)
            logger.info("Batch %d latent sum: %f", i, latent_sum)
@@ -105,18 +100,14 @@ def main() -> None:

        # Recreate dataloader and load state
        dataset, dataloader = build_parquet_map_style_dataloader(
-            args.path,
-            args.batch_size,
-            parquet_schema=pyarrow_schema_t2v,
-            num_data_workers=args.num_data_workers)
+            args.path, args.batch_size, args.num_data_workers)
        load_states = {"dataloader": dataloader}
        dist_cp.load(load_states, checkpoint_id=checkpoint_dir.as_posix())
        logger.info("Rank %d: Loaded dataloader state from %s",
                    get_world_rank(), checkpoint_dir)

-        for i, batch in enumerate(dataloader):
-            latents = batch['vae_latent']
-            embeddings = batch['text_embedding']
+        for i, (latents, embeddings, masks,
+                caption_text) in enumerate(dataloader):
            latent_sum = latents.sum().item()
            first_pass_sums.append(latent_sum)
            logger.info("Batch %d latent sum: %f",
@@ -125,16 +116,11 @@ def main() -> None:
                break

        dataset, dataloader = build_parquet_map_style_dataloader(
-            args.path,
-            args.batch_size,
-            parquet_schema=pyarrow_schema_t2v,
-            num_data_workers=args.num_data_workers)
+            args.path, args.batch_size, args.num_data_workers)

        # Second pass - verify latent sums match
        second_pass_sums = []
-        for i, batch in enumerate(dataloader):
-            latents = batch['vae_latent']
-            embeddings = batch['text_embedding']
+        for i, (latents, embeddings, masks) in enumerate(dataloader):
            latent_sum = latents.sum().item()
            second_pass_sums.append(latent_sum)
            logger.info("Batch %d latent sum: %f (should match first pass: %f)",
@@ -158,15 +144,14 @@ def main() -> None:
    total_samples = 0
    total_batches = 0
    for _ in range(args.num_epoch):
-        for i, batch in enumerate(dataloader):
-            latents = batch['vae_latent']
-            embeddings = batch['text_embedding']
+        for i, (latents, embeddings, masks,
+                caption_text) in enumerate(dataloader):
            if i >= args.num_batches_per_epoch:
                break

            # Move data to device
-            latents = latents.to(get_local_torch_device())
-            embeddings = embeddings.to(get_local_torch_device())
+            latents = latents.to(get_torch_device())
+            embeddings = embeddings.to(get_torch_device())

            # Calculate actual batch size
            batch_size = latents.size(0)
@@ -26,47 +26,15 @@ pyarrow_schema_i2v = pa.schema([
    pa.field("text_embedding_shape", pa.list_(pa.int64())),
    # e.g., 'bfloat16' or 'float32'
    pa.field("text_embedding_dtype", pa.string()),
+    pa.field("text_attention_mask_bytes", pa.binary()),
+    # e.g., [SeqLen]
+    pa.field("text_attention_mask_shape", pa.list_(pa.int64())),
+    # e.g., 'bool' or 'int8'
+    pa.field("text_attention_mask_dtype", pa.string()),
    #I2V
    pa.field("clip_feature_bytes", pa.binary()),
    pa.field("clip_feature_shape", pa.list_(pa.int64())),
    pa.field("clip_feature_dtype", pa.string()),
-    pa.field("first_frame_latent_bytes", pa.binary()),
-    pa.field("first_frame_latent_shape", pa.list_(pa.int64())),
-    pa.field("first_frame_latent_dtype", pa.string()),
-    # I2V Validation
-    pa.field("pil_image_bytes", pa.binary()),
-    pa.field("pil_image_shape", pa.list_(pa.int64())),
-    pa.field("pil_image_dtype", pa.string()),
-    # --- Metadata ---
-    pa.field("file_name", pa.string()),
-    pa.field("caption", pa.string()),
-    pa.field("media_type", pa.string()),  # 'image' or 'video'
-    pa.field("width", pa.int64()),
-    pa.field("height", pa.int64()),
-    # -- Video-specific (can be null/default for images) ---
-    # Number of frames processed (e.g., 1 for image, N for video)
-    pa.field("num_frames", pa.int64()),
-    pa.field("duration_sec", pa.float64()),
-    pa.field("fps", pa.float64()),
-])
-
-pyarrow_schema_i2v_validation = pa.schema([
-    pa.field("id", pa.string()),
-    # --- Text encoder output tensor ---
-    # Tensors are stored as raw bytes with shape and dtype info for loading
-    pa.field("text_embedding_bytes", pa.binary()),
-    # e.g., [SeqLen, Dim]
-    pa.field("text_embedding_shape", pa.list_(pa.int64())),
-    # e.g., 'bfloat16' or 'float32'
-    pa.field("text_embedding_dtype", pa.string()),
-    #I2V
-    pa.field("clip_feature_bytes", pa.binary()),
-    pa.field("clip_feature_shape", pa.list_(pa.int64())),
-    pa.field("clip_feature_dtype", pa.string()),
-    # I2V Validation
-    pa.field("pil_image_bytes", pa.binary()),
-    pa.field("pil_image_shape", pa.list_(pa.int64())),
-    pa.field("pil_image_dtype", pa.string()),
    # --- Metadata ---
    pa.field("file_name", pa.string()),
    pa.field("caption", pa.string()),
@@ -96,6 +64,11 @@ pyarrow_schema_t2v = pa.schema([
    pa.field("text_embedding_shape", pa.list_(pa.int64())),
    # e.g., 'bfloat16' or 'float32'
    pa.field("text_embedding_dtype", pa.string()),
+    pa.field("text_attention_mask_bytes", pa.binary()),
+    # e.g., [SeqLen]
+    pa.field("text_attention_mask_shape", pa.list_(pa.int64())),
+    # e.g., 'bool' or 'int8'
+    pa.field("text_attention_mask_dtype", pa.string()),
    # --- Metadata ---
    pa.field("file_name", pa.string()),
    pa.field("caption", pa.string()),
@@ -107,26 +80,4 @@ pyarrow_schema_t2v = pa.schema([
    pa.field("num_frames", pa.int64()),
    pa.field("duration_sec", pa.float64()),
    pa.field("fps", pa.float64()),
-])
-
-pyarrow_schema_t2v_validation = pa.schema([
-    pa.field("id", pa.string()),
-    # --- Text encoder output tensor ---
-    # Tensors are stored as raw bytes with shape and dtype info for loading
-    pa.field("text_embedding_bytes", pa.binary()),
-    # e.g., [SeqLen, Dim]
-    pa.field("text_embedding_shape", pa.list_(pa.int64())),
-    # e.g., 'bfloat16' or 'float32'
-    pa.field("text_embedding_dtype", pa.string()),
-    # --- Metadata ---
-    pa.field("file_name", pa.string()),
-    pa.field("caption", pa.string()),
-    pa.field("media_type", pa.string()),  # 'image' or 'video'
-    pa.field("width", pa.int64()),
-    pa.field("height", pa.int64()),
-    # -- Video-specific (can be null/default for images) ---
-    # Number of frames processed (e.g., 1 for image, N for video)
-    pa.field("num_frames", pa.int64()),
-    pa.field("duration_sec", pa.float64()),
-    pa.field("fps", pa.float64()),
-])
+])
@@ -4,7 +4,6 @@ import random
 from typing import Dict, List, Tuple

 import numpy as np
-import pyarrow as pa
 import pyarrow.parquet as pq
 import torch
 import tqdm
@@ -71,12 +70,10 @@ class LatentsParquetIterStyleDataset(IterableDataset):
                 drop_last: bool = True,
                 text_padding_length: int = 512,
                 seed: int = 42,
-                 read_batch_size: int = 32,
-                 parquet_schema: pa.Schema = None):
+                 read_batch_size: int = 32):
        super().__init__()
        self.path = str(path)
        self.batch_size = batch_size
-        self.parquet_schema = parquet_schema
        self.cfg_rate = cfg_rate
        self.text_padding_length = text_padding_length
        self.seed = seed
@@ -3,7 +3,6 @@ import os
 import pickle
 from typing import Any, Dict, List, Tuple

-import pyarrow as pa
 import pyarrow.parquet as pq
 # Torch in general
 import torch
@@ -12,7 +11,7 @@ import tqdm
 from torch.utils.data import Dataset, Sampler
 from torchdata.stateful_dataloader import StatefulDataLoader

-from fastvideo.v1.dataset.utils import collate_rows_from_parquet_schema
+from fastvideo.v1.dataset.utils import collate_latents_embs_masks
 from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
                                      get_world_size)
 from fastvideo.v1.logger import init_logger
@@ -185,12 +184,13 @@ class LatentsParquetMapStyleDataset(Dataset):
    Note: 
    Using parquet for map style dataset is not efficient, we mainly keep it for backward compatibility and debugging.
    """
+    # Modify this in the future if we want to add more keys, for example, in image to video.
+    keys = [("vae_latent", "latent"), "text_embedding"]

    def __init__(
        self,
        path: str,
        batch_size: int,
-        parquet_schema: pa.Schema,
        cfg_rate: float = 0.0,
        seed: int = 42,
        drop_last: bool = True,
@@ -200,12 +200,24 @@ class LatentsParquetMapStyleDataset(Dataset):
        super().__init__()
        self.path = path
        self.cfg_rate = cfg_rate
-        self.parquet_schema = parquet_schema
+        if cfg_rate > 0.0:
+            raise ValueError(
+                "cfg_rate > 0.0 is not supported for now because it will trigger bug when num_data_workers > 0"
+            )
        logger.info("Initializing LatentsParquetMapStyleDataset with path: %s",
                    path)
        self.parquet_files, self.lengths = get_parquet_files_and_length(path)
        self.batch = batch_size
        self.text_padding_length = text_padding_length
+        self._cols = [
+            "vae_latent_bytes",
+            "vae_latent_shape",
+            "text_embedding_bytes",
+            "text_embedding_shape",
+            "text_embedding_dtype",
+            "height",
+            "width",
+        ]
        self.sampler = DP_SP_BatchSampler(
            batch_size=batch_size,
            dataset_size=sum(self.lengths),
@@ -220,7 +232,7 @@ class LatentsParquetMapStyleDataset(Dataset):
                    len(self.parquet_files), sum(self.lengths))

    def get_validation_negative_prompt(
-            self) -> tuple[torch.Tensor, torch.Tensor, str]:
+            self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, str]:
        """
        Get the negative prompt for validation. 
        This method ensures the negative prompt is loaded and cached properly.
@@ -234,23 +246,19 @@ class LatentsParquetMapStyleDataset(Dataset):
        row_dict = read_row_from_parquet_file([file_path], row_idx,
                                              [self.lengths[0]])

-        batch = collate_rows_from_parquet_schema([row_dict],
-                                                 self.parquet_schema,
-                                                 self.text_padding_length,
-                                                 cfg_rate=0.0)
-        negative_prompt = batch['info_list'][0]['prompt']
-        negative_prompt_embedding = batch['text_embedding']
-        negative_prompt_attention_mask = batch['text_attention_mask']
-        if len(negative_prompt_embedding.shape) == 2:
-            negative_prompt_embedding = negative_prompt_embedding.unsqueeze(0)
-        if len(negative_prompt_attention_mask.shape) == 1:
-            negative_prompt_attention_mask = negative_prompt_attention_mask.unsqueeze(
-                0).unsqueeze(0)
-
-        return negative_prompt_embedding, negative_prompt_attention_mask, negative_prompt
+        all_latents_list, all_embs_list, all_masks_list, caption_text_list = collate_latents_embs_masks(
+            [row_dict], self.text_padding_length, self.keys)
+        all_latents, all_embs, all_masks, caption_text = all_latents_list[
+            0], all_embs_list[0], all_masks_list[0], caption_text_list[0]
+        # add batch dimension
+        if len(all_embs.shape) == 2:
+            all_embs = all_embs.unsqueeze(0)
+        if len(all_masks.shape) == 1:
+            all_masks = all_masks.unsqueeze(0).unsqueeze(0)
+        return all_latents, all_embs, all_masks, caption_text

    # PyTorch calls this ONLY because the batch_sampler yields a list
-    def __getitems__(self, indices: List[int]) -> Dict[str, Any]:
+    def __getitems__(self, indices: List[int]):
        """
        Batch fetch using read_row_from_parquet_file for each index.
        """
@@ -259,11 +267,9 @@ class LatentsParquetMapStyleDataset(Dataset):
            for idx in indices
        ]

-        batch = collate_rows_from_parquet_schema(rows,
-                                                 self.parquet_schema,
-                                                 self.text_padding_length,
-                                                 cfg_rate=self.cfg_rate)
-        return batch
+        all_latents, all_embs, all_masks, caption_text = collate_latents_embs_masks(
+            rows, self.text_padding_length, self.keys)
+        return all_latents, all_embs, all_masks, caption_text

    def __len__(self):
        return sum(self.lengths)
@@ -280,7 +286,6 @@ def build_parquet_map_style_dataloader(
        path,
        batch_size,
        num_data_workers,
-        parquet_schema,
        cfg_rate=0.0,
        drop_last=True,
        drop_first_row=False,
@@ -293,7 +298,6 @@ def build_parquet_map_style_dataloader(
        drop_last=drop_last,
        drop_first_row=drop_first_row,
        text_padding_length=text_padding_length,
-        parquet_schema=parquet_schema,
        seed=seed)

    loader = StatefulDataLoader(
@@ -1,615 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import json
-import math
-import os
-import random
-from abc import ABC, abstractmethod
-from collections import Counter
-from dataclasses import dataclass
-from os.path import join as opj
-from typing import Any, Dict, List, Optional, Union
-
-import numpy as np
-import torch
-import torchvision
-from einops import rearrange
-from PIL import Image
-from transformers import AutoTokenizer
-
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-@dataclass
-class PreprocessBatch:
-    """
-    Batch information for dataset processing stages.
-    
-    This class holds all the information about a video-caption or image-caption pair
-    as it moves through the processing pipeline. Fields are populated by different stages.
-    """
-    # Raw metadata
-    path: str
-    cap: Union[str, List[str]]
-    resolution: Optional[Dict] = None
-    fps: Optional[float] = None
-    duration: Optional[float] = None
-
-    # Processed metadata
-    num_frames: Optional[int] = None
-    sample_frame_index: Optional[List[int]] = None
-    sample_num_frames: Optional[int] = None
-
-    # Processed data
-    pixel_values: Optional[torch.Tensor] = None
-    text: Optional[str] = None
-    input_ids: Optional[torch.Tensor] = None
-    cond_mask: Optional[torch.Tensor] = None
-
-    @property
-    def is_video(self) -> bool:
-        """Check if this is a video item."""
-        return self.path.endswith(".mp4")
-
-    @property
-    def is_image(self) -> bool:
-        """Check if this is an image item."""
-        return self.path.endswith(".jpg")
-
-
-class DatasetStage(ABC):
-    """
-    Abstract base class for dataset processing stages.
-    
-    Similar to PipelineStage but designed for dataset preprocessing operations.
-    """
-
-    @abstractmethod
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """
-        Process the dataset batch.
-        
-        Args:
-            batch: Dataset batch to process
-            **kwargs: Additional processing parameters
-            
-        Returns:
-            Processed batch
-        """
-        raise NotImplementedError
-
-
-class DatasetFilterStage(ABC):
-    """
-    Abstract base class for dataset filtering stages.
-    
-    These stages can filter out items during metadata processing.
-    """
-
-    @abstractmethod
-    def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
-        """
-        Check if batch should be kept.
-        
-        Args:
-            batch: Dataset batch to check
-            **kwargs: Additional parameters
-            
-        Returns:
-            True if batch should be kept, False otherwise
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """
-        Process the dataset batch (for non-filtering operations).
-        
-        Args:
-            batch: Dataset batch to process
-            **kwargs: Additional processing parameters
-            
-        Returns:
-            Processed batch
-        """
-        raise NotImplementedError
-
-
-class DataValidationStage(DatasetFilterStage):
-    """Stage for validating data items."""
-
-    def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
-        """
-        Validate data item.
-        
-        Args:
-            batch: Dataset batch to validate
-            
-        Returns:
-            True if valid, False if invalid
-        """
-        # Check for caption
-        if batch.cap is None:
-            return False
-
-        if batch.is_video:
-            # Validate video-specific fields
-            if batch.duration is None or batch.fps is None:
-                return False
-        elif not batch.is_image:
-            return False
-
-        return True
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """Process does nothing for validation - filtering is handled by should_keep."""
-        return batch
-
-
-class ResolutionFilterStage(DatasetFilterStage):
-    """Stage for filtering data items based on resolution constraints."""
-
-    def __init__(self,
-                 max_h_div_w_ratio: float = 17 / 16,
-                 min_h_div_w_ratio: float = 8 / 16,
-                 max_height: int = 1024,
-                 max_width: int = 1024):
-        self.max_h_div_w_ratio = max_h_div_w_ratio
-        self.min_h_div_w_ratio = min_h_div_w_ratio
-        self.max_height = max_height
-        self.max_width = max_width
-
-    def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
-        """
-        Check if data item passes resolution filtering.
-        
-        Args:
-            batch: Dataset batch with resolution information
-            
-        Returns:
-            True if passes filter, False otherwise
-        """
-        # Only apply to videos
-        if not batch.is_video:
-            return True
-
-        if batch.resolution is None:
-            return False
-
-        height = batch.resolution.get("height", None)
-        width = batch.resolution.get("width", None)
-        if height is None or width is None:
-            return False
-
-        # Check aspect ratio
-        aspect = self.max_height / self.max_width
-        hw_aspect_thr = 1.5
-
-        return self.filter_resolution(
-            height,
-            width,
-            max_h_div_w_ratio=hw_aspect_thr * aspect,
-            min_h_div_w_ratio=1 / hw_aspect_thr * aspect,
-        )
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """Process does nothing for resolution filtering - filtering is handled by should_keep."""
-        return batch
-
-    def filter_resolution(self, h: int, w: int, max_h_div_w_ratio: float,
-                          min_h_div_w_ratio: float) -> bool:
-        """Filter based on height/width ratio."""
-        return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio
-
-
-class FrameSamplingStage(DatasetFilterStage):
-    """Stage for temporal frame sampling and indexing."""
-
-    def __init__(self,
-                 num_frames: int,
-                 train_fps: int,
-                 speed_factor: int = 1,
-                 video_length_tolerance_range: float = 5.0,
-                 drop_short_ratio: float = 0.0):
-        self.num_frames = num_frames
-        self.train_fps = train_fps
-        self.speed_factor = speed_factor
-        self.video_length_tolerance_range = video_length_tolerance_range
-        self.drop_short_ratio = drop_short_ratio
-
-    def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
-        """
-        Check if video should be kept based on length constraints.
-        
-        Args:
-            batch: Dataset batch
-            
-        Returns:
-            True if should be kept, False otherwise
-        """
-        if batch.is_image:
-            return True
-
-        if batch.duration is None or batch.fps is None:
-            return False
-
-        num_frames = math.ceil(batch.fps * batch.duration)
-
-        # Check if video is too long
-        if (num_frames / batch.fps > self.video_length_tolerance_range *
-            (self.num_frames / self.train_fps * self.speed_factor)):
-            return False
-
-        # Resample frame indices to check length
-        frame_interval = batch.fps / self.train_fps
-        start_frame_idx = 0
-        frame_indices = np.arange(start_frame_idx, num_frames,
-                                  frame_interval).astype(int)
-
-        # Filter short videos
-        return not (len(frame_indices) < self.num_frames
-                    and random.random() < self.drop_short_ratio)
-
-    def process(self,
-                batch: PreprocessBatch,
-                temporal_sample_fn=None,
-                **kwargs) -> PreprocessBatch:
-        """
-        Process frame sampling for video data items.
-        
-        Args:
-            batch: Dataset batch
-            temporal_sample_fn: Function for temporal sampling
-            
-        Returns:
-            Updated batch with frame sampling info
-        """
-        if batch.is_image:
-            # For images, just add sample info
-            batch.sample_frame_index = [0]
-            batch.sample_num_frames = 1
-            return batch
-
-        assert batch.duration is not None and batch.fps is not None
-        batch.num_frames = math.ceil(batch.fps * batch.duration)
-
-        # Resample frame indices
-        frame_interval = batch.fps / self.train_fps
-        start_frame_idx = 0
-        frame_indices = np.arange(start_frame_idx, batch.num_frames,
-                                  frame_interval).astype(int)
-
-        # Temporal crop if too long
-        if len(frame_indices
-               ) > self.num_frames and temporal_sample_fn is not None:
-            begin_index, end_index = temporal_sample_fn(len(frame_indices))
-            frame_indices = frame_indices[begin_index:end_index]
-
-        batch.sample_frame_index = frame_indices.tolist()
-        batch.sample_num_frames = len(frame_indices)
-
-        return batch
-
-
-class VideoTransformStage(DatasetStage):
-    """Stage for video data transformation."""
-
-    def __init__(self, transform) -> None:
-        self.transform = transform
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """
-        Transform video data.
-        
-        Args:
-            batch: Dataset batch with video information
-            
-        Returns:
-            Batch with transformed video tensor
-        """
-        if not batch.is_video:
-            return batch
-
-        assert os.path.exists(batch.path), f"file {batch.path} do not exist!"
-        assert batch.sample_frame_index is not None, "Frame indices must be set before transformation"
-
-        torchvision_video, _, metadata = torchvision.io.read_video(
-            batch.path, output_format="TCHW")
-        video = torchvision_video[batch.sample_frame_index]
-        if self.transform is not None:
-            video = self.transform(video)
-        video = rearrange(video, "t c h w -> c t h w")
-        video = video.to(torch.uint8)
-
-        h, w = video.shape[-2:]
-        assert (
-            h / w <= 17 / 16 and h / w >= 8 / 16
-        ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({batch.path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}"
-
-        video = video.float() / 127.5 - 1.0
-        batch.pixel_values = video
-        return batch
-
-
-class ImageTransformStage(DatasetStage):
-    """Stage for image data transformation."""
-
-    def __init__(self, transform, transform_topcrop) -> None:
-        self.transform = transform
-        self.transform_topcrop = transform_topcrop
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """
-        Transform image data.
-        
-        Args:
-            batch: Dataset batch with image information
-            
-        Returns:
-            Batch with transformed image tensor
-        """
-        if not batch.is_image:
-            return batch
-
-        image = Image.open(batch.path).convert("RGB")
-        image = torch.from_numpy(np.array(image))
-        image = rearrange(image, "h w c -> c h w").unsqueeze(0)
-
-        if self.transform_topcrop is not None:
-            image = self.transform_topcrop(image)
-        elif self.transform is not None:
-            image = self.transform(image)
-
-        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
-        image = image.float() / 127.5 - 1.0
-        batch.pixel_values = image
-        return batch
-
-
-class TextEncodingStage(DatasetStage):
-    """Stage for text tokenization and encoding."""
-
-    def __init__(self, tokenizer, text_max_length: int, cfg_rate: float = 0.0):
-        self.tokenizer = tokenizer
-        self.text_max_length = text_max_length
-        self.cfg_rate = cfg_rate
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """
-        Process text data.
-        
-        Args:
-            batch: Dataset batch with caption information
-            
-        Returns:
-            Batch with encoded text information
-        """
-        text = batch.cap
-        if not isinstance(text, list):
-            text = [text]
-        text = [random.choice(text)]
-
-        text = text[0] if random.random() > self.cfg_rate else ""
-        text_tokens_and_mask = self.tokenizer(
-            text,
-            max_length=self.text_max_length,
-            padding="max_length",
-            truncation=True,
-            return_attention_mask=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-
-        batch.text = text
-        batch.input_ids = text_tokens_and_mask["input_ids"]
-        batch.cond_mask = text_tokens_and_mask["attention_mask"]
-        return batch
-
-
-class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
-                                torch.distributed.checkpoint.stateful.Stateful):
-    """
-    Merged dataset for video and caption data with stage-based processing.
-    Assumes that data_merge_path is a txt file with the following format:
-    <folder_path>,<json_file_path>
-
-        The folder should contain videos.
-
-        The json file should be a list of dictionaries with the following format:
-        [
-        {
-            "path": "1gGQy4nxyUo-Scene-016.mp4",
-            "resolution": {
-            "width": 1920,
-            "height": 1080
-            },
-            "size": 2439112,
-            "fps": 25.0,
-            "duration": 6.88,
-            "num_frames": 172,
-            "cap": [
-            "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open."
-            ]
-        },
-        ...
-        ]
-    
-    This dataset processes video and image data through a series of stages:
-    - Data validation
-    - Resolution filtering  
-    - Frame sampling
-    - Transformation
-    - Text encoding
-    """
-
-    def __init__(self,
-                 data_merge_path: str,
-                 args,
-                 transform,
-                 temporal_sample,
-                 transform_topcrop,
-                 start_idx: int = 0):
-        self.data_merge_path = data_merge_path
-        self.start_idx = start_idx
-        self.args = args
-        self.temporal_sample = temporal_sample
-
-        # Initialize tokenizer
-        tokenizer_path = os.path.join(args.model_path, "tokenizer")
-        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
-                                                  cache_dir=args.cache_dir)
-
-        # Initialize processing stages
-        self._init_stages(args, transform, transform_topcrop, tokenizer)
-
-        # Process metadata
-        self.processed_batches = self._process_metadata()
-
-    def _init_stages(self, args, transform, transform_topcrop,
-                     tokenizer) -> None:
-        """Initialize all processing stages."""
-        self.validation_stage = DataValidationStage()
-        self.resolution_filter_stage = ResolutionFilterStage(
-            max_height=args.max_height, max_width=args.max_width)
-        self.frame_sampling_stage = FrameSamplingStage(
-            num_frames=args.num_frames,
-            train_fps=args.train_fps,
-            speed_factor=args.speed_factor,
-            video_length_tolerance_range=args.video_length_tolerance_range,
-            drop_short_ratio=args.drop_short_ratio)
-        self.video_transform_stage = VideoTransformStage(transform)
-        self.image_transform_stage = ImageTransformStage(
-            transform, transform_topcrop)
-        self.text_encoding_stage = TextEncodingStage(
-            tokenizer=tokenizer,
-            text_max_length=args.text_max_length,
-            cfg_rate=args.training_cfg_rate)
-
-    def _load_raw_data(self) -> List[Dict]:
-        """Load raw data from JSON files."""
-        # Read folder-annotation pairs
-        with open(self.data_merge_path) as f:
-            folder_anno_pairs = [
-                line.strip().split(",") for line in f if line.strip()
-            ]
-        assert len(
-            folder_anno_pairs) == 1, "Only support one folder-annotation pair"
-        assert len(folder_anno_pairs[0]
-                   ) == 2, "Folder-annotation pair should have two elements"
-        folder, annotation_file = folder_anno_pairs[0]
-
-        data_items: List[Dict] = []
-        with open(annotation_file) as f:
-            data_items = json.load(f)
-
-        # Update paths with folder prefix
-        for item in data_items:
-            item["path"] = opj(folder, item["path"])
-
-        return data_items
-
-    def _process_metadata(self) -> List[PreprocessBatch]:
-        """Process the raw metadata through all filtering stages."""
-        raw_data = self._load_raw_data()
-        processed_batches = []
-
-        # Initialize counters
-        filter_counts = {
-            "validation_failed": 0,
-            "resolution_failed": 0,
-            "frame_sampling_failed": 0
-        }
-        sample_num_frames: List[int] = []
-
-        for item in raw_data:
-            batch = PreprocessBatch(path=item["path"],
-                                    cap=item["cap"],
-                                    resolution=item.get("resolution"),
-                                    fps=item.get("fps"),
-                                    duration=item.get("duration"))
-
-            # Apply filtering stages
-            if not self._apply_filter_stages(batch, filter_counts):
-                continue
-
-            # Apply frame sampling processing
-            batch = self.frame_sampling_stage.process(
-                batch, temporal_sample_fn=self.temporal_sample)
-
-            processed_batches.append(batch)
-            assert batch.sample_num_frames is not None
-            sample_num_frames.append(batch.sample_num_frames)
-
-        self._log_filtering_stats(filter_counts, sample_num_frames,
-                                  len(raw_data), len(processed_batches))
-        return processed_batches
-
-    def _apply_filter_stages(self, batch: PreprocessBatch,
-                             filter_counts: Dict[str, int]) -> bool:
-        """Apply all filter stages and update counters. Returns True if batch should be kept."""
-        if not self.validation_stage.should_keep(batch):
-            filter_counts["validation_failed"] += 1
-            return False
-
-        if not self.resolution_filter_stage.should_keep(batch):
-            filter_counts["resolution_failed"] += 1
-            return False
-
-        if not self.frame_sampling_stage.should_keep(batch):
-            filter_counts["frame_sampling_failed"] += 1
-            return False
-
-        return True
-
-    def _log_filtering_stats(self, filter_counts: Dict[str, int],
-                             sample_num_frames: List[int], before_count: int,
-                             after_count: int):
-        """Log filtering statistics."""
-        logger.info(
-            "validation_failed: %d, resolution_failed: %d, frame_sampling_failed: %d, "
-            "Counter(sample_num_frames): %s, before filter: %d, after filter: %d",
-            filter_counts['validation_failed'],
-            filter_counts['resolution_failed'],
-            filter_counts['frame_sampling_failed'], Counter(sample_num_frames),
-            before_count, after_count)
-
-    def __iter__(self):
-        """Iterate through processed data items."""
-        for idx in range(len(self.processed_batches)):
-            yield self._get_item(idx)
-
-    def __len__(self):
-        return len(self.processed_batches)
-
-    def _get_item(self, idx: int) -> Dict:
-        """Get a single processed data item."""
-        batch = self.processed_batches[idx]
-
-        # Apply transformation stages
-        batch = self.video_transform_stage.process(batch)
-        batch = self.image_transform_stage.process(batch)
-        batch = self.text_encoding_stage.process(batch)
-
-        # Build result dictionary
-        result = {
-            "pixel_values": batch.pixel_values,
-            "text": batch.text,
-            "input_ids": batch.input_ids,
-            "cond_mask": batch.cond_mask,
-            "path": batch.path,
-        }
-
-        # Add video-specific fields
-        if batch.is_video:
-            result.update({"fps": batch.fps, "duration": batch.duration})
-
-        return result
-
-    def state_dict(self) -> Dict[str, Any]:
-        """Return state dict for checkpointing."""
-        return {"processed_batches": self.processed_batches}
-
-    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        """Load state dict from checkpoint."""
-        self.processed_batches = state_dict["processed_batches"]
@@ -0,0 +1,352 @@
+# SPDX-License-Identifier: Apache-2.0
+import json
+import math
+import os
+import random
+from collections import Counter
+from os.path import join as opj
+
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+from PIL import Image
+from torch.utils.data import Dataset
+
+from fastvideo.utils.dataset_utils import DecordInit
+from fastvideo.utils.logging_ import main_print
+
+
+class SingletonMeta(type):
+    _instances: dict[type, 'SingletonMeta'] = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            instance = super().__call__(*args, **kwargs)
+            cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+class DataSetProg(metaclass=SingletonMeta):
+
+    def __init__(self) -> None:
+        self.cap_list: list[dict] = []
+        self.elements: list[int] = []
+        self.num_workers = 1
+        self.n_elements = 0
+        self.worker_elements: dict[int, list[int]] = {}
+        self.n_used_elements: dict[int, int] = {}
+
+    def set_cap_list(self, num_workers, cap_list, n_elements) -> None:
+        self.num_workers = num_workers
+        self.cap_list = cap_list
+        self.n_elements = n_elements
+        self.elements = list(range(n_elements))
+        random.shuffle(self.elements)
+        print(f"n_elements: {len(self.elements)}", flush=True)
+
+        for i in range(self.num_workers):
+            self.n_used_elements[i] = 0
+            per_worker = int(
+                math.ceil(len(self.elements) / float(self.num_workers)))
+            start = i * per_worker
+            end = min(start + per_worker, len(self.elements))
+            self.worker_elements[i] = self.elements[start:end]
+
+    def get_item(self, work_info) -> int:
+        worker_id = 0 if work_info is None else work_info.id
+
+        idx = self.worker_elements[worker_id][
+            self.n_used_elements[worker_id] %
+            len(self.worker_elements[worker_id])]
+        self.n_used_elements[worker_id] += 1
+        return idx
+
+
+dataset_prog = DataSetProg()
+
+
+def filter_resolution(h: int,
+                      w: int,
+                      max_h_div_w_ratio: float = 17 / 16,
+                      min_h_div_w_ratio: float = 8 / 16) -> bool:
+    return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio
+
+
+class T2V_dataset(Dataset):
+
+    def __init__(self,
+                 args,
+                 transform,
+                 temporal_sample,
+                 tokenizer,
+                 transform_topcrop,
+                 start_idx=0) -> None:
+        self.start_idx = start_idx
+        self.data = args.data_merge_path
+        self.num_frames = args.num_frames
+        self.train_fps = args.train_fps
+        self.use_image_num = args.use_image_num
+        self.transform = transform
+        self.transform_topcrop = transform_topcrop
+        self.temporal_sample = temporal_sample
+        self.tokenizer = tokenizer
+        self.text_max_length = args.text_max_length
+        self.cfg = args.cfg
+        self.speed_factor = args.speed_factor
+        self.max_height = args.max_height
+        self.max_width = args.max_width
+        self.drop_short_ratio = args.drop_short_ratio
+        assert self.speed_factor >= 1
+        self.v_decoder = DecordInit()
+        self.video_length_tolerance_range = args.video_length_tolerance_range
+        self.support_Chinese = True
+        if "mt5" not in args.text_encoder_name:
+            self.support_Chinese = False
+
+        cap_list = self.get_cap_list()
+
+        assert len(cap_list) > 0
+        cap_list, self.sample_num_frames = self.define_frame_index(cap_list)
+        self.lengths = self.sample_num_frames
+
+        n_elements = len(cap_list)
+        dataset_prog.set_cap_list(args.dataloader_num_workers, cap_list,
+                                  n_elements)
+
+        print(f"video length: {len(dataset_prog.cap_list)}", flush=True)
+
+    def set_checkpoint(self, n_used_elements):
+        for i in range(len(dataset_prog.n_used_elements)):
+            dataset_prog.n_used_elements[i] = n_used_elements
+
+    def __len__(self):
+        return dataset_prog.n_elements
+
+    def __getitem__(self, idx):
+
+        data = self.get_data(idx)
+        return data
+
+    def get_data(self, idx) -> dict:
+        path = dataset_prog.cap_list[idx]["path"]
+        if path.endswith(".mp4"):
+            return self.get_video(idx)
+        else:
+            return self.get_image(idx)
+
+    def get_video(self, idx) -> dict:
+        video_path = dataset_prog.cap_list[idx]["path"]
+        assert os.path.exists(video_path), f"file {video_path} do not exist!"
+        frame_indices = dataset_prog.cap_list[idx]["sample_frame_index"]
+
+        torchvision_video, _, metadata = torchvision.io.read_video(
+            video_path, output_format="TCHW")
+        video = torchvision_video[frame_indices]
+        video = self.transform(video)
+        video = rearrange(video, "t c h w -> c t h w")
+        video = video.to(torch.uint8)
+        assert video.dtype == torch.uint8
+
+        h, w = video.shape[-2:]
+        assert (
+            h / w <= 17 / 16 and h / w >= 8 / 16
+        ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({video_path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}"
+
+        video = video.float() / 127.5 - 1.0
+
+        text = dataset_prog.cap_list[idx]["cap"]
+        if not isinstance(text, list):
+            text = [text]
+        text = [random.choice(text)]
+
+        text = text[0] if random.random() > self.cfg else ""
+        text_tokens_and_mask = self.tokenizer(
+            text,
+            max_length=self.text_max_length,
+            padding="max_length",
+            truncation=True,
+            return_attention_mask=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        input_ids = text_tokens_and_mask["input_ids"]
+        cond_mask = text_tokens_and_mask["attention_mask"]
+        return dict(pixel_values=video,
+                    text=text,
+                    input_ids=input_ids,
+                    cond_mask=cond_mask,
+                    path=video_path,
+                    fps=dataset_prog.cap_list[idx]["fps"],
+                    duration=dataset_prog.cap_list[idx]["duration"])
+
+    def get_image(self, idx) -> dict:
+        image_data = dataset_prog.cap_list[
+            idx]  # [{'path': path, 'cap': cap}, ...]
+
+        image = Image.open(image_data["path"]).convert("RGB")  # [h, w, c]
+        image = torch.from_numpy(np.array(image))  # [h, w, c]
+        image = rearrange(image, "h w c -> c h w").unsqueeze(0)  #  [1 c h w]
+        # for i in image:
+        #     h, w = i.shape[-2:]
+        #     assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only image with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But found ratio is {round(h / w, 2)} with the shape of {i.shape}'
+
+        image = (self.transform_topcrop(image) if "human_images"
+                 in image_data["path"] else self.transform(image)
+                 )  #  [1 C H W] -> num_img [1 C H W]
+        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
+
+        image = image.float() / 127.5 - 1.0
+
+        caps: list[str] = (image_data["cap"] if isinstance(
+            image_data["cap"], list) else [image_data["cap"]])
+        caps = [random.choice(caps)]
+        text = caps
+        input_ids, cond_mask = [], []
+        single_text = text[0] if random.random() > self.cfg else ""
+        text_tokens_and_mask = self.tokenizer(
+            single_text,
+            max_length=self.text_max_length,
+            padding="max_length",
+            truncation=True,
+            return_attention_mask=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        input_ids = text_tokens_and_mask["input_ids"]  # 1, l
+        cond_mask = text_tokens_and_mask["attention_mask"]  # 1, l
+        return dict(
+            pixel_values=image,
+            text=text,
+            input_ids=input_ids,
+            cond_mask=cond_mask,
+            path=image_data["path"],
+        )
+
+    def define_frame_index(self, cap_list) -> tuple[list[dict], list[int]]:
+        new_cap_list = []
+        sample_num_frames = []
+        cnt_too_long = 0
+        cnt_too_short = 0
+        cnt_no_cap = 0
+        cnt_no_resolution = 0
+        cnt_resolution_mismatch = 0
+        cnt_movie = 0
+        cnt_img = 0
+        for i in cap_list:
+            path = i["path"]
+            cap = i.get("cap", None)
+            # ======no caption=====
+            if cap is None:
+                cnt_no_cap += 1
+                continue
+            if path.endswith(".mp4"):
+                # ======no fps and duration=====
+                duration = i.get("duration", None)
+                fps = i.get("fps", None)
+                if fps is None or duration is None:
+                    continue
+
+                # ======resolution mismatch=====
+                resolution = i.get("resolution", None)
+                if resolution is None:
+                    cnt_no_resolution += 1
+                    continue
+                else:
+                    if (resolution.get("height", None) is None
+                            or resolution.get("width", None) is None):
+                        cnt_no_resolution += 1
+                        continue
+                    height, width = i["resolution"]["height"], i["resolution"][
+                        "width"]
+                    aspect = self.max_height / self.max_width
+                    hw_aspect_thr = 1.5
+                    is_pick = filter_resolution(
+                        height,
+                        width,
+                        max_h_div_w_ratio=hw_aspect_thr * aspect,
+                        min_h_div_w_ratio=1 / hw_aspect_thr * aspect,
+                    )
+                    if not is_pick:
+                        print("resolution mismatch")
+                        cnt_resolution_mismatch += 1
+                        continue
+
+                # if path == 'finetrainers/3dgs-dissolve/videos/1.mp4':
+                #     from IPython import embed; embed()
+                i["num_frames"] = math.ceil(fps * duration)
+                # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration.
+                if i["num_frames"] / fps > self.video_length_tolerance_range * (
+                        self.num_frames / self.train_fps * self.speed_factor
+                ):  # too long video is not suitable for this training stage (self.num_frames)
+                    cnt_too_long += 1
+                    continue
+
+                # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24)
+                frame_interval = fps / self.train_fps
+                start_frame_idx = 0
+                frame_indices = np.arange(start_frame_idx, i["num_frames"],
+                                          frame_interval).astype(int)
+
+                # comment out it to enable dynamic frames training
+                if (len(frame_indices) < self.num_frames
+                        and random.random() < self.drop_short_ratio):
+                    cnt_too_short += 1
+                    continue
+
+                #  too long video will be temporal-crop randomly
+                if len(frame_indices) > self.num_frames:
+                    begin_index, end_index = self.temporal_sample(
+                        len(frame_indices))
+                    frame_indices = frame_indices[begin_index:end_index]
+                    # frame_indices = frame_indices[:self.num_frames]  # head crop
+                i["sample_frame_index"] = frame_indices.tolist()
+                new_cap_list.append(i)
+                i["sample_num_frames"] = len(
+                    i["sample_frame_index"]
+                )  # will use in dataloader(group sampler)
+                sample_num_frames.append(i["sample_num_frames"])
+            elif path.endswith(".jpg"):  # image
+                cnt_img += 1
+                new_cap_list.append(i)
+                i["sample_num_frames"] = 1
+                sample_num_frames.append(i["sample_num_frames"])
+            else:
+                raise NameError(
+                    f"Unknown file extension {path.split('.')[-1]}, only support .mp4 for video and .jpg for image"
+                )
+        # import ipdb;ipdb.set_trace()
+        main_print(
+            f"no_cap: {cnt_no_cap}, too_long: {cnt_too_long}, too_short: {cnt_too_short}, "
+            f"no_resolution: {cnt_no_resolution}, resolution_mismatch: {cnt_resolution_mismatch}, "
+            f"Counter(sample_num_frames): {Counter(sample_num_frames)}, cnt_movie: {cnt_movie}, cnt_img: {cnt_img}, "
+            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}"
+        )
+        return new_cap_list, sample_num_frames
+
+    def decord_read(self, path, frame_indices) -> torch.Tensor:
+        decord_vr = self.v_decoder(path)
+        video_data = decord_vr.get_batch(frame_indices).asnumpy()
+        video_data = torch.from_numpy(video_data)
+        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
+        return video_data
+
+    def read_jsons(self, data) -> list[dict]:
+        cap_lists = []
+        with open(data) as f:
+            folder_anno = [
+                i.strip().split(",") for i in f.readlines()
+                if len(i.strip()) > 0
+            ]
+        print(folder_anno)
+        for folder, anno in folder_anno:
+            with open(anno) as f:
+                sub_list = json.load(f)
+            for i in range(len(sub_list)):
+                sub_list[i]["path"] = opj(folder, sub_list[i]["path"])
+            cap_lists += sub_list
+        return cap_lists
+
+    def get_cap_list(self) -> list:
+        cap_lists = self.read_jsons(self.data)[self.start_idx:]
+        return cap_lists
@@ -1,5 +1,4 @@
-import random
-from typing import Any, Dict, List, cast
+from typing import Any, Dict, List

 import numpy as np
 import torch
@@ -21,7 +20,7 @@ def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
        return t[:padding_length], torch.ones(padding_length)


-def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> Dict[str, Any]:
+def get_torch_tensors_from_row_dict(row_dict, keys) -> Dict[str, Any]:
    """
    Get the latents and prompts from a row dictionary.
    """
@@ -43,10 +42,7 @@ def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> Dict[str, Any]:
            bytes = row_dict[f"{key}_bytes"]

        # TODO (peiyuan): read precision
-        if key == 'text_embedding' and random.random() < cfg_rate:
-            data = np.zeros((512, 4096), dtype=np.float32)
-        else:
-            data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy()
+        data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy()
        data = torch.from_numpy(data)
        if len(data.shape) == 3:
            B, L, D = data.shape
@@ -57,11 +53,8 @@ def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> Dict[str, Any]:


 def collate_latents_embs_masks(
-        batch_to_process,
-        text_padding_length,
-        keys,
-        cfg_rate=0.0
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str]]:
+        batch_to_process, text_padding_length,
+        keys) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str]]:
    # Initialize tensors to hold padded embeddings and masks
    all_latents = []
    all_embs = []
@@ -70,7 +63,7 @@ def collate_latents_embs_masks(
    # Process each row individually
    for i, row in enumerate(batch_to_process):
        # Get tensors from row
-        data = get_torch_tensors_from_row_dict(row, keys, cfg_rate)
+        data = get_torch_tensors_from_row_dict(row, keys)
        latents, emb = data["vae_latent"], data["text_embedding"]

        padded_emb, mask = pad(emb, text_padding_length)
@@ -90,131 +83,3 @@ def collate_latents_embs_masks(
    all_masks = torch.stack(all_masks)

    return all_latents, all_embs, all_masks, caption_text
-
-
-def collate_rows_from_parquet_schema(rows,
-                                     parquet_schema,
-                                     text_padding_length,
-                                     cfg_rate=0.0) -> Dict[str, Any]:
-    """
-    Collate rows from parquet files based on the provided schema.
-    Dynamically processes tensor fields based on schema and returns batched data.
-    
-    Args:
-        rows: List of row dictionaries from parquet files
-        parquet_schema: PyArrow schema defining the structure of the data
-    
-    Returns:
-        Dict containing batched tensors and metadata
-    """
-    if not rows:
-        return cast(Dict[str, Any], {})
-
-    # Initialize containers for different data types
-    batch_data: Dict[str, Any] = {}
-
-    # Get tensor and metadata field names from schema (fields ending with '_bytes')
-    tensor_fields = []
-    metadata_fields = []
-    for field in parquet_schema.names:
-        if field.endswith('_bytes'):
-            shape_field = field.replace('_bytes', '_shape')
-            dtype_field = field.replace('_bytes', '_dtype')
-            tensor_name = field.replace('_bytes', '')
-            tensor_fields.append(tensor_name)
-            assert shape_field in parquet_schema.names, f"Shape field {shape_field} not found in schema for field {field}. Currently we only support *_bytes fields for tensors."
-            assert dtype_field in parquet_schema.names, f"Dtype field {dtype_field} not found in schema for field {field}. Currently we only support *_bytes fields for tensors."
-        elif not field.endswith('_shape') and not field.endswith('_dtype'):
-            # Only add actual metadata fields, not the shape/dtype helper fields
-            metadata_fields.append(field)
-
-    # Process each tensor field
-    for tensor_name in tensor_fields:
-        tensor_list = []
-
-        for row in rows:
-            # Get tensor data from row using the existing helper function pattern
-            shape_key = f"{tensor_name}_shape"
-            bytes_key = f"{tensor_name}_bytes"
-
-            if shape_key in row and bytes_key in row:
-                shape = row[shape_key]
-                bytes_data = row[bytes_key]
-
-                if len(bytes_data) == 0:
-                    tensor = torch.zeros(0, dtype=torch.bfloat16)
-                else:
-                    # Convert bytes to tensor using float32 as default
-                    if tensor_name == 'text_embedding' and random.random(
-                    ) < cfg_rate:
-                        data = np.zeros((512, 4096), dtype=np.float32)
-                    else:
-                        data = np.frombuffer(
-                            bytes_data, dtype=np.float32).reshape(shape).copy()
-                    tensor = torch.from_numpy(data)
-                    # if len(data.shape) == 3:
-                    #     B, L, D = tensor.shape
-                    #     assert B == 1, "Batch size must be 1"
-                    #     tensor = tensor.squeeze(0)
-
-                tensor_list.append(tensor)
-            else:
-                # Handle missing tensor data
-                tensor_list.append(torch.zeros(0, dtype=torch.bfloat16))
-
-        # Stack tensors with special handling for text embeddings
-        if tensor_name == 'text_embedding':
-            # Handle text embeddings with padding
-            padded_tensors = []
-            attention_masks = []
-
-            for tensor in tensor_list:
-                if tensor.numel() > 0:
-                    padded_tensor, mask = pad(tensor, text_padding_length)
-                    padded_tensors.append(padded_tensor)
-                    attention_masks.append(mask)
-                else:
-                    # Handle empty embeddings - assume default embedding dimension
-                    padded_tensors.append(
-                        torch.zeros(text_padding_length,
-                                    768,
-                                    dtype=torch.bfloat16))
-                    attention_masks.append(torch.zeros(text_padding_length))
-
-            batch_data[tensor_name] = torch.stack(padded_tensors)
-            batch_data['text_attention_mask'] = torch.stack(attention_masks)
-        else:
-            # Stack all tensors to preserve batch consistency
-            # Don't filter out None or empty tensors as this breaks batch sizing
-            try:
-                batch_data[tensor_name] = torch.stack(tensor_list)
-            except ValueError as e:
-                shapes = [
-                    t.shape
-                    if t is not None and hasattr(t, 'shape') else 'None/Invalid'
-                    for t in tensor_list
-                ]
-                raise ValueError(
-                    f"Failed to stack tensors for field '{tensor_name}'. "
-                    f"Tensor shapes: {shapes}. "
-                    f"All tensors in a batch must have compatible shapes. "
-                    f"Original error: {e}") from e
-
-    # Process metadata fields into info_list
-    info_list = []
-    for row in rows:
-        info = {}
-        for field in metadata_fields:
-            info[field] = row.get(field, "")
-
-        # Add prompt field for backward compatibility
-        info["prompt"] = info.get("caption", "")
-        info_list.append(info)
-
-    batch_data['info_list'] = info_list
-
-    # Add caption_text for backward compatibility
-    if info_list and 'caption' in info_list[0]:
-        batch_data['caption_text'] = [info['caption'] for info in info_list]
-
-    return batch_data
@@ -1,103 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# adapted from: https://github.com/a-r-r-o-w/finetrainers/blob/main/finetrainers/data/dataset.py
-import os
-import pathlib
-
-import datasets
-import torch
-
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.vision_utils import load_image, load_video
-
-logger = init_logger(__name__)
-
-
-class ValidationDataset(torch.utils.data.IterableDataset):
-
-    def __init__(self, filename: str):
-        super().__init__()
-
-        self.filename = pathlib.Path(filename)
-        # get directory of filename
-        self.dir = os.path.abspath(self.filename.parent)
-
-        if not self.filename.exists():
-            raise FileNotFoundError(
-                f"File {self.filename.as_posix()} does not exist")
-
-        if self.filename.suffix == ".csv":
-            data = datasets.load_dataset("csv",
-                                         data_files=self.filename.as_posix(),
-                                         split="train")
-        elif self.filename.suffix == ".json":
-            data = datasets.load_dataset("json",
-                                         data_files=self.filename.as_posix(),
-                                         split="train",
-                                         field="data")
-        elif self.filename.suffix == ".parquet":
-            data = datasets.load_dataset("parquet",
-                                         data_files=self.filename.as_posix(),
-                                         split="train")
-        elif self.filename.suffix == ".arrow":
-            data = datasets.load_dataset("arrow",
-                                         data_files=self.filename.as_posix(),
-                                         split="train")
-        else:
-            _SUPPORTED_FILE_FORMATS = [".csv", ".json", ".parquet", ".arrow"]
-            raise ValueError(
-                f"Unsupported file format {self.filename.suffix} for validation dataset. Supported formats are: {_SUPPORTED_FILE_FORMATS}"
-            )
-
-        self._data = data.to_iterable_dataset()
-
-    def __iter__(self):
-        for sample in self._data:
-            # For consistency reasons, we mandate that "caption" is always present in the validation dataset.
-            # However, since the model specifications use "prompt", we create an alias here.
-            sample["prompt"] = sample["caption"]
-
-            # Load image or video if the path is provided
-            # TODO(aryan): need to handle custom columns here for control conditions
-            sample["image"] = None
-            sample["video"] = None
-
-            if sample.get("image_path", None) is not None:
-                image_path = sample["image_path"]
-                image_path = os.path.join(self.dir, image_path)
-                if not pathlib.Path(image_path).is_file(
-                ) and not image_path.startswith("http"):
-                    logger.warning("Image file %s does not exist.", image_path)
-                else:
-                    sample["image"] = load_image(image_path)
-
-            if sample.get("video_path", None) is not None:
-                video_path = sample["video_path"]
-                video_path = os.path.join(self.dir, video_path)
-                if not pathlib.Path(video_path).is_file(
-                ) and not video_path.startswith("http"):
-                    logger.warning("Video file %s does not exist.", video_path)
-                else:
-                    sample["video"] = load_video(video_path)
-
-            if sample.get("control_image_path", None) is not None:
-                control_image_path = sample["control_image_path"]
-                control_image_path = os.path.join(self.dir, control_image_path)
-                if not pathlib.Path(control_image_path).is_file(
-                ) and not control_image_path.startswith("http"):
-                    logger.warning("Control Image file %s does not exist.",
-                                   control_image_path)
-                else:
-                    sample["control_image"] = load_image(control_image_path)
-
-            if sample.get("control_video_path", None) is not None:
-                control_video_path = sample["control_video_path"]
-                control_video_path = os.path.join(self.dir, control_video_path)
-                if not pathlib.Path(control_video_path).is_file(
-                ) and not control_video_path.startswith("http"):
-                    logger.warning("Control Video file %s does not exist.",
-                                   control_video_path)
-                else:
-                    sample["control_video"] = load_video(control_video_path)
-
-            sample = {k: v for k, v in sample.items() if v is not None}
-            yield sample
@@ -3,10 +3,10 @@
 from fastvideo.v1.distributed.communication_op import *
 from fastvideo.v1.distributed.parallel_state import (
    cleanup_dist_env_and_memory, get_dp_group, get_dp_rank, get_dp_world_size,
-    get_local_torch_device, get_sp_group, get_sp_parallel_rank,
-    get_sp_world_size, get_tp_group, get_tp_rank, get_tp_world_size,
-    get_world_group, get_world_rank, get_world_size,
-    init_distributed_environment, initialize_model_parallel,
+    get_sp_group, get_sp_parallel_rank, get_sp_world_size, get_torch_device,
+    get_tp_group, get_tp_rank, get_tp_world_size, get_world_group,
+    get_world_rank, get_world_size, init_distributed_environment,
+    initialize_model_parallel,
    maybe_init_distributed_environment_and_model_parallel,
    model_parallel_is_initialized)
 from fastvideo.v1.distributed.utils import *
@@ -40,5 +40,5 @@ __all__ = [
    "get_tp_world_size",

    # Get torch device
-    "get_local_torch_device",
+    "get_torch_device",
 ]
@@ -36,7 +36,6 @@ from unittest.mock import patch

 import torch
 import torch.distributed
-import torch.distributed as dist
 from torch.distributed import Backend, ProcessGroup, ReduceOp

 import fastvideo.v1.envs as envs
@@ -693,7 +692,6 @@ class GroupCoordinator:


 _WORLD: Optional[GroupCoordinator] = None
-_NODE: Optional[GroupCoordinator] = None


 def get_world_group() -> GroupCoordinator:
@@ -701,11 +699,6 @@ def get_world_group() -> GroupCoordinator:
    return _WORLD


-def get_node_group() -> GroupCoordinator:
-    assert _NODE is not None, ("node group is not initialized")
-    return _NODE
-
-
 def init_world_group(ranks: List[int], local_rank: int,
                     backend: str) -> GroupCoordinator:
    return GroupCoordinator(
@@ -717,18 +710,6 @@ def init_world_group(ranks: List[int], local_rank: int,
    )


-def init_node_group(local_rank: int, backend: str):
-    cpu_group = get_world_group().cpu_group
-    node_ranks = same_node_ranks(cpu_group)
-    node_size = len(node_ranks)
-    all_node_ranks = [
-        list(range(i * node_size, (i + 1) * node_size))
-        for i in range(dist.get_world_size() // node_size)
-    ]
-    global _NODE
-    _NODE = init_model_parallel_group(all_node_ranks, local_rank, backend)
-
-
 def init_model_parallel_group(
    group_ranks: List[List[int]],
    local_rank: int,
@@ -801,8 +782,6 @@ def init_distributed_environment(
    else:
        assert _WORLD.world_size == torch.distributed.get_world_size(), (
            "world group already initialized with a different world size")
-    # Init a group for each node
-    init_node_group(local_rank, backend)


 _SP: Optional[GroupCoordinator] = None
@@ -925,7 +904,7 @@ def get_dp_rank() -> int:
    return get_dp_group().rank_in_group


-def get_local_torch_device() -> torch.device:
+def get_torch_device() -> torch.device:
    """Return the torch device for the current rank."""
    return torch.device(f"cuda:{envs.LOCAL_RANK}")

@@ -1042,22 +1021,17 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
            "torch._C._host_emptyCache() only available in Pytorch >=2.5")


-def same_node_ranks(pg: Union[ProcessGroup, StatelessProcessGroup],
-                    source_rank: int = 0) -> List[int]:
+def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
+                        source_rank: int = 0) -> List[bool]:
    """
-    This is a collective operation that returns ranks that are in the same node
+    This is a collective operation that returns if each rank is in the same node
    as the source rank. It tests if processes are attached to the same
    memory system (shared access to shared memory).
-    Args:
-        pg: the global process group to test
-        source_rank: the rank to test against
-    Returns:
-        A list of ranks that are in the same node as the source rank.
    """
    if isinstance(pg, ProcessGroup):
        assert torch.distributed.get_backend(
            pg) != torch.distributed.Backend.NCCL, (
-                "same_node_ranks should be tested with a non-NCCL group.")
+                "in_the_same_node_as should be tested with a non-NCCL group.")
        # local rank inside the group
        rank = torch.distributed.get_rank(group=pg)
        world_size = torch.distributed.get_world_size(group=pg)
@@ -1129,7 +1103,7 @@ def same_node_ranks(pg: Union[ProcessGroup, StatelessProcessGroup],
            rank_data = pg.broadcast_obj(is_in_the_same_node, src=i)
            aggregated_data += rank_data

-    return [i for i, x in enumerate(aggregated_data.tolist()) if x == 1]
+    return [x == 1 for x in aggregated_data.tolist()]


 def initialize_tensor_parallel_group(
@@ -58,10 +58,8 @@ class FastVideoArgs:

    output_type: str = "pil"

-    use_cpu_offload: bool = True  # For DiT
+    use_cpu_offload: bool = True
    use_fsdp_inference: bool = True
-    text_encoder_offload: bool = True
-    pin_cpu_memory: bool = True

    # STA (Sliding Tile Attention) parameters
    mask_strategy_file_path: Optional[str] = None
@@ -210,7 +208,7 @@ class FastVideoArgs:
            "--use-cpu-offload",
            action=StoreBoolean,
            help=
-            "Use CPU offload for DiT inference. Enable if run out of memory with FSDP.",
+            "Use CPU offload for model inference. Enable if run out of memory with FSDP.",
        )
        parser.add_argument(
            "--use-fsdp-inference",
@@ -218,19 +216,7 @@ class FastVideoArgs:
            help=
            "Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.",
        )
-        parser.add_argument(
-            "--text-encoder-cpu-offload",
-            action=StoreBoolean,
-            help=
-            "Use CPU offload for text encoder. Enable if run out of memory.",
-        )
-        parser.add_argument(
-            "--pin-cpu-memory",
-            action=StoreBoolean,
-            help=
-            "Pin memory for CPU offload. Only added as a temp workaround if it throws \"CUDA error: invalid argument\". "
-            "Should be enabled in almost all cases",
-        )
+
        parser.add_argument(
            "--disable-autocast",
            action=StoreBoolean,
@@ -398,12 +384,11 @@ class TrainingArgs(FastVideoArgs):
    # diffusion setting
    ema_decay: float = 0.0
    ema_start_step: int = 0
-    training_cfg_rate: float = 0.0
+    cfg: float = 0.0
    precondition_outputs: bool = False

    # validation & logs
-    validation_dataset_file: str = ""
-    validation_preprocessed_path: str = ""
+    validation_prompt_dir: str = ""
    validation_sampling_steps: str = ""
    validation_guidance_scale: str = ""
    validation_steps: float = 0.0
@@ -542,7 +527,7 @@ class TrainingArgs(FastVideoArgs):
                            type=int,
                            default=0,
                            help="Step to start EMA")
-        parser.add_argument("--training-cfg-rate",
+        parser.add_argument("--cfg",
                            type=float,
                            help="Classifier-free guidance scale")
        parser.add_argument(
@@ -551,12 +536,9 @@ class TrainingArgs(FastVideoArgs):
            help="Whether to precondition the outputs of the model")

        # Validation and logging
-        parser.add_argument("--validation-dataset-file",
+        parser.add_argument("--validation-prompt-dir",
                            type=str,
-                            help="Path to unprocessed validation dataset")
-        parser.add_argument("--validation-preprocessed-path",
-                            type=str,
-                            help="Path to processed validation dataset")
+                            help="Directory containing validation prompts")
        parser.add_argument("--validation-sampling-steps",
                            type=str,
                            help="Validation sampling steps")
@@ -6,7 +6,6 @@ from typing import Optional, Tuple, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.distributed.tensor import DTensor

 from fastvideo.v1.layers.custom_op import CustomOp

@@ -71,12 +70,7 @@ class RMSNorm(CustomOp):
        x = x * torch.rsqrt(variance + self.variance_epsilon)
        x = x.to(orig_dtype)
        if self.has_weight:
-            # TODO(wenxuan): When using CPU offload, FSDP has a bug that doesn't unwrap DTensor in final_layer_norm.
-            # Report this
-            if isinstance(self.weight, DTensor):
-                x = x * self.weight.to(x.device).full_tensor()
-            else:
-                x = x * self.weight
+            x = x * self.weight
        if residual is None:
            return x
        else:
@@ -14,7 +14,6 @@ class BaseDiT(nn.Module, ABC):
    _fsdp_shard_conditions: list = []
    _compile_conditions: list = []
    _param_names_mapping: dict
-    _reverse_param_names_mapping: dict
    hidden_size: int
    num_attention_heads: int
    num_channels_latents: int
@@ -79,7 +78,6 @@ 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 = {}
    # Ensure these instance attributes are properly defined in subclasses
    hidden_size: int
@@ -442,8 +442,6 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
    _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

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

 class StepVideoModel(BaseDiT):
    # (Optional) Keep the same attribute for compatibility with splitting, etc.
-    _fsdp_shard_conditions = StepVideoConfig()._fsdp_shard_conditions
+    _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
    _supported_attention_backends = StepVideoConfig(
    )._supported_attention_backends
@@ -518,7 +518,6 @@ class WanTransformer3DModel(CachableDiT):
    _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

    def __init__(self, config: WanVideoConfig, hf_config: dict[str,
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from abc import ABC, abstractmethod
-from dataclasses import field
-from typing import List, Optional, Tuple
+from typing import Optional, Tuple

 import torch
 from torch import nn
@@ -13,9 +12,6 @@ from fastvideo.v1.platforms import AttentionBackendEnum


 class TextEncoder(nn.Module, ABC):
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [])
-    _stacked_params_mapping: List[Tuple[str, str,
-                                        str]] = field(default_factory=list)
    _supported_attention_backends: Tuple[
        AttentionBackendEnum,
        ...] = TextEncoderConfig()._supported_attention_backends
@@ -23,8 +19,6 @@ class TextEncoder(nn.Module, ABC):
    def __init__(self, config: TextEncoderConfig) -> None:
        super().__init__()
        self.config = config
-        self._fsdp_shard_conditions = config._fsdp_shard_conditions
-        self._stacked_params_mapping = config.arch_config.stacked_params_mapping
        if not self.supported_attention_backends:
            raise ValueError(
                f"Subclass {self.__class__.__name__} must define _supported_attention_backends"
@@ -596,7 +596,12 @@ class CLIPVisionModel(ImageEncoder):
    #        ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
-
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
        params_dict = dict(self.named_parameters())
        loaded_params: Set[str] = set()
        layer_count = len(self.vision_model.encoder.layers)
@@ -615,8 +620,7 @@ class CLIPVisionModel(ImageEncoder):
                if layer_idx >= layer_count:
                    continue

-            for (param_name, weight_name,
-                 shard_id) in self.config.arch_config.stacked_params_mapping:
+            for (param_name, weight_name, shard_id) in stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
@@ -369,7 +369,14 @@ class LlamaModel(TextEncoder):

    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
-
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
        params_dict = dict(self.named_parameters())
        loaded_params: Set[str] = set()
        for name, loaded_weight in weights:
@@ -399,7 +406,7 @@ class LlamaModel(TextEncoder):
                    continue
                else:
                    name = kv_scale_name
-            for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
@@ -494,7 +494,7 @@ class T5Stack(nn.Module):
                attention_mask=attention_mask,
                attn_metadata=attn_metadata,
            )
-        hidden_states = self.final_layer_norm.forward(hidden_states)
+        hidden_states = self.final_layer_norm.forward_native(hidden_states)
        return hidden_states


@@ -631,13 +631,19 @@ class UMT5EncoderModel(TextEncoder):

    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q", "q"),
+            (".qkv_proj", ".k", "k"),
+            (".qkv_proj", ".v", "v"),
+        ]
        params_dict = dict(self.named_parameters())
        loaded_params: Set[str] = set()
        for name, loaded_weight in weights:
            loaded = False
            if "decoder" in name or "lm_head" in name:
                continue
-            for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
@@ -10,20 +10,17 @@ from copy import deepcopy
 from typing import Any, Generator, Iterable, List, Optional, Tuple, cast

 import torch
-import torch.distributed as dist
 import torch.nn as nn
 from safetensors.torch import load_file as safetensors_load_file
 from transformers import AutoImageProcessor, AutoTokenizer
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME

 from fastvideo.v1.configs.models import EncoderConfig
-from fastvideo.v1.distributed import get_local_torch_device
+from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.hf_transformer_utils import get_diffusers_config
-from fastvideo.v1.models.loader.fsdp_load import (init_device_mesh,
-                                                  maybe_load_fsdp_model,
-                                                  shard_model)
+from fastvideo.v1.models.loader.fsdp_load import maybe_load_fsdp_model
 from fastvideo.v1.models.loader.utils import set_default_torch_dtype
 from fastvideo.v1.models.loader.weight_utils import (
    filter_duplicate_safetensors_files, filter_files_not_needed_for_inference,
@@ -166,19 +163,16 @@ class TextEncoderLoader(ComponentLoader):
        return hf_folder, hf_weights_files, use_safetensors

    def _get_weights_iterator(
-            self,
-            source: "Source",
-            to_cpu: bool = True
+            self, source: "Source"
    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
        """Get an iterator for the model weights based on the load format."""
        hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
            source.model_or_path, source.fall_back_to_pt,
            source.allow_patterns_overrides)
        if use_safetensors:
-            weights_iterator = safetensors_weights_iterator(
-                hf_weights_files, to_cpu)
+            weights_iterator = safetensors_weights_iterator(hf_weights_files)
        else:
-            weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu)
+            weights_iterator = pt_weights_iterator(hf_weights_files)

        if self.counter_before_loading_weights == 0.0:
            self.counter_before_loading_weights = time.perf_counter()
@@ -187,11 +181,10 @@ class TextEncoderLoader(ComponentLoader):
                for (name, tensor) in weights_iterator)

    def _get_all_weights(
-            self,
-            model_config: Any,
-            model: nn.Module,
-            model_path: str,
-            to_cpu: bool = True
+        self,
+        model_config: Any,
+        model: nn.Module,
+        model_path: str,
    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
        primary_weights = TextEncoderLoader.Source(
            model_path,
@@ -200,14 +193,14 @@ class TextEncoderLoader(ComponentLoader):
            allow_patterns_overrides=getattr(model, "allow_patterns_overrides",
                                             None),
        )
-        yield from self._get_weights_iterator(primary_weights, to_cpu)
+        yield from self._get_weights_iterator(primary_weights)

        secondary_weights = cast(
            Iterable[TextEncoderLoader.Source],
            getattr(model, "secondary_weights", ()),
        )
        for source in secondary_weights:
-            yield from self._get_weights_iterator(source, to_cpu)
+            yield from self._get_weights_iterator(source)

    def load(self, model_path: str, architecture: str,
             fastvideo_args: FastVideoArgs):
@@ -240,22 +233,16 @@ class TextEncoderLoader(ComponentLoader):
            encoder_precision = fastvideo_args.pipeline_config.text_encoder_precisions[
                1]

-        target_device = get_local_torch_device()
+        target_device = get_torch_device()
        # TODO(will): add support for other dtypes
        return self.load_model(model_path, encoder_config, target_device,
-                               fastvideo_args, encoder_precision)
+                               encoder_precision)

    def load_model(self,
                   model_path: str,
                   model_config: EncoderConfig,
                   target_device: torch.device,
-                   fastvideo_args: FastVideoArgs,
                   dtype: str = "fp16"):
-        use_cpu_offload = fastvideo_args.text_encoder_offload and len(
-            getattr(model_config, "_fsdp_shard_conditions", [])) > 0
-
-        if fastvideo_args.text_encoder_offload:
-            target_device = torch.device("cpu")
        with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]):
            with target_device:
                architectures = getattr(model_config, "architectures", [])
@@ -264,26 +251,12 @@ class TextEncoderLoader(ComponentLoader):

            weights_to_load = {name for name, _ in model.named_parameters()}
            loaded_weights = model.load_weights(
-                self._get_all_weights(model_config, model, model_path,
-                                      use_cpu_offload))
+                self._get_all_weights(model_config, model, model_path))
            self.counter_after_loading_weights = time.perf_counter()
            logger.info(
                "Loading weights took %.2f seconds",
                self.counter_after_loading_weights -
                self.counter_before_loading_weights)
-
-            if use_cpu_offload:
-                mesh = init_device_mesh(
-                    "cuda",
-                    mesh_shape=(1, dist.get_world_size()),
-                    mesh_dim_names=("offload", "replicate"),
-                )
-                shard_model(model,
-                            cpu_offload=True,
-                            reshard_after_forward=True,
-                            mesh=mesh["offload"],
-                            fsdp_shard_conditions=model._fsdp_shard_conditions,
-                            pin_cpu_memory=fastvideo_args.pin_cpu_memory)
            # We only enable strict check for non-quantized models
            # that have loaded weights tracking currently.
            # if loaded_weights is not None:
@@ -317,10 +290,10 @@ class ImageEncoderLoader(TextEncoderLoader):
        encoder_config = fastvideo_args.pipeline_config.image_encoder_config
        encoder_config.update_model_arch(model_config)

-        target_device = get_local_torch_device()
+        target_device = get_torch_device()
        # TODO(will): add support for other dtypes
        return self.load_model(
-            model_path, encoder_config, target_device, fastvideo_args,
+            model_path, encoder_config, target_device,
            fastvideo_args.pipeline_config.image_encoder_precision)


@@ -373,7 +346,7 @@ class VAELoader(ComponentLoader):
        with set_default_torch_dtype(PRECISION_TO_TYPE[
                fastvideo_args.pipeline_config.vae_precision]):
            vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-            vae = vae_cls(vae_config).to(get_local_torch_device())
+            vae = vae_cls(vae_config).to(get_torch_device())

        # Find all safetensors files
        safetensors_list = glob.glob(
@@ -432,7 +405,7 @@ class TransformerLoader(ComponentLoader):
                "hf_config": hf_config
            },
            weight_dir_list=safetensors_list,
-            device=get_local_torch_device(),
+            device=get_torch_device(),
            hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
            hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
            cpu_offload=fastvideo_args.use_cpu_offload,
@@ -69,7 +69,6 @@ def maybe_load_fsdp_model(
    fsdp_inference: bool = False,
    output_dtype: Optional[torch.dtype] = None,
    training_mode: bool = True,
-    pin_cpu_memory: bool = True,
 ) -> torch.nn.Module:
    """
    Load the model with FSDP if is training, else load the model without FSDP.
@@ -102,12 +101,9 @@ def maybe_load_fsdp_model(
                cpu_offload=cpu_offload,
                reshard_after_forward=True,
                mp_policy=mp_policy,
-                mesh=device_mesh,
-                fsdp_shard_conditions=model._fsdp_shard_conditions,
-                pin_cpu_memory=pin_cpu_memory)
+                mesh=device_mesh)

-    weight_iterator = safetensors_weights_iterator(
-        weight_dir_list, to_cpu=cpu_offload, async_broadcast=not cpu_offload)
+    weight_iterator = safetensors_weights_iterator(weight_dir_list)
    param_names_mapping_fn = get_param_names_mapping(model._param_names_mapping)
    load_model_from_full_model_state_dict(
        model,
@@ -130,13 +126,12 @@ def maybe_load_fsdp_model(

 def shard_model(
    model,
+    *,
    cpu_offload: bool,
    reshard_after_forward: bool = True,
-    mp_policy: Optional[MixedPrecisionPolicy] = MixedPrecisionPolicy(),  # noqa
+    mp_policy: Optional[MixedPrecisionPolicy] = None,
+    dp_mesh: Optional[DeviceMesh] = None,
    mesh: Optional[DeviceMesh] = None,
-    fsdp_shard_conditions: Optional[List[Callable[[str, nn.Module],
-                                                  bool]]] = None,
-    pin_cpu_memory: bool = True,
 ) -> None:
    """
    Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API.
@@ -155,28 +150,19 @@ def shard_model(
        reshard_after_forward (bool): Whether to reshard parameters and buffers after
            the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy
            from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy.
-        mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism.
+        dp_mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism.
            Default to None.
-        fsdp_shard_conditions (Optional[List[Callable[[str, nn.Module], bool]]]): A list of functions to determine
-            which modules to shard with FSDP. 

    Raises:
        ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered.
    """
-    if fsdp_shard_conditions is None or len(fsdp_shard_conditions) == 0:
-        logger.warning(
-            "The FSDP shard condition list is empty or None. No modules will be sharded in %s",
-            type(model).__name__)
-        return
-
    fsdp_kwargs = {
        "reshard_after_forward": reshard_after_forward,
        "mesh": mesh,
        "mp_policy": mp_policy,
    }
    if cpu_offload:
-        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(
-            pin_memory=pin_cpu_memory)
+        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy()

    # iterating in reverse to start with
    # lowest-level modules first
@@ -186,7 +172,7 @@ def shard_model(
    for n, m in reversed(list(model.named_modules())):
        if any([
                shard_condition(n, m)
-                for shard_condition in fsdp_shard_conditions
+                for shard_condition in model._fsdp_shard_conditions
        ]):
            fully_shard(m, **fsdp_kwargs)
            num_layers_sharded += 1
@@ -195,6 +181,7 @@ def shard_model(
        raise ValueError(
            "No layer modules were sharded. Please check if shard conditions are working as expected."
        )
+
    # Finally shard the entire model to account for any stragglers
    fully_shard(model, **fsdp_kwargs)

@@ -235,17 +222,10 @@ def load_model_from_full_model_state_dict(
    used_keys = set()
    sharded_sd = {}
    to_merge_params: DefaultDict[str, Dict[Any, Any]] = defaultdict(dict)
-    reverse_param_names_mapping = {}
-    assert param_names_mapping is not None
-
-    # iterate over all the weights to sync broadcast before use
-    full_sd_iterator = list(full_sd_iterator)  # type: ignore
    for source_param_name, full_tensor in full_sd_iterator:
+        assert param_names_mapping is not None
        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
@@ -280,7 +260,6 @@ 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()) - used_keys
    if unused_keys:
        logger.warning("Found new parameters in meta state dict: %s",
@@ -11,11 +11,9 @@ from typing import Generator, List, Optional, Tuple, Union
 import filelock
 import huggingface_hub.constants
 import torch
-import torch.distributed as dist
 from safetensors.torch import safe_open
 from tqdm.auto import tqdm

-from fastvideo.v1.distributed.parallel_state import get_node_group
 from fastvideo.v1.logger import init_logger

 logger = init_logger(__name__)
@@ -120,77 +118,36 @@ _BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elap


 def safetensors_weights_iterator(
-    hf_weights_files: List[str],
-    to_cpu: bool = False,
-    async_broadcast: bool = False
+    hf_weights_files: List[str]
 ) -> Generator[Tuple[str, torch.Tensor], None, None]:
-    """Iterate over the weights in the model safetensor files.
-    Args:
-        hf_weights_files: List of safetensor files to load.
-        to_cpu: Whether to load the weights to CPU. If False, will load to the GPU device bound to the current process.
-        async_broadcast: Whether to overlap loading from disk and broadcasting to other ranks. If True,
-            must iterate over all the weights before use. Only use if to_cpu is False.
-    """
-    local_rank = get_node_group().rank
-    device = f"cuda:{local_rank}" if not to_cpu else "cpu"
-    enable_tqdm = not torch.distributed.is_initialized() or get_node_group(
-    ).rank == 0
-    assert not (async_broadcast
-                and to_cpu), "Cannot broadcast weights when loading to CPU"
-
-    handles = []
+    """Iterate over the weights in the model safetensor files."""
+    enable_tqdm = not torch.distributed.is_initialized(
+    ) or torch.distributed.get_rank() == 0
    for st_file in tqdm(
            hf_weights_files,
            desc="Loading safetensors checkpoint shards",
            disable=not enable_tqdm,
            bar_format=_BAR_FORMAT,
    ):
-        with safe_open(st_file, framework="pt", device=device) as f:
+        with safe_open(st_file, framework="pt") as f:
            for name in f.keys():  # noqa: SIM118
-                if to_cpu:
-                    param = f.get_tensor(name)
-                else:
-                    if local_rank == 0:
-                        param = f.get_tensor(name)
-                    else:
-                        shape = f.get_slice(name).get_shape()
-                        param = torch.empty(shape, device=device)
-                    # broadcast to local ranks
-                    # TODO(Wenxuan): scatter instead of broadcast
-                    if get_node_group().world_size > 1:
-                        group = get_node_group().device_group
-                        if async_broadcast:
-                            handle = dist.broadcast(param,
-                                                    src=dist.get_global_rank(
-                                                        group, 0),
-                                                    async_op=True)
-                            handles.append(handle)
-                        else:
-                            dist.broadcast(param,
-                                           src=dist.get_global_rank(group, 0))
+                param = f.get_tensor(name)
                yield name, param

-        if async_broadcast:
-            for handle in handles:
-                handle.wait()
-

 def pt_weights_iterator(
-    hf_weights_files: List[str],
-    to_cpu: bool = True  # default to CPU for text encoder
+    hf_weights_files: List[str]
 ) -> Generator[Tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model bin/pt files."""
-    local_rank = get_node_group().rank
-    device = f"cuda:{local_rank}" if not to_cpu else "cpu"
-    enable_tqdm = not torch.distributed.is_initialized() or get_node_group(
-    ).rank == 0
+    enable_tqdm = not torch.distributed.is_initialized(
+    ) or torch.distributed.get_rank() == 0
    for bin_file in tqdm(
            hf_weights_files,
            desc="Loading pt checkpoint shards",
            disable=not enable_tqdm,
            bar_format=_BAR_FORMAT,
    ):
-        state = torch.load(bin_file, map_location=device, weights_only=True)
+        state = torch.load(bin_file, map_location="cpu", weights_only=True)
        yield from state.items()
        del state

@@ -1,11 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0

 import os
-import tempfile
 from typing import Callable, List, Optional, Tuple, Union
-from urllib.parse import unquote, urlparse

-import imageio
 import numpy as np
 import PIL.Image
 import PIL.ImageOps
@@ -89,7 +86,6 @@ def normalize(
    return 2.0 * images - 1.0


-# adapted from diffusers.utils import load_image
 def load_image(
    image: Union[str, PIL.Image.Image],
    convert_method: Optional[Callable[[PIL.Image.Image],
@@ -135,85 +131,6 @@ def load_image(
    return image


-# adapted from diffusers.utils import load_video
-def load_video(
-    video: str,
-    convert_method: Optional[Callable[[List[PIL.Image.Image]],
-                                      List[PIL.Image.Image]]] = None,
-) -> List[PIL.Image.Image]:
-    """
-    Loads `video` to a list of PIL Image.
-    Args:
-        video (`str`):
-            A URL or Path to a video to convert to a list of PIL Image format.
-        convert_method (Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]], *optional*):
-            A conversion method to apply to the video after loading it. When set to `None` the images will be converted
-            to "RGB".
-    Returns:
-        `List[PIL.Image.Image]`:
-            The video as a list of PIL images.
-    """
-    is_url = video.startswith("http://") or video.startswith("https://")
-    is_file = os.path.isfile(video)
-    was_tempfile_created = False
-
-    if not (is_url or is_file):
-        raise ValueError(
-            f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path."
-        )
-
-    if is_url:
-        response = requests.get(video, stream=True)
-        if response.status_code != 200:
-            raise ValueError(
-                f"Failed to download video. Status code: {response.status_code}"
-            )
-
-        parsed_url = urlparse(video)
-        file_name = os.path.basename(unquote(parsed_url.path))
-
-        suffix = os.path.splitext(file_name)[1] or ".mp4"
-        with tempfile.NamedTemporaryFile(suffix=suffix,
-                                         delete=False) as temp_file:
-            video_path = temp_file.name
-            video_data = response.iter_content(chunk_size=8192)
-            for chunk in video_data:
-                temp_file.write(chunk)
-
-        video = video_path
-
-    pil_images = []
-    if video.endswith(".gif"):
-        gif = PIL.Image.open(video)
-        try:
-            while True:
-                pil_images.append(gif.copy())
-                gif.seek(gif.tell() + 1)
-        except EOFError:
-            pass
-
-    else:
-        try:
-            imageio.plugins.ffmpeg.get_exe()
-        except AttributeError:
-            raise AttributeError(
-                "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg"
-            ) from None
-
-        with imageio.get_reader(video) as reader:
-            # Read all frames
-            for frame in reader:
-                pil_images.append(PIL.Image.fromarray(frame))
-
-    if was_tempfile_created:
-        os.remove(video_path)
-
-    if convert_method is not None:
-        pil_images = convert_method(pil_images)
-
-    return pil_images
-
-
 def get_default_height_width(
    image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
    vae_scale_factor: int,
@@ -11,7 +11,6 @@ import pprint
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, List, Optional, Union

-import PIL.Image
 import torch

 from fastvideo.v1.attention import AttentionMetadata
@@ -38,8 +37,6 @@ class ForwardBatch:
    # Image inputs
    image_path: Optional[str] = None
    image_embeds: List[torch.Tensor] = field(default_factory=list)
-    pil_image: Optional[PIL.Image.Image] = None
-    preprocessed_image: Optional[torch.Tensor] = None

    # Text inputs
    prompt: Optional[Union[str, List[str]]] = None
@@ -151,11 +148,7 @@ class TrainingBatch:
    latents: Optional[torch.Tensor] = None
    encoder_hidden_states: Optional[torch.Tensor] = None
    encoder_attention_mask: Optional[torch.Tensor] = None
-    # i2v
-    preprocessed_image: Optional[torch.Tensor] = None
-    image_embeds: Optional[torch.Tensor] = None
-    image_latents: Optional[torch.Tensor] = None
-    infos: Optional[List[Dict[str, Any]]] = None
+    info: Optional[Dict[str, Any]] = None

    # Transformer inputs
    noisy_model_input: Optional[torch.Tensor] = None
@@ -165,9 +158,6 @@ class TrainingBatch:

    attn_metadata: Optional[AttentionMetadata] = None

-    # input kwargs
-    input_kwargs: Optional[Dict[str, Any]] = None
-
    # Training loss
    loss: torch.Tensor | None = None

@@ -2,23 +2,20 @@
 import gc
 import multiprocessing
 import os
-from collections import defaultdict
 from concurrent.futures import ProcessPoolExecutor
-from itertools import chain
 from typing import Any, Dict, List, Optional

 import numpy as np
-import PIL.Image
 import pyarrow as pa
 import pyarrow.parquet as pq
 import torch
 from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
 from tqdm import tqdm

 from fastvideo.v1.configs.sample import SamplingParam
-from fastvideo.v1.dataset import ValidationDataset, getdataset
-from fastvideo.v1.dataset.preprocessing_datasets import PreprocessBatch
-from fastvideo.v1.distributed import get_local_torch_device
+from fastvideo.v1.dataset import getdataset
+from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
@@ -49,7 +46,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
        # Initialize class variables for data sharing
        self.video_data: Dict[str, Any] = {}  # Store video metadata and paths
        self.latent_data: Dict[str, Any] = {}  # Store latent tensors
-        self.preprocess_validation(fastvideo_args, args)
+        self.preprocess_validation_text(fastvideo_args, args)
        self.preprocess_video_and_text(fastvideo_args, args)

    def get_extra_features(self, valid_data: Dict[str, Any],
@@ -61,206 +58,39 @@ class BasePreprocessPipeline(ComposedPipelineBase):
        """Get the schema fields for the pipeline type. Override in subclasses."""
        raise NotImplementedError

-    def create_record_for_schema(self,
-                                 preprocess_batch: PreprocessBatch,
-                                 schema: pa.Schema,
-                                 strict: bool = False) -> Dict[str, Any]:
-        """Create a record for the Parquet dataset using a generic schema-based approach.
-        
-        Args:
-            preprocess_batch: The batch containing the data to extract
-            schema: PyArrow schema defining the expected fields
-            strict: If True, raises an exception when required fields are missing or unfilled
-            
-        Returns:
-            Dictionary record matching the schema
-            
-        Raises:
-            ValueError: If strict=True and required fields are missing or unfilled
-        """
-        record = {}
-        unfilled_fields = []
-
-        for field in schema.names:
-            field_filled = False
-
-            if field.endswith('_bytes'):
-                # Handle binary tensor data - convert numpy array or tensor to bytes
-                tensor_name = field.replace('_bytes', '')
-                tensor_data = getattr(preprocess_batch, tensor_name, None)
-                if tensor_data is not None:
-                    try:
-                        if hasattr(tensor_data, 'numpy'):  # torch tensor
-                            record[field] = tensor_data.cpu().numpy().tobytes()
-                            field_filled = True
-                        elif hasattr(tensor_data, 'tobytes'):  # numpy array
-                            record[field] = tensor_data.tobytes()
-                            field_filled = True
-                        else:
-                            raise ValueError(
-                                f"Unsupported tensor type for field {field}: {type(tensor_data)}"
-                            )
-                    except Exception as e:
-                        if strict:
-                            raise ValueError(
-                                f"Failed to convert tensor {tensor_name} to bytes: {e}"
-                            ) from e
-                        record[field] = b''  # Empty bytes for missing data
-                else:
-                    record[field] = b''  # Empty bytes for missing data
-
-            elif field.endswith('_shape'):
-                # Handle tensor shape info
-                tensor_name = field.replace('_shape', '')
-                tensor_data = getattr(preprocess_batch, tensor_name, None)
-                if tensor_data is not None and hasattr(tensor_data, 'shape'):
-                    record[field] = list(tensor_data.shape)
-                    field_filled = True
-                else:
-                    record[field] = []
-
-            elif field.endswith('_dtype'):
-                # Handle tensor dtype info
-                tensor_name = field.replace('_dtype', '')
-                tensor_data = getattr(preprocess_batch, tensor_name, None)
-                if tensor_data is not None and hasattr(tensor_data, 'dtype'):
-                    record[field] = str(tensor_data.dtype)
-                    field_filled = True
-                else:
-                    record[field] = 'unknown'
-
-            elif field in ['width', 'height', 'num_frames']:
-                # Handle integer metadata fields
-                value = getattr(preprocess_batch, field, None)
-                if value is not None:
-                    try:
-                        record[field] = int(value)
-                        field_filled = True
-                    except (ValueError, TypeError) as e:
-                        if strict:
-                            raise ValueError(
-                                f"Failed to convert field {field} to int: {e}"
-                            ) from e
-                        record[field] = 0
-                else:
-                    record[field] = 0
-
-            elif field in ['duration_sec', 'fps']:
-                # Handle float metadata fields
-                # Map schema field names to batch attribute names
-                attr_name = 'duration' if field == 'duration_sec' else field
-                value = getattr(preprocess_batch, attr_name, None)
-                if value is not None:
-                    try:
-                        record[field] = float(value)
-                        field_filled = True
-                    except (ValueError, TypeError) as e:
-                        if strict:
-                            raise ValueError(
-                                f"Failed to convert field {field} to float: {e}"
-                            ) from e
-                        record[field] = 0.0
-                else:
-                    record[field] = 0.0
-
-            else:
-                # Handle string fields (id, file_name, caption, media_type, etc.)
-                # Map common schema field names to batch attribute names
-                attr_name = field
-                if field == 'caption':
-                    attr_name = 'text'
-                elif field == 'file_name':
-                    attr_name = 'path'
-                elif field == 'id':
-                    # Generate ID from path if available
-                    path_value = getattr(preprocess_batch, 'path', None)
-                    if path_value:
-                        import os
-                        record[field] = os.path.basename(path_value).split(
-                            '.')[0]
-                        field_filled = True
-                    else:
-                        record[field] = ""
-                    continue
-                elif field == 'media_type':
-                    # Determine media type from path
-                    path_value = getattr(preprocess_batch, 'path', None)
-                    if path_value:
-                        record[field] = 'video' if path_value.endswith(
-                            '.mp4') else 'image'
-                        field_filled = True
-                    else:
-                        record[field] = ""
-                    continue
-
-                value = getattr(preprocess_batch, attr_name, None)
-                if value is not None:
-                    record[field] = str(value)
-                    field_filled = True
-                else:
-                    record[field] = ""
-
-            # Track unfilled fields
-            if not field_filled:
-                unfilled_fields.append(field)
-
-        # Handle strict mode
-        if strict and unfilled_fields:
-            raise ValueError(
-                f"Required fields were not filled: {unfilled_fields}")
-
-        # Log unfilled fields as warning if not in strict mode
-        if unfilled_fields:
-            logger.warning(
-                "Some fields were not filled and got default values: %s",
-                unfilled_fields)
-
-        return record
-
    def create_record(
            self,
            video_name: str,
            vae_latent: np.ndarray,
            text_embedding: np.ndarray,
-            valid_data: Dict[str, Any],
+            text_attention_mask: np.ndarray,
+            valid_data: Optional[Dict[str, Any]],
            idx: int,
            extra_features: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Create a record for the Parquet dataset."""
        record = {
-            "id":
-            video_name,
-            "vae_latent_bytes":
-            vae_latent.tobytes(),
-            "vae_latent_shape":
-            list(vae_latent.shape),
-            "vae_latent_dtype":
-            str(vae_latent.dtype),
-            "text_embedding_bytes":
-            text_embedding.tobytes(),
-            "text_embedding_shape":
-            list(text_embedding.shape),
-            "text_embedding_dtype":
-            str(text_embedding.dtype),
-            "file_name":
-            video_name,
-            "caption":
-            valid_data["text"][idx] if len(valid_data["text"]) > 0 else "",
-            "media_type":
-            "video",
+            "id": video_name,
+            "vae_latent_bytes": vae_latent.tobytes(),
+            "vae_latent_shape": list(vae_latent.shape),
+            "vae_latent_dtype": str(vae_latent.dtype),
+            "text_embedding_bytes": text_embedding.tobytes(),
+            "text_embedding_shape": list(text_embedding.shape),
+            "text_embedding_dtype": str(text_embedding.dtype),
+            "text_attention_mask_bytes": text_attention_mask.tobytes(),
+            "text_attention_mask_shape": list(text_attention_mask.shape),
+            "text_attention_mask_dtype": str(text_attention_mask.dtype),
+            "file_name": video_name,
+            "caption": valid_data["text"][idx] if valid_data else "",
+            "media_type": "video",
            "width":
-            valid_data["pixel_values"][idx].shape[-2]
-            if len(valid_data["pixel_values"]) > 0 else 0,
+            valid_data["pixel_values"][idx].shape[-2] if valid_data else 0,
            "height":
-            valid_data["pixel_values"][idx].shape[-1]
-            if len(valid_data["pixel_values"]) > 0 else 0,
+            valid_data["pixel_values"][idx].shape[-1] if valid_data else 0,
            "num_frames":
            vae_latent.shape[1] if len(vae_latent.shape) > 1 else 0,
            "duration_sec":
-            float(valid_data["duration"][idx])
-            if len(valid_data["duration"]) > 0 else 0.0,
-            "fps":
-            float(valid_data["fps"][idx])
-            if len(valid_data["fps"]) > 0 else 0.0,
+            float(valid_data["duration"][idx]) if valid_data else 0.0,
+            "fps": float(valid_data["fps"][idx]) if valid_data else 0.0,
        }
        if extra_features:
            record.update(extra_features)
@@ -273,6 +103,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                                            "combined_parquet_dataset")
        os.makedirs(combined_parquet_dir, exist_ok=True)
        local_rank = int(os.getenv("RANK", 0))
+        world_size = int(os.getenv("WORLD_SIZE", 1))

        # Get how many samples have already been processed
        start_idx = 0
@@ -283,10 +114,14 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                    start_idx += table.num_rows

        # Loading dataset
-        train_dataset = getdataset(args)
-
+        train_dataset = getdataset(args, start_idx=start_idx)
+        sampler = DistributedSampler(train_dataset,
+                                     rank=local_rank,
+                                     num_replicas=world_size,
+                                     shuffle=False)
        train_dataloader = DataLoader(
            train_dataset,
+            sampler=sampler,
            batch_size=args.preprocess_video_batch_size,
            num_workers=args.dataloader_num_workers,
        )
@@ -328,8 +163,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                # VAE
                with torch.autocast("cuda", dtype=torch.float32):
                    latents = self.get_module("vae").encode(
-                        valid_data["pixel_values"].to(
-                            get_local_torch_device())).mean
+                        valid_data["pixel_values"].to(get_torch_device())).mean

                # Get extra features if needed
                extra_features = self.get_extra_features(
@@ -381,6 +215,8 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                # Convert tensors to numpy arrays
                vae_latent = latent.cpu().numpy()
                text_embedding = prompt_embeds[idx].cpu().numpy()
+                text_attention_mask = prompt_attention_mask[idx].cpu().numpy(
+                ).astype(np.uint8)

                # Get extra features for this sample if needed
                sample_extra_features = {}
@@ -397,6 +233,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                    video_name=video_name,
                    vae_latent=vae_latent,
                    text_embedding=text_embedding,
+                    text_attention_mask=text_attention_mask,
                    valid_data=valid_data,
                    idx=idx,
                    extra_features=sample_extra_features)
@@ -448,7 +285,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                num_processed_samples = 0
                self.all_tables = []

-    def preprocess_validation(self, fastvideo_args: FastVideoArgs, args):
+    def preprocess_validation_text(self, fastvideo_args: FastVideoArgs, args):
        """Process validation text prompts and save them to parquet files.
        
        This base implementation handles the common validation text processing logic.
@@ -459,32 +296,22 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                                              "validation_parquet_dataset")
        os.makedirs(validation_parquet_dir, exist_ok=True)

-        validation_dataset = ValidationDataset(args.validation_dataset_file)
+        with open(args.validation_prompt_txt, encoding="utf-8") as file:
+            lines = file.readlines()
+        prompts = [line.strip() for line in lines]

        # Prepare batch data for Parquet dataset
        batch_data = []
        sampling_param = SamplingParam.from_pretrained(
            fastvideo_args.model_path)
        if sampling_param.negative_prompt:
-            negative_prompt = {
-                'caption': sampling_param.negative_prompt,
-                'image_path': None,
-                'video_path': None,
-            }
-            validation_iterable = chain([negative_prompt], validation_dataset)
-        else:
-            negative_prompt = None
-            validation_iterable = validation_dataset
-
+            prompts = [sampling_param.negative_prompt] + prompts
        # Add progress bar for validation text preprocessing
-        pbar = tqdm(enumerate(validation_iterable),
+        pbar = tqdm(enumerate(prompts),
                    desc="Processing validation prompts",
                    unit="prompt")
-        for idx, sample in pbar:
+        for prompt_idx, prompt in pbar:
            with torch.inference_mode():
-                prompt = sample["caption"]
-                is_negative_prompt = idx == 0
-
                # Text Encoder
                batch = ForwardBatch(
                    data_type="video",
@@ -511,43 +338,15 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                "Shape after removing padding - Embeddings: %s, Mask: %s",
                text_embedding.shape, text_attention_mask.shape)

-            extra_features = {}
-            if not is_negative_prompt:
-                height = sample["height"]
-                width = sample["width"]
-                if "image_path" in sample and "video_path" in sample:
-                    raise ValueError(
-                        "Only one of image_path or video_path should be provided"
-                    )
-
-                if "image" in sample:
-                    extra_features = self.preprocess_image(
-                        sample["image"], height, width, fastvideo_args)
-
-                if "video" in sample:
-                    extra_features = self.preprocess_video(
-                        sample["video"], height, width, fastvideo_args)
-
-            # Get extra features for this sample if needed
-            sample_extra_features = {}
-            if extra_features:
-                for key, value in extra_features.items():
-                    if isinstance(value, torch.Tensor):
-                        sample_extra_features[key] = value.cpu().numpy()
-                    else:
-                        sample_extra_features[key] = value
-
-            valid_data = defaultdict(list)
-            valid_data["text"] = [prompt]
-
            # Create record for Parquet dataset
            record = self.create_record(video_name=file_name,
                                        vae_latent=np.array([],
                                                            dtype=np.float32),
                                        text_embedding=text_embedding,
-                                        valid_data=valid_data,
+                                        text_attention_mask=text_attention_mask,
+                                        valid_data=None,
                                        idx=0,
-                                        extra_features=sample_extra_features)
+                                        extra_features=None)
            batch_data.append(record)

            logger.info("Saved validation sample: %s", file_name)
@@ -621,15 +420,6 @@ class BasePreprocessPipeline(ComposedPipelineBase):
            del table
            gc.collect()  # Force garbage collection

-    def preprocess_image(self, image: PIL.Image.Image, height: int, width: int,
-                         fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
-        return {}
-
-    def preprocess_video(self, video: list[PIL.Image.Image], height: int,
-                         width: int,
-                         fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
-        return {}
-
    def _flush_tables(self, num_processed_samples: int, args,
                      combined_parquet_dir: str):
        """Flush collected tables to disk."""
@@ -8,21 +8,15 @@ using the modular pipeline architecture.
 from typing import Any, Dict, List, Optional

 import numpy as np
-import PIL
 import torch
 from PIL import Image

 from fastvideo.v1.dataset.dataloader.schema import pyarrow_schema_i2v
-from fastvideo.v1.distributed import get_local_torch_device
+from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
-from fastvideo.v1.models.vision_utils import (get_default_height_width,
-                                              normalize, numpy_to_pt,
-                                              pil_to_numpy, resize)
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.preprocess.preprocess_pipeline_base import (
    BasePreprocessPipeline)
-from fastvideo.v1.pipelines.stages import ImageEncodingStage, TextEncodingStage


 class PreprocessPipeline_I2V(BasePreprocessPipeline):
@@ -32,73 +26,18 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
        "text_encoder", "tokenizer", "vae", "image_encoder", "image_processor"
    ]

-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        self.add_stage(stage_name="prompt_encoding_stage",
-                       stage=TextEncodingStage(
-                           text_encoders=[self.get_module("text_encoder")],
-                           tokenizers=[self.get_module("tokenizer")],
-                       ))
-
-        self.add_stage(stage_name="image_encoding_stage",
-                       stage=ImageEncodingStage(
-                           image_encoder=self.get_module("image_encoder"),
-                           image_processor=self.get_module("image_processor"),
-                       ))
-
-    def preprocess_image(self, image: PIL.Image.Image, height: int, width: int,
-                         fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
-        assert hasattr(
-            self,
-            "image_encoding_stage"), "Image encoding stage must be created"
-
-        batch = ForwardBatch(
-            data_type="video",
-            pil_image=image,
-        )
-        result_batch = self.image_encoding_stage(batch, fastvideo_args)
-        clip_features = result_batch.image_embeds[0]
-
-        image = self.preprocess(
-            image,
-            vae_scale_factor=self.get_module("vae").spatial_compression_ratio,
-            height=height,
-            width=width)
-
-        return {
-            "clip_feature": clip_features[0],
-            "pil_image": image,
-        }
-
-    def preprocess_video(self, video: list[PIL.Image.Image], height: int,
-                         width: int,
-                         fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
-        return self.preprocess_image(video[0], height, width, fastvideo_args)
-
    def get_schema_fields(self) -> List[str]:
        """Get the schema fields for I2V pipeline."""
        return [f.name for f in pyarrow_schema_i2v]

    def get_extra_features(self, valid_data: Dict[str, Any],
                           fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
-
-        # TODO(will): move these to cpu at some point
-        self.get_module("image_encoder").to(get_local_torch_device())
-        self.get_module("vae").to(get_local_torch_device())
-
-        features = {}
        """Get CLIP features from the first frame of each video."""
        first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(
            0, 2, 3, 1)  # (B, C, T, H, W) -> (B, H, W, C)
-        batch_size, _, num_frames, height, width = valid_data[
-            "pixel_values"].shape
-        latent_height = height // self.get_module(
-            "vae").spatial_compression_ratio
-        latent_width = width // self.get_module("vae").spatial_compression_ratio

        processed_images = []
-        # Frame has values between -1 and 1
        for frame in first_frame:
-            frame = (frame + 1) * 127.5
            frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
            processed_img = self.get_module("image_processor")(
                images=frame_pil, return_tensors="pt")
@@ -107,91 +46,29 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
        # Get CLIP features
        pixel_values = torch.cat(
            [img['pixel_values'] for img in processed_images],
-            dim=0).to(get_local_torch_device())
+            dim=0).to(get_torch_device())
        with torch.no_grad():
            image_inputs = {'pixel_values': pixel_values}
            with set_forward_context(current_timestep=0, attn_metadata=None):
                clip_features = self.get_module("image_encoder")(**image_inputs)
            clip_features = clip_features.last_hidden_state

-        features["clip_feature"] = clip_features
-        """Get VAE features from the first frame of each video"""
-        video_conditions = []
-        for frame in first_frame:
-            processed_img = frame.to(device="cpu", dtype=torch.float32)
-            processed_img = processed_img.unsqueeze(0).permute(0, 3, 1,
-                                                               2).unsqueeze(2)
-            # (B, H, W, C) -> (B, C, 1, H, W)
-            video_condition = torch.cat([
-                processed_img,
-                processed_img.new_zeros(processed_img.shape[0],
-                                        processed_img.shape[1], num_frames - 1,
-                                        height, width)
-            ],
-                                        dim=2)
-            video_condition = video_condition.to(
-                device=get_local_torch_device(), dtype=torch.float32)
-            video_conditions.append(video_condition)
-
-        video_conditions = torch.cat(video_conditions, dim=0)
-
-        with torch.autocast(device_type="cuda",
-                            dtype=torch.float32,
-                            enabled=True):
-            encoder_outputs = self.get_module("vae").encode(video_conditions)
-
-        latent_condition = encoder_outputs.mean
-        if (hasattr(self.get_module("vae"), "shift_factor")
-                and self.get_module("vae").shift_factor is not None):
-            if isinstance(self.get_module("vae").shift_factor, torch.Tensor):
-                latent_condition -= self.get_module("vae").shift_factor.to(
-                    latent_condition.device, latent_condition.dtype)
-            else:
-                latent_condition -= self.get_module("vae").shift_factor
-
-        if isinstance(self.get_module("vae").scaling_factor, torch.Tensor):
-            latent_condition = latent_condition * self.get_module(
-                "vae").scaling_factor.to(latent_condition.device,
-                                         latent_condition.dtype)
-        else:
-            latent_condition = latent_condition * self.get_module(
-                "vae").scaling_factor
-
-        mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height,
-                                   latent_width)
-        mask_lat_size[:, :, list(range(1, num_frames))] = 0
-        first_frame_mask = mask_lat_size[:, :, 0:1]
-        first_frame_mask = torch.repeat_interleave(
-            first_frame_mask,
-            dim=2,
-            repeats=self.get_module("vae").temporal_compression_ratio)
-        mask_lat_size = torch.concat(
-            [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
-        mask_lat_size = mask_lat_size.view(
-            batch_size, -1,
-            self.get_module("vae").temporal_compression_ratio, latent_height,
-            latent_width)
-        mask_lat_size = mask_lat_size.transpose(1, 2)
-        mask_lat_size = mask_lat_size.to(latent_condition.device)
-
-        image_latent = torch.concat([mask_lat_size, latent_condition], dim=1)
-
-        features["first_frame_latent"] = image_latent
-
-        return features
+        return {"clip_feature": clip_features}

    def create_record(
            self,
            video_name: str,
            vae_latent: np.ndarray,
            text_embedding: np.ndarray,
-            valid_data: Dict[str, Any],
+            text_attention_mask: np.ndarray,
+            valid_data: Optional[Dict[str, Any]],
            idx: int,
            extra_features: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Create a record for the Parquet dataset with CLIP features."""
        record = super().create_record(video_name=video_name,
                                       vae_latent=vae_latent,
                                       text_embedding=text_embedding,
+                                       text_attention_mask=text_attention_mask,
                                       valid_data=valid_data,
                                       idx=idx,
                                       extra_features=extra_features)
@@ -210,69 +87,7 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
                "clip_feature_dtype": "",
            })

-        if extra_features and "first_frame_latent" in extra_features:
-            first_frame_latent = extra_features["first_frame_latent"]
-            record.update({
-                "first_frame_latent_bytes":
-                first_frame_latent.tobytes(),
-                "first_frame_latent_shape":
-                list(first_frame_latent.shape),
-                "first_frame_latent_dtype":
-                str(first_frame_latent.dtype),
-            })
-        else:
-            record.update({
-                "first_frame_latent_bytes": b"",
-                "first_frame_latent_shape": [],
-                "first_frame_latent_dtype": "",
-            })
-
-        if extra_features and "pil_image" in extra_features:
-            pil_image = extra_features["pil_image"]
-            record.update({
-                "pil_image_bytes": pil_image.tobytes(),
-                "pil_image_shape": list(pil_image.shape),
-                "pil_image_dtype": str(pil_image.dtype),
-            })
-        else:
-            record.update({
-                "pil_image_bytes": b"",
-                "pil_image_shape": [],
-                "pil_image_dtype": "",
-            })
-
-        return record
-
-    def pil_to_tensor(self, image: PIL.Image.Image) -> torch.Tensor:
-        image = image
-
-        image = np.array(image).astype(np.float32)
-        image = torch.from_numpy(image)
-        return image
-
-    def preprocess(self,
-                   image: PIL.Image.Image,
-                   vae_scale_factor: int,
-                   height: int,
-                   width: int,
-                   resize_mode: str = "default") -> torch.Tensor:
-        image = [image]
-
-        height, width = get_default_height_width(image[0], vae_scale_factor,
-                                                 height, width)
-        image = [
-            resize(i, height, width, resize_mode=resize_mode) for i in image
-        ]
-        image = pil_to_numpy(image)  # to np
-        image = numpy_to_pt(image)  # to pt
-
-        do_normalize = True
-        if image.min() < 0:
-            do_normalize = False
-        if do_normalize:
-            image = normalize(image)
-
-        return image
+        return record  # type: ignore


 EntryClass = PreprocessPipeline_I2V
@@ -44,7 +44,7 @@ if __name__ == "__main__":
    parser.add_argument("--model_path", type=str, default="data/mochi")
    parser.add_argument("--model_type", type=str, default="mochi")
    parser.add_argument("--data_merge_path", type=str, required=True)
-    parser.add_argument("--validation_dataset_file", type=str)
+    parser.add_argument("--validation_prompt_txt", type=str)
    parser.add_argument("--num_frames", type=int, default=163)
    parser.add_argument(
        "--dataloader_num_workers",
@@ -59,6 +59,12 @@ if __name__ == "__main__":
        default=2,
        help="Batch size (per device) for the training dataloader.",
    )
+    parser.add_argument(
+        "--preprocess_text_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
    parser.add_argument("--samples_per_file", type=int, default=64)
    parser.add_argument("--flush_frequency",
                        type=int,
@@ -73,6 +79,7 @@ if __name__ == "__main__":
    parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
    parser.add_argument("--group_frame", action="store_true")  # TODO
    parser.add_argument("--group_resolution", action="store_true")  # TODO
+    parser.add_argument("--dataset", default="t2v")
    parser.add_argument("--preprocess_task", type=str, default="t2v")
    parser.add_argument("--train_fps", type=int, default=30)
    parser.add_argument("--use_image_num", type=int, default=0)
@@ -84,7 +91,7 @@ if __name__ == "__main__":
                        type=str,
                        default="google/t5-v1_1-xxl")
    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-    parser.add_argument("--training_cfg_rate", type=float, default=0.0)
+    parser.add_argument("--cfg", type=float, default=0.0)
    parser.add_argument(
        "--output_dir",
        type=str,
@@ -5,7 +5,7 @@ Decoding stage for diffusion pipelines.

 import torch

-from fastvideo.v1.distributed import get_local_torch_device
+from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.vaes.common import ParallelTiledVAE
@@ -61,7 +61,7 @@ class DecodingStage(PipelineStage):
        Returns:
            The batch with decoded outputs.
        """
-        self.vae = self.vae.to(get_local_torch_device())
+        self.vae = self.vae.to(get_torch_device())

        latents = batch.latents
        # TODO(will): remove this once we add input/output validation for stages
@@ -12,9 +12,8 @@ from tqdm.auto import tqdm

 from fastvideo.v1.attention import get_attn_backend
 from fastvideo.v1.configs.pipelines.base import STA_Mode
-from fastvideo.v1.distributed import (get_local_torch_device,
-                                      get_sp_parallel_rank, get_sp_world_size,
-                                      get_world_group)
+from fastvideo.v1.distributed import (get_sp_parallel_rank, get_sp_world_size,
+                                      get_torch_device, get_world_group)
 from fastvideo.v1.distributed.communication_op import (
    sequence_model_parallel_all_gather)
 from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -193,7 +192,7 @@ class DenoisingStage(PipelineStage):
                        [fastvideo_args.pipeline_config.embedded_cfg_scale] *
                        latent_model_input.shape[0],
                        dtype=torch.float32,
-                        device=get_local_torch_device(),
+                        device=get_torch_device(),
                    ).to(target_dtype) *
                    1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale
                    is not None else None)
@@ -7,13 +7,13 @@ from typing import Optional
 import PIL.Image
 import torch

-from fastvideo.v1.distributed import get_local_torch_device
+from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 from fastvideo.v1.models.vision_utils import (get_default_height_width,
-                                              normalize, numpy_to_pt,
-                                              pil_to_numpy, resize)
+                                              load_image, normalize,
+                                              numpy_to_pt, pil_to_numpy, resize)
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.pipelines.stages.validators import V  # Import validators
@@ -49,37 +49,31 @@ class EncodingStage(PipelineStage):
        Returns:
            The batch with encoded outputs.
        """
-        self.vae = self.vae.to(get_local_torch_device())
+        self.vae = self.vae.to(get_torch_device())

+        image_path = batch.image_path
+        # TODO(will): remove this once we add input/output validation for stages
+        if image_path is None:
+            raise ValueError("Image Path must be provided")
        assert batch.height is not None
        assert batch.width is not None
        latent_height = batch.height // self.vae.spatial_compression_ratio
        latent_width = batch.width // self.vae.spatial_compression_ratio

-        image = batch.preprocessed_image
-        # TODO(will)
-        if image is None:
-            assert batch.pil_image is not None
-            image = batch.pil_image
-            image = self.preprocess(
-                image,
-                vae_scale_factor=self.vae.spatial_compression_ratio,
-                height=batch.height,
-                width=batch.width).to(get_local_torch_device(),
-                                      dtype=torch.float32)
-
-            image = image.unsqueeze(2)
-        else:
-            # assumes image is loaded from parquet file and used for validation
-            image = image.transpose(1, 2)
-        logger.info("image: %s", image.shape)
+        image = load_image(image_path)
+        image = self.preprocess(
+            image,
+            vae_scale_factor=self.vae.spatial_compression_ratio,
+            height=batch.height,
+            width=batch.width).to(get_torch_device(), dtype=torch.float32)
+        image = image.unsqueeze(2)
        video_condition = torch.cat([
            image,
            image.new_zeros(image.shape[0], image.shape[1],
                            batch.num_frames - 1, batch.height, batch.width)
        ],
                                    dim=2)
-        video_condition = video_condition.to(device=get_local_torch_device(),
+        video_condition = video_condition.to(device=get_torch_device(),
                                             dtype=torch.float32)

        # Setup VAE precision
@@ -103,8 +97,7 @@ class EncodingStage(PipelineStage):
        generator = batch.generator
        if generator is None:
            raise ValueError("Generator must be provided")
-        # latent_condition = self.retrieve_latents(encoder_output, generator, sample_mode="argmax")
-        latent_condition = self.retrieve_latents(encoder_output, generator)
+        latent_condition = self.retrieve_latents(encoder_output, generator[0])

        # Apply shifting if needed
        if (hasattr(self.vae, "shift_factor")
@@ -188,7 +181,7 @@ class EncodingStage(PipelineStage):
                     fastvideo_args: FastVideoArgs) -> VerificationResult:
        """Verify encoding stage inputs."""
        result = VerificationResult()
-        # result.add_check("pil_image", batch.pil_image)
+        result.add_check("image_path", batch.image_path, V.string_not_empty)
        result.add_check("height", batch.height, V.positive_int)
        result.add_check("width", batch.width, V.positive_int)
        result.add_check("generator", batch.generator,
@@ -7,10 +7,11 @@ This module contains implementations of image encoding stages for diffusion pipe

 import torch

-from fastvideo.v1.distributed import get_local_torch_device
+from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.vision_utils import load_image
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.pipelines.stages.validators import StageValidators as V
@@ -55,12 +56,12 @@ class ImageEncodingStage(PipelineStage):
            The batch with encoded prompt embeddings.
        """
        if fastvideo_args.use_cpu_offload:
-            self.image_encoder = self.image_encoder.to(get_local_torch_device())
+            self.image_encoder = self.image_encoder.to(get_torch_device())

-        image = batch.pil_image
+        image = load_image(batch.image_path)

        image_inputs = self.image_processor(
-            images=image, return_tensors="pt").to(get_local_torch_device())
+            images=image, return_tensors="pt").to(get_torch_device())
        with set_forward_context(current_timestep=0, attn_metadata=None):
            outputs = self.image_encoder(**image_inputs)
            image_embeds = outputs.last_hidden_state
@@ -77,7 +78,7 @@ class ImageEncodingStage(PipelineStage):
                     fastvideo_args: FastVideoArgs) -> VerificationResult:
        """Verify image encoding stage inputs."""
        result = VerificationResult()
-        result.add_check("pil_image", batch.pil_image, V.not_none)
+        result.add_check("image_path", batch.image_path, V.string_not_empty)
        result.add_check("image_embeds", batch.image_embeds, V.is_list)
        return result

@@ -7,7 +7,6 @@ import torch

 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.vision_utils import load_image
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.pipelines.stages.validators import (StageValidators,
@@ -92,11 +91,6 @@ class InputValidationStage(PipelineStage):
                f"Guidance scale must be positive, but got {batch.guidance_scale}"
            )

-        # for i2v, get image from image_path
-        if batch.image_path is not None:
-            image = load_image(batch.image_path)
-            batch.pil_image = image
-
        return batch

    def verify_input(self, batch: ForwardBatch,
@@ -5,7 +5,7 @@ Latent preparation stage for diffusion pipelines.

 from diffusers.utils.torch_utils import randn_tensor

-from fastvideo.v1.distributed import get_local_torch_device
+from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
@@ -62,7 +62,7 @@ class LatentPreparationStage(PipelineStage):

        # Get required parameters
        dtype = batch.prompt_embeds[0].dtype
-        device = get_local_torch_device()
+        device = get_torch_device()
        generator = batch.generator
        latents = batch.latents
        num_frames = latent_num_frames if latent_num_frames is not None else batch.num_frames
@@ -5,7 +5,9 @@ Prompt encoding stages for diffusion pipelines.
 This module contains implementations of prompt encoding stages for diffusion pipelines.
 """

-from fastvideo.v1.distributed import get_local_torch_device
+import torch
+
+from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
@@ -60,6 +62,8 @@ class TextEncodingStage(PipelineStage):
                fastvideo_args.pipeline_config.text_encoder_configs,
                fastvideo_args.pipeline_config.preprocess_text_funcs,
                fastvideo_args.pipeline_config.postprocess_text_funcs):
+            if fastvideo_args.use_cpu_offload:
+                text_encoder = text_encoder.to(get_torch_device())

            assert isinstance(batch.prompt, (str, list))
            if isinstance(batch.prompt, str):
@@ -67,9 +71,8 @@ class TextEncodingStage(PipelineStage):
            texts = []
            for prompt_str in batch.prompt:
                texts.append(preprocess_func(prompt_str))
-            text_inputs = tokenizer(texts,
-                                    **encoder_config.tokenizer_kwargs).to(
-                                        get_local_torch_device())
+            text_inputs = tokenizer(
+                texts, **encoder_config.tokenizer_kwargs).to(get_torch_device())
            input_ids = text_inputs["input_ids"]
            attention_mask = text_inputs["attention_mask"]
            with set_forward_context(current_timestep=0, attn_metadata=None):
@@ -88,8 +91,8 @@ class TextEncodingStage(PipelineStage):
                assert isinstance(batch.negative_prompt, str)
                negative_text = preprocess_func(batch.negative_prompt)
                negative_text_inputs = tokenizer(
-                    negative_text, **encoder_config.tokenizer_kwargs).to(
-                        get_local_torch_device())
+                    negative_text,
+                    **encoder_config.tokenizer_kwargs).to(get_torch_device())
                negative_input_ids = negative_text_inputs["input_ids"]
                negative_attention_mask = negative_text_inputs["attention_mask"]
                with set_forward_context(current_timestep=0,
@@ -107,8 +110,9 @@ class TextEncodingStage(PipelineStage):
                    batch.negative_attention_mask.append(
                        negative_attention_mask)

-            if fastvideo_args.text_encoder_offload:
+            if fastvideo_args.use_cpu_offload:
                text_encoder.to('cpu')
+                torch.cuda.empty_cache()

        return batch

@@ -7,7 +7,7 @@ This module contains implementations of timestep preparation stages for diffusio

 import inspect

-from fastvideo.v1.distributed import get_local_torch_device
+from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
@@ -45,7 +45,7 @@ class TimestepPreparationStage(PipelineStage):
            The batch with prepared timesteps.
        """
        scheduler = self.scheduler
-        device = get_local_torch_device()
+        device = get_torch_device()
        num_inference_steps = batch.num_inference_steps
        timesteps = batch.timesteps
        sigmas = batch.sigmas
@@ -14,7 +14,7 @@ from typing import Any, Dict
 import torch
 from huggingface_hub import hf_hub_download

-from fastvideo.v1.distributed import get_local_torch_device
+from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.encoders.bert import HunyuanClip  # type: ignore
@@ -78,7 +78,7 @@ class StepVideoPipeline(LoRAPipeline, ComposedPipelineBase):
        """
        Initialize the pipeline.
        """
-        target_device = get_local_torch_device()
+        target_device = get_torch_device()
        llm_dir = os.path.join(self.model_path, "step_llm")
        clip_dir = os.path.join(self.model_path, "hunyuan_clip")
        text_enc = self.build_llm(llm_dir, target_device)
@@ -76,36 +76,4 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
                       stage=DecodingStage(vae=self.get_module("vae")))


-class WanImageToVideoValidationPipeline(ComposedPipelineBase):
-    """
-    I2V Validation pipeline for Wan2.1, assumes that the input are preprocess latents.
-    """
-    _required_config_modules = ["vae", "scheduler", "transformer"]
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.modules["scheduler"] = FlowUniPCMultistepScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift)
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self.add_stage(stage_name="timestep_preparation_stage",
-                       stage=TimestepPreparationStage(
-                           scheduler=self.get_module("scheduler")))
-
-        self.add_stage(stage_name="latent_preparation_stage",
-                       stage=LatentPreparationStage(
-                           scheduler=self.get_module("scheduler"),
-                           transformer=self.get_module("transformer")))
-
-        self.add_stage(stage_name="image_latent_preparation_stage",
-                       stage=EncodingStage(vae=self.get_module("vae")))
-
-        self.add_stage(stage_name="denoising_stage",
-                       stage=DenoisingStage(
-                           transformer=self.get_module("transformer"),
-                           scheduler=self.get_module("scheduler")))
-
-        self.add_stage(stage_name="decoding_stage",
-                       stage=DecodingStage(vae=self.get_module("vae")))
-
-
 EntryClass = WanImageToVideoPipeline
@@ -6,7 +6,7 @@ import numpy as np
 import pytest
 import torch
 from transformers import AutoConfig
-import gc
+
 from fastvideo.models.hunyuan.text_encoder import (load_text_encoder,
                                                   load_tokenizer)
 # from fastvideo.v1.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer
@@ -16,8 +16,6 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import maybe_download_model
 from fastvideo.v1.configs.models.encoders import CLIPTextConfig
-from torch.distributed.tensor import DTensor
-from torch.testing import assert_close

 logger = init_logger(__name__)

@@ -68,6 +66,7 @@ def test_clip_encoder():
    # Load the HuggingFace implementation directly
    # model2 = CLIPTextModel(hf_config)
    # model2 = model2.to(torch.float16)
+    model2 = model2.to(device)
    model2.eval()

    # Sanity check weights between the two models
@@ -79,20 +78,19 @@ def test_clip_encoder():
    logger.info("Model1 has %d parameters", len(params1))
    logger.info("Model2 has %d parameters", len(params2))

-    for name1, param1 in sorted(params1.items()):
-        name2 = name1
-        skip = False
-        for param_name, weight_name, shard_id in model2.config.arch_config.stacked_params_mapping:
-            if weight_name not in name1:
-                skip = True
-        # stacked params are more troublesome
-        if skip:
-            continue
-        param2 = params2[name2]
-        param2 = param2.to_local().to(device) if isinstance(param2, DTensor) else param2.to(device)
-        assert_close(param1, param2, atol=1e-4, rtol=1e-4)
-    gc.collect()
-    torch.cuda.empty_cache()
+    # Compare a few key parameters
+
+    # weight_diffs = []
+    # for (name1, param1), (name2, param2) in zip(
+    #     sorted(params1.items()), sorted(params2.items())
+    # ):
+    #     # if len(weight_diffs) < 5:  # Just check a few parameters
+    #     max_diff = torch.max(torch.abs(param1 - param2)).item()
+    #     mean_diff = torch.mean(torch.abs(param1 - param2)).item()
+    #     weight_diffs.append((name1, name2, max_diff, mean_diff))
+    #     logger.info(f"Parameter: {name1} vs {name2}")
+    #     logger.info(f"  Max diff: {max_diff}, Mean diff: {mean_diff}")
+
    # Load tokenizer
    tokenizer, _ = load_tokenizer(tokenizer_type="clipL",
                                  tokenizer_path=args.model_path,
@@ -170,5 +168,5 @@ def test_clip_encoder():
                f"Pooler outputs differ significantly: mean diff = {mean_diff_pooler.item()}"
            assert max_diff_hidden < 1e-1, \
                f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}"
-            assert max_diff_pooler < 2e-2, \
+            assert max_diff_pooler < 1e-2, \
                f"Pooler outputs differ significantly: max diff = {max_diff_pooler.item()}"
@@ -5,7 +5,7 @@ import numpy as np
 import pytest
 import torch
 from transformers import AutoConfig
-import gc
+
 from fastvideo.models.hunyuan.text_encoder import (load_text_encoder,
                                                   load_tokenizer)
 from fastvideo.v1.configs.pipelines import PipelineConfig
@@ -15,8 +15,7 @@ from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.loader.component_loader import TextEncoderLoader
 from fastvideo.v1.utils import maybe_download_model
 from fastvideo.v1.configs.models.encoders import LlamaConfig
-from torch.distributed.tensor import DTensor
-from torch.testing import assert_close
+
 logger = init_logger(__name__)

 os.environ["MASTER_ADDR"] = "localhost"
@@ -63,6 +62,7 @@ def test_llama_encoder():

    # Convert to float16 and move to device
    # model2 = model2.to(torch.float16)
+    model2 = model2.to(device)
    model2.eval()

    # Sanity check weights between the two models
@@ -77,28 +77,34 @@ def test_llama_encoder():
    # Compare a few key parameters
    weight_diffs = []
    # check if embed_tokens are the same
-    device = model1.embed_tokens.weight.device
+    print(model1.embed_tokens.weight.shape, model2.embed_tokens.weight.shape)
    assert torch.allclose(model1.embed_tokens.weight,
-                          model2.embed_tokens.weight.to_local().to(device) if isinstance(model2.embed_tokens.weight, DTensor) else model2.embed_tokens.weight.to(device))
+                          model2.embed_tokens.weight)
    weights = [
        "layers.{}.input_layernorm.weight",
        "layers.{}.post_attention_layernorm.weight"
    ]
-
-    for name1, param1 in sorted(params1.items()):
-        name2 = name1
-        skip = False
-        for param_name, weight_name, shard_id in model2.config.arch_config.stacked_params_mapping:
-            if weight_name not in name1:
-                skip = True
-        # stacked params are more troublesome
-        if skip:
-            continue
-        param2 = params2[name2]
-        param2 = param2.to_local().to(device) if isinstance(param2, DTensor) else param2.to(device)
-        assert_close(param1, param2, atol=1e-4, rtol=1e-4)
-    gc.collect()
-    torch.cuda.empty_cache()
+    # for (name1, param1), (name2, param2) in zip(
+    #     sorted(params1.items()), sorted(params2.items())
+    # ):
+    for layer_idx in range(hf_config.num_hidden_layers):
+        for w in weights:
+            name1 = w.format(layer_idx)
+            name2 = w.format(layer_idx)
+            p1 = params1[name1]
+            p2 = params2[name2]
+            # print(type(p2))
+            if "gate_up" in name2:
+                # print("skipping gate_up")
+                continue
+            try:
+                # logger.info(f"Parameter: {name1} vs {name2}")
+                max_diff = torch.max(torch.abs(p1 - p2)).item()
+                mean_diff = torch.mean(torch.abs(p1 - p2)).item()
+                weight_diffs.append((name1, name2, max_diff, mean_diff))
+                # logger.info(f"  Max diff: {max_diff}, Mean diff: {mean_diff}")
+            except Exception as e:
+                logger.info("Error comparing %s and %s: %s", name1, name2, e)

    tokenizer, _ = load_tokenizer(tokenizer_type="llm",
                                  tokenizer_path=TOKENIZER_PATH,
@@ -4,8 +4,6 @@ import os
 import numpy as np
 import pytest
 import torch
-from torch.distributed.tensor import DTensor
-from torch.testing import assert_close
 from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel

 from fastvideo.v1.configs.pipelines import PipelineConfig
@@ -43,13 +41,13 @@ def test_t5_encoder():
    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)


-    args = FastVideoArgs(model_path=TEXT_ENCODER_PATH,
-                        pipeline_config=PipelineConfig(text_encoder_configs=(T5Config(),),
-                        text_encoder_precisions=(precision_str,)),
-                        pin_cpu_memory=False)
+    args = FastVideoArgs(model_path=TEXT_ENCODER_PATH, pipeline_config=PipelineConfig(text_encoder_configs=(T5Config(),), text_encoder_precisions=(precision_str,)))
    loader = TextEncoderLoader()
    model2 = loader.load(TEXT_ENCODER_PATH, "", args)
-    model2 = model2.to(precision)
+
+    # Convert to float16 and move to device
+    # model2 = model2.to(precision)
+    model2 = model2.to(device)
    model2.eval()

    # Sanity check weights between the two models
@@ -66,17 +64,23 @@ def test_t5_encoder():
    weights = ["encoder.block.{}.layer.0.layer_norm.weight", "encoder.block.{}.layer.0.SelfAttention.relative_attention_bias.weight", \
               "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\
                "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \
-                "encoder.block.{}.layer.1.layer_norm.weight", "encoder.final_layer_norm.weight"]
-    
+                "encoder.block.{}.layer.1.layer_norm.weight", "encoder.final_layer_norm.weight", "shared.weight"]
    for idx in range(hf_config.num_hidden_layers):
        for w in weights:
            name1 = w.format(idx)
            name2 = w.format(idx)
            p1 = params1[name1]
            p2 = params2[name2]
-            p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1)
-            assert_close(p1, p2, atol=1e-4, rtol=1e-4)
-    
+            assert p1.dtype == p2.dtype
+            try:
+                logger.info("Parameter: %s vs %s", name1, name2)
+                max_diff = torch.max(torch.abs(p1 - p2)).item()
+                mean_diff = torch.mean(torch.abs(p1 - p2)).item()
+                weight_diffs.append((name1, name2, max_diff, mean_diff))
+                logger.info("  Max diff: %s, Mean diff: %s", max_diff,
+                            mean_diff)
+            except Exception as e:
+                logger.info("Error comparing %s and %s: %s", name1, name2, e)

    # Test with some sample prompts
    prompts = [
@@ -4,43 +4,31 @@ app = modal.App()

 import os

-image_version = os.getenv("IMAGE_VERSION")
+image_version = os.getenv("IMAGE_VERSION", "latest")
 image_tag = f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{image_version}"
 print(f"Using image: {image_tag}")

 image = (
    modal.Image.from_registry(image_tag, add_python="3.12")
-    .run_commands("rm -rf /FastVideo")
    .apt_install("cmake", "pkg-config", "build-essential", "curl", "libssl-dev")
    .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable")
    .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
-    .env({
-        "PATH": "/root/.cargo/bin:$PATH",
-        "BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""),
-        "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
-    })
+    .env({"PATH": "/root/.cargo/bin:$PATH"})
+    .run_commands("/bin/bash -c 'source $HOME/.local/bin/env && source /opt/venv/bin/activate && cd /FastVideo && uv pip install -e .[test]'")
 )

-def run_test(pytest_command: str):
-    """Helper function to run a test suite with custom pytest command"""
+@app.function(gpu="L40S:1", image=image, timeout=1800)
+def run_encoder_tests():
+    """Run encoder tests on L40S GPU"""
    import subprocess
    import sys
    import os
    
-    git_repo = os.environ.get("BUILDKITE_REPO")
-    git_commit = os.environ.get("BUILDKITE_COMMIT")
+    os.chdir("/FastVideo")
    
-    print(f"Cloning repository: {git_repo}")
-    print(f"Checking out commit: {git_commit}")
-    
-    command = f"""
-    source $HOME/.local/bin/env &&
-    source /opt/venv/bin/activate &&
-    git clone {git_repo} /FastVideo &&
-    cd /FastVideo &&
-    git checkout {git_commit} &&
-    uv pip install -e .[test] &&
-    {pytest_command}
+    command = """
+    source /opt/venv/bin/activate && 
+    pytest ./fastvideo/v1/tests/encoders -s
    """
    
    result = subprocess.run([
@@ -49,38 +37,62 @@ def run_test(pytest_command: str):
    
    sys.exit(result.returncode)

-@app.function(gpu="L40S:1", image=image, timeout=1800)
-def run_encoder_tests():
-    run_test("pytest ./fastvideo/v1/tests/encoders -vs")
-
@app.function(gpu="L40S:1", image=image, timeout=1800)
 def run_vae_tests():
-    run_test("pytest ./fastvideo/v1/tests/vaes -vs")
+    """Run VAE tests on L40S GPU"""
+    import subprocess
+    import sys
+    import os
+    
+    os.chdir("/FastVideo")
+    
+    command = """
+    source /opt/venv/bin/activate && 
+    pytest ./fastvideo/v1/tests/vaes -s
+    """
+    
+    result = subprocess.run([
+        "/bin/bash", "-c", command
+    ], stdout=sys.stdout, stderr=sys.stderr, check=False)
+    
+    sys.exit(result.returncode)

@app.function(gpu="L40S:1", image=image, timeout=1800)
 def run_transformer_tests():
-    run_test("pytest ./fastvideo/v1/tests/transformers -vs")
+    """Run transformer tests on L40S GPU"""
+    import subprocess
+    import sys
+    import os
+    
+    os.chdir("/FastVideo")
+    
+    command = """
+    source /opt/venv/bin/activate && 
+    pytest ./fastvideo/v1/tests/transformers -s
+    """
+    
+    result = subprocess.run([
+        "/bin/bash", "-c", command
+    ], stdout=sys.stdout, stderr=sys.stderr, check=False)
+    
+    sys.exit(result.returncode)

@app.function(gpu="L40S:2", image=image, timeout=3600)
 def run_ssim_tests():
-    run_test("pytest ./fastvideo/v1/tests/ssim -vs")
-
-@app.function(gpu="L40S:4", image=image, timeout=1800, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
-def run_training_tests():
-    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/Vanilla -srP")
-
-@app.function(gpu="H100:1", image=image, timeout=1800, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
-def run_training_tests_VSA():
-    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/VSA -srP")
-
-@app.function(gpu="H100:1", image=image, timeout=1800)
-def run_inference_tests_STA():
-    run_test("pytest ./fastvideo/v1/tests/inference/STA -srP")
-
-@app.function(gpu="H100:1", image=image, timeout=1800)
-def run_precision_tests_STA():
-    run_test("python csrc/attn/tests/test_sta.py")
-
-@app.function(gpu="H100:1", image=image, timeout=1800)
-def run_precision_tests_VSA():
-    run_test("python csrc/attn/tests/test_block_sparse.py")
+    """Run SSIM tests on 2x L40S GPUs"""
+    import subprocess
+    import sys
+    import os
+    
+    os.chdir("/FastVideo")
+    
+    command = """
+    source /opt/venv/bin/activate && 
+    pytest ./fastvideo/v1/tests/ssim -vs
+    """
+    
+    result = subprocess.run([
+        "/bin/bash", "-c", command
+    ], stdout=sys.stdout, stderr=sys.stderr, check=False)
+    
+    sys.exit(result.returncode)
@@ -1 +0,0 @@
-{"step_time":2.245914653001819,"_wandb":{"runtime":1434},"learning_rate":1e-05,"grad_norm":0.57421875,"avg_step_time":1.1814782944297622,"train_loss":0.07932619750499725,"vsa_sparsity":0,"_timestamp":1.750578625921253e+09,"validation_videos_40_steps":{"count":1,"videos":[{"size":420969,"path":"media/videos/validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4","_type":"video-file","sha256":"581ff5eae2909d3a7b362dcb24d060c006c09e4d4deb44b82f4aa697f6789ba7"}],"captions":false,"_type":"videos"},"_runtime":1434.62395329,"_step":901}
@@ -1,177 +0,0 @@
-import os
-from pathlib import Path
-from huggingface_hub import snapshot_download
-import shutil
-import subprocess
-import sys
-from fastvideo.v1.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
-
-# Import the training pipeline
-sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-
-NUM_NODES = "1"
-MODEL_PATH = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-
-# preprocessing
-DATA_DIR = "data"
-LOCAL_RAW_DATA_DIR = Path(os.path.join(DATA_DIR, "cats"))
-NUM_GPUS_PER_NODE_PREPROCESSING = "1"
-PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/v1/pipelines/preprocess/v1_preprocess.py"
-
-LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "cats_preprocessed_data_i2v"))
-
-
-# training
-NUM_GPUS_PER_NODE_TRAINING = "8"
-TRAINING_ENTRY_FILE_PATH = "fastvideo/v1/training/wan_i2v_training_pipeline.py"
-LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_parquet_dataset")
-LOCAL_VALIDATION_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "validation_parquet_dataset")
-LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs"))
-
-def download_data():
-    # create the data dir if it doesn't exist
-    data_dir = Path(DATA_DIR)
-    # if data_dir.exists():
-    #     print(f"Removing existing data directory at {data_dir}")
-    #     shutil.rmtree(data_dir)
-
-    print(f"Creating data directory at {data_dir}")
-    os.makedirs(data_dir)
-
-    print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
-    try:
-        # result = snapshot_download(
-        #     repo_id="wlsaidhi/cats-overfit-merged",
-        #     local_dir=str(LOCAL_RAW_DATA_DIR),
-        #     repo_type="dataset",
-        #     resume_download=True,
-        #     token=os.environ.get("HF_TOKEN"),  # In case authentication is needed
-        # )
-        print(f"Download completed successfully. Files downloaded to: {result}")
-        
-        # Verify the download
-        if not LOCAL_RAW_DATA_DIR.exists():
-            raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist")
-            
-        # List downloaded files
-        print("Downloaded files:")
-        for file in LOCAL_RAW_DATA_DIR.rglob("*"):
-            if file.is_file():
-                print(f"  - {file.relative_to(LOCAL_RAW_DATA_DIR)}")
-                
-    except Exception as e:
-        print(f"Error during download: {str(e)}")
-        raise
-
-
-def run_preprocessing():
-    # Run torchrun command
-    cmd = [
-        "torchrun",
-        "--nnodes", NUM_NODES,
-        "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING,
-        PREPROCESSING_ENTRY_FILE_PATH,
-        "--model_path", MODEL_PATH,
-        "--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge_1_sample.txt"),
-        "--preprocess_video_batch_size", "1",
-        "--max_height", "480",
-        "--max_width", "832",
-        "--num_frames", "77",
-        "--dataloader_num_workers", "0",
-        "--output_dir", LOCAL_PREPROCESSED_DATA_DIR,
-        "--train_fps", "16",
-        "--validation_dataset_file", os.path.join(LOCAL_RAW_DATA_DIR, "validation_i2v_prompt_1_sample.json"),
-        "--samples_per_file", "1",
-        "--flush_frequency", "1",
-        "--video_length_tolerance_range", "5",
-        "--preprocess_task", "i2v",
-    ]
-
-    process = subprocess.run(cmd, check=True)
-
-
-def run_training():
-    cmd = [
-        "torchrun",
-        "--nnodes", NUM_NODES,
-        "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING,
-        TRAINING_ENTRY_FILE_PATH,
-        "--model_path", MODEL_PATH,
-        "--inference_mode", "False",
-        "--pretrained_model_name_or_path", MODEL_PATH,
-        "--data_path", LOCAL_TRAINING_DATA_DIR,
-        "--validation_preprocessed_path", LOCAL_VALIDATION_DATA_DIR,
-        "--train_batch_size", "1",
-        "--num_latent_t", "8",
-        "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
-        "--sp_size", NUM_GPUS_PER_NODE_TRAINING,
-        "--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,
-        "--train_sp_batch_size", "1",
-        "--dataloader_num_workers", "10",
-        "--gradient_accumulation_steps", "1",
-        "--max_train_steps", "901",
-        "--learning_rate", "1e-5",
-        "--mixed_precision", "bf16",
-        "--checkpointing_steps", "6000",
-        "--validation_steps", "100",
-        "--validation_sampling_steps", "40",
-        "--log_validation",
-        "--checkpoints_total_limit", "3",
-        "--allow_tf32",
-        "--ema_start_step", "0",
-        "--training_cfg_rate", "0.1",
-        "--output_dir", LOCAL_OUTPUT_DIR,
-        "--tracker_project_name", "wan_i2v_finetune_overfit_ci",
-        "--num_height", "480",
-        "--num_width", "832",
-        "--num_frames", "81",
-        "--validation_guidance_scale", "1.0",
-        "--num_euler_timesteps", "50",
-        "--multi_phased_distill_schedule", "4000-1",
-        "--weight_decay", "0.01",
-        "--not_apply_cfg_solver",
-        "--dit_precision", "fp32",
-        "--max_grad_norm", "1.0",
-    ]
-
-    print(f"Running training with command: {cmd}")
-    process = subprocess.run(cmd, check=True)
-
-
-def test_e2e_overfit_single_sample():
-    os.environ["WANDB_MODE"] = "online"
-
-    # download_data()
-    run_preprocessing()
-    run_training()
-
-    reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4")
-    print(f"reference_video_file: {reference_video_file}")
-    final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4")
-    print(f"final_validation_video_file: {final_validation_video_file}")
-
-
-    # Ensure both files exist
-    assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}"
-    assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}"
-
-    # Compute SSIM
-    mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
-        reference_video_file,
-        final_validation_video_file,
-        use_ms_ssim=True  # Using MS-SSIM for better quality assessment
-    )
-
-    print("\n===== SSIM Results for Step 900 Validation =====")
-    print(f"Mean MS-SSIM: {mean_ssim:.4f}")
-    print(f"Min MS-SSIM: {min_ssim:.4f}")
-    print(f"Max MS-SSIM: {max_ssim:.4f}")
-
-    assert max_ssim > 0.5, f"Max SSIM is below 0.5: {max_ssim}"
-
-
-if __name__ == "__main__":
-    test_e2e_overfit_single_sample()
@@ -31,9 +31,12 @@ LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs"))
 def download_data():
    # create the data dir if it doesn't exist
    data_dir = Path(DATA_DIR)
+    if data_dir.exists():
+        print(f"Removing existing data directory at {data_dir}")
+        shutil.rmtree(data_dir)

    print(f"Creating data directory at {data_dir}")
-    os.makedirs(data_dir, exist_ok=True)
+    os.makedirs(data_dir)

    print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
    try:
@@ -77,7 +80,7 @@ def run_preprocessing():
        "--dataloader_num_workers", "0",
        "--output_dir", LOCAL_PREPROCESSED_DATA_DIR,
        "--train_fps", "16",
-        "--validation_dataset_file", os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt_1_sample.json"),
+        "--validation_prompt_txt", os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt_1_sample.txt"),
        "--samples_per_file", "1",
        "--flush_frequency", "1",
        "--video_length_tolerance_range", "5",
@@ -97,7 +100,7 @@ def run_training():
        "--inference_mode", "False",
        "--pretrained_model_name_or_path", MODEL_PATH,
        "--data_path", LOCAL_TRAINING_DATA_DIR,
-        "--validation_preprocessed_path", LOCAL_VALIDATION_DATA_DIR,
+        "--validation_prompt_dir", LOCAL_VALIDATION_DATA_DIR,
        "--train_batch_size", "1",
        "--num_latent_t", "8",
        "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
@@ -119,7 +122,7 @@ def run_training():
        "--checkpoints_total_limit", "3",
        "--allow_tf32",
        "--ema_start_step", "0",
-        "--training_cfg_rate", "0.0",
+        "--cfg", "0.0",
        "--output_dir", LOCAL_OUTPUT_DIR,
        "--tracker_project_name", "wan_finetune_overfit_ci",
        "--num_height", "480",
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Edenzzzz	6f88ec45f1	fix	2025-06-20 22:47:06 -07:00
Edenzzzz	716f78ee5a	trigger ci	2025-06-20 22:46:56 -07:00
				`@@ -1 +0,0 @@`
				{"step_time":2.245914653001819,"_wandb":{"runtime":1434},"learning_rate":1e-05,"grad_norm":0.57421875,"avg_step_time":1.1814782944297622,"train_loss":0.07932619750499725,"vsa_sparsity":0,"_timestamp":1.750578625921253e+09,"validation_videos_40_steps":{"count":1,"videos":[{"size":420969,"path":"media/videos/validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4","_type":"video-file","sha256":"581ff5eae2909d3a7b362dcb24d060c006c09e4d4deb44b82f4aa697f6789ba7"}],"captions":false,"_type":"videos"},"_runtime":1434.62395329,"_step":901}