fix

Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: Wenxuan Tan <wenxuan.tan@wisc.edu>

.buildkite/pipeline.yml

@@ -0,0 +1,66 @@
+env:
+  IMAGE_VERSION: "py3.12-latest"
+
+steps:
+  - 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/loaders/**"
+                - "fastvideo/v1/tests/encoders/**"
+              config:
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                label: "Encoder Tests"
+                env:
+                  - BUILDKITE_CLEAN_CHECKOUT=true
+                  - TEST_TYPE=encoder
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/models/vaes/**"
+                - "fastvideo/v1/models/loaders/**"
+                - "fastvideo/v1/tests/vaes/**"
+              config:
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                label: "VAE Tests"
+                env:
+                  - BUILDKITE_CLEAN_CHECKOUT=true
+                  - TEST_TYPE=vae
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/models/dits/**"
+                - "fastvideo/v1/models/loaders/**"
+                - "fastvideo/v1/tests/transformers/**"
+                - "fastvideo/v1/layers/**"
+                - "fastvideo/v1/attention/**"
+              config:
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                label: "Transformer Tests"
+                env:
+                  - BUILDKITE_CLEAN_CHECKOUT=true
+                  - TEST_TYPE=transformer
+                agents:
+                  queue: "default"
+            - path: "fastvideo/v1/**/*.py"
+              config:
+                command: "timeout 60m .buildkite/scripts/pr_test.sh"
+                label: "SSIM Tests"
+                env:
+                  - BUILDKITE_CLEAN_CHECKOUT=true
+                  - TEST_TYPE=ssim
+                agents:
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -0,0 +1,91 @@
+#!/bin/bash
+set -uo pipefail
+
+log() {
+  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
+}
+
+log "=== Starting Modal test execution ==="
+
+# Change to the project directory
+cd "$(dirname "$0")/../.."
+PROJECT_ROOT=$(pwd)
+log "Project root: $PROJECT_ROOT"
+
+# Install Modal if not available
+if ! python3 -m modal --version &> /dev/null; then
+    log "Modal not found, installing..."
+    python3 -m pip install modal
+    
+    # Verify installation
+    if ! python3 -m modal --version &> /dev/null; then
+        log "Error: Failed to install modal. Please install it manually."
+        exit 1
+    fi
+fi
+
+log "modal version: $(python3 -m modal --version)"
+
+# Set up Modal authentication using Buildkite secrets
+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)
+
+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
+    if [ $? -eq 0 ]; then
+        log "Modal authentication successful"
+    else
+        log "Error: Failed to set Modal credentials"
+        exit 1
+    fi
+else
+    log "Error: Could not retrieve Modal credentials from Buildkite secrets."
+    log "Please ensure 'modal_token_id' and 'modal_token_secret' secrets are set in Buildkite."
+    exit 1
+fi
+
+MODAL_TEST_FILE="fastvideo/v1/tests/modal/pr_test.py"
+
+if [ -z "${TEST_TYPE:-}" ]; then
+    log "Error: TEST_TYPE environment variable is not set"
+    exit 1
+fi
+log "Test type: $TEST_TYPE"
+
+case "$TEST_TYPE" in
+    "encoder")
+        log "Running encoder tests..."
+        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
+        ;;
+    "vae")
+        log "Running VAE tests..."
+        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
+        ;;
+    "transformer")
+        log "Running transformer tests..."
+        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
+        ;;
+    "ssim")
+        log "Running SSIM tests..."
+        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_ssim_tests"
+        ;;
+    *)
+        log "Error: Unknown test type: $TEST_TYPE"
+        exit 1
+        ;;
+esac
+
+log "Executing: $MODAL_COMMAND"
+eval "$MODAL_COMMAND"
+TEST_EXIT_CODE=$?
+
+if [ $TEST_EXIT_CODE -eq 0 ]; then
+    log "Modal test completed successfully"
+else
+    log "Error: Modal test failed with exit code: $TEST_EXIT_CODE"
+fi
+
+log "=== Test execution completed with exit code: $TEST_EXIT_CODE ==="
+exit $TEST_EXIT_CODE

.github/scripts/runpod_api.py

@@ -160,8 +160,7 @@ def execute_command(pod_id):
     setup_steps = [
         "tar -xzf /tmp/repo.tar.gz --no-same-owner -C /workspace/",
         f"cd /workspace/{repo_name}",
-        "source /opt/conda/etc/profile.d/conda.sh",
-        "conda activate fastvideo-dev",
+        "source $HOME/.local/bin/env && source /opt/venv/bin/activate",
         args.test_command
     ]
     

.github/workflows/pr-test.yml

@@ -12,12 +12,14 @@ on:
     paths:
       - "fastvideo/**/*.py"
       - ".github/workflows/pr-test.yml"
+      - "pyproject.toml"
+      - "docker/Dockerfile.python3.12"
   workflow_dispatch:
     inputs:
       custom_image:
-        description: "Custom image from this repository (default: fastvideo-dev:latest)"
+        description: "Custom image from this repository (default: fastvideo-dev:py3.12-latest)"
         required: false
-        default: "fastvideo-dev:latest"
+        default: "fastvideo-dev:py3.12-latest"
         type: string
       run_encoder_test:
         description: "Run encoder-test"
@@ -44,6 +46,16 @@ on:
         required: false
         default: false
         type: boolean
+      run_training_test_VSA:
+        description: "Run training-test-VSA"
+        required: false
+        default: false
+        type: boolean
+      run_inference_test_STA:
+        description: "Run inference-test-STA"
+        required: false
+        default: false
+        type: boolean
       run_nightly_test:
         description: "Run nightly-test"
         required: false
@@ -70,6 +82,8 @@ jobs:
       vae-test: ${{ steps.filter.outputs.vae-test }}
       transformer-test: ${{ steps.filter.outputs.transformer-test }}
       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 }}
     steps:
       - uses: actions/checkout@v4
       - uses: dorny/paths-filter@v3
@@ -80,18 +94,34 @@ jobs:
               - 'fastvideo/v1/models/encoders/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/encoders/**'
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
             vae-test:
               - 'fastvideo/v1/models/vaes/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/vaes/**'
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
             transformer-test:
               - 'fastvideo/v1/models/dits/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/transformers/**'
               - 'fastvideo/v1/layers/**'
               - 'fastvideo/v1/attention/**'
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
             training-test:
               - 'fastvideo/v1/**'
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
+            training-test-VSA:
+              - 'fastvideo/v1/**'
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
+            inference-test-STA:
+              - 'fastvideo/v1/**'
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
 
   encoder-test:
     needs: change-filter
@@ -104,8 +134,8 @@ jobs:
       gpu_type: "NVIDIA A40"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
+      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:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -122,8 +152,8 @@ jobs:
       gpu_type: "NVIDIA A40"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
+      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:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -140,8 +170,8 @@ jobs:
       gpu_type: "NVIDIA L40S"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
+      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:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -168,7 +198,7 @@ jobs:
       volume_size: 200
       disk_size: 200
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
       timeout_minutes: 60
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -186,8 +216,48 @@ jobs:
       gpu_count: 4
       volume_size: 100
       disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/training -srP"
+      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:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
+      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+
+  training-test-VSA:
+    needs: change-filter
+    if: >-
+      (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:
+      job_id: "training-test-VSA"
+      gpu_type: "NVIDIA H100 NVL"
+      gpu_count: 1
+      volume_size: 100
+      disk_size: 100
+      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:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
+      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+
+  inference-test-STA:
+    needs: change-filter
+    if: >-
+      (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:
+      job_id: "inference-test-STA"
+      gpu_type: "NVIDIA H100 NVL"
+      gpu_count: 1
+      volume_size: 100
+      disk_size: 100
+      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 }}
@@ -203,8 +273,8 @@ jobs:
       gpu_count: 4
       volume_size: 100
       disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs"
+      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:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}

README.md

@@ -91,7 +91,7 @@ For a more detailed guide, please see our [inference quick start](https://hao-ai
 
 ## Distillation and Finetuning
 - [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/training/distillation.html)
-- [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetuning.html)
+- [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html)
 
 ## 📑 Development Plan
 
@@ -111,7 +111,7 @@ For a more detailed guide, please see our [inference quick start](https://hao-ai
 
 ## 🤝 Contributing
 
-We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/developer_guide/overview.html)
+We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview.html)
 
 ## Acknowledgement
 We learned and reused code from the following projects:

docker/Dockerfile.python3.12

@@ -1,7 +1,9 @@
-FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
+SHELL ["/bin/bash", "-c"]
+
 WORKDIR /FastVideo
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -9,17 +11,25 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     git \
     ca-certificates \
     openssh-server \
+    zsh \
+    vim \
+    curl \
+    gcc-11 \
+    g++-11 \
+    clang-11 \
     && rm -rf /var/lib/apt/lists/*
 
-RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
-    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
-    rm Miniconda3-latest-Linux-x86_64.sh
+# Set up C++20 compilers for ThunderKittens
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
 
-ENV PATH=/opt/conda/bin:$PATH
+# Set CUDA environment variables
+ENV CUDA_HOME=/usr/local/cuda-12.8
+ENV PATH=${CUDA_HOME}/bin:${PATH}
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 
-RUN conda create --name fastvideo-dev python=3.12.9 -y
-
-SHELL ["/bin/bash", "-c"]
+# Install uv and source its environment
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    echo 'source $HOME/.local/bin/env' >> /root/.bashrc
 
 # Copy just the pyproject.toml first to leverage Docker cache
 COPY pyproject.toml ./
@@ -27,22 +37,36 @@ COPY pyproject.toml ./
 # Create a dummy README to satisfy the installation
 RUN echo "# Placeholder" > README.md
 
-RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
-    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
-    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
-    conda clean -afy
+# Create and activate virtual environment with specific Python version and seed
+RUN source $HOME/.local/bin/env && \
+    uv venv --python 3.12 --seed /opt/venv && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir --upgrade pip && \
+    uv pip install --no-cache-dir .[dev] && \
+    uv pip install --no-cache-dir flash-attn==2.8.0.post2 --no-build-isolation
 
 COPY . .
 
-RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
+# Install dependencies using uv and set up shell configuration
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir -e .[dev] && \
+    git config --unset-all http.https://github.com/.extraheader || true && \
+    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
+    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Remove authentication headers
-RUN git config --unset-all http.https://github.com/.extraheader || true
+# Install STA (Sliding Tile Attention)
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn && \
+    git submodule update --init --recursive && \
+    python setup_sta.py install
 
-# Set up automatic conda environment activation for all shells
-RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
-    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
-    # Ensure .bashrc is sourced for SSH login shells
-    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+# Install VSA 
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn && \
+    git submodule update --init --recursive && \
+    python setup_vsa.py install
 
 EXPOSE 22

docs/source/sliding_tile_attention/demo.md

@@ -1,7 +1,7 @@
 (sta-demo)=
 
 # 🔍 Demo
-There is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
+This is is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
 
 <div style="text-align: center;">
   <video controls width="800">
@@ -9,3 +9,9 @@ There is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
     Your browser does not support the video tag.
   </video>
 </div>
+
+You can run STA using the following command:
+
+```bash
+bash scripts/inference/v1_inference_wan_STA.sh
+```

docs/source/training/finetune.md

@@ -16,6 +16,13 @@ bash scripts/finetune/finetune_mochi.sh # for mochi
 ```
 
 **Note that for finetuning, we did not tune the hyperparameters in the provided script.**
+## ⚡ Finetune with VSA
+Follow [data_preprocess.md](#v0-data-preprocess) to get parquet files for preproccessed latent, and then run:
+
+```bash
+bash scripts/finetune/finetune_v1_VSA.sh
+```
+
 ## ⚡ Lora Finetune
 
 Hunyuan supports Lora fine-tuning of videos up to 720p. Demos and prompts of Black-Myth-Wukong can be found in [here](https://huggingface.co/FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight). You can download the Lora weight through:

examples/inference/sta_mask_search/inference_wan_sta.sh

@@ -5,7 +5,7 @@ export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
 export MODEL_BASE=Wan-AI/Wan2.1-T2V-14B-Diffusers
 
 base_port=29503
-num_gpu=$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l)
+num_gpu=1
 gpu_ids=$(seq 0 $((num_gpu-1)))
 skip_time_steps=12
 
@@ -14,7 +14,7 @@ STA_mode="STA_searching"
 for i in $gpu_ids; do
     port=$((base_port+i))
     CUDA_VISIBLE_DEVICES=$i MASTER_PORT=$port python examples/inference/sta_mask_search/wan_example.py \
-        --prompt_path ./assets/prompt_extend_${i}.txt \
+        --prompt_path ./assets/prompt_${i}.txt \
         --output_path $output_path \
         --STA_mode $STA_mode &
     sleep 1
@@ -27,7 +27,7 @@ STA_mode="STA_tuning"
 for i in $gpu_ids; do
     port=$((base_port+i))
     CUDA_VISIBLE_DEVICES=$i MASTER_PORT=$port python examples/inference/sta_mask_search/wan_example.py \
-        --prompt_path ./assets/prompt_extend_${i}.txt \
+        --prompt_path ./assets/prompt_${i}.txt \
         --output_path $output_path \
         --STA_mode $STA_mode \
         --skip_time_steps $skip_time_steps &

fastvideo/v1/STA_configuration.py

@@ -6,6 +6,8 @@ from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 
+from fastvideo.v1.utils import dict_to_3d_list
+
 
 def configure_sta(mode: str = 'STA_searching',
                   layer_num: int = 40,
@@ -349,21 +351,6 @@ def select_best_mask_strategy(
     return best_mask_strategy, overall_sparsity, strategy_counts
 
 
-def dict_to_3d_list(mask_strategy: Optional[Dict[str, List[int]]],
-                    t_max: int = 50,
-                    l_max: int = 60,
-                    h_max: int = 24) -> List[List[List[Optional[List[int]]]]]:
-    result: List[List[List[Optional[List[int]]]]] = [[[
-        None for _ in range(h_max)
-    ] for _ in range(l_max)] for _ in range(t_max)]
-    if mask_strategy is None:
-        return result
-    for key, value in mask_strategy.items():
-        t, layer_idx, h = map(int, key.split('_'))
-        result[t][layer_idx][h] = value
-    return result
-
-
 def save_mask_search_results(
         mask_search_final_result: List[Dict[str, List[float]]],
         prompt: str,

fastvideo/v1/attention/__init__.py

@@ -10,8 +10,8 @@ from fastvideo.v1.attention.selector import get_attn_backend
 
 __all__ = [
     "DistributedAttention",
-    "DistributedAttention_VSA",
     "LocalAttention",
+    "DistributedAttention_VSA",
     "AttentionBackend",
     "AttentionMetadata",
     "AttentionMetadataBuilder",

fastvideo/v1/attention/backends/sliding_tile_attn.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import json
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, List, Optional, Type
 
 import torch
 from einops import rearrange
@@ -17,33 +17,11 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import ForwardContext, get_forward_context
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.utils import dict_to_3d_list
 
 logger = init_logger(__name__)
 
 
-# TODO(will-refactor): move this to a utils file
-def dict_to_3d_list(
-        mask_strategy: Dict[str,
-                            Any]) -> List[List[List[Optional[torch.Tensor]]]]:
-    indices = [tuple(map(int, key.split('_'))) for key in mask_strategy]
-
-    max_timesteps_idx = max(
-        timesteps_idx for timesteps_idx, layer_idx, head_idx in indices) + 1
-    max_layer_idx = max(layer_idx
-                        for timesteps_idx, layer_idx, head_idx in indices) + 1
-    max_head_idx = max(head_idx
-                       for timesteps_idx, layer_idx, head_idx in indices) + 1
-
-    result = [[[None for _ in range(max_head_idx)]
-               for _ in range(max_layer_idx)] for _ in range(max_timesteps_idx)]
-
-    for key, value in mask_strategy.items():
-        timesteps_idx, layer_idx, head_idx = map(int, key.split('_'))
-        result[timesteps_idx][layer_idx][head_idx] = value
-
-    return result
-
-
 class RangeDict(dict):
 
     def __getitem__(self, item: int) -> str:

fastvideo/v1/configs/pipelines/base.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import json
 from dataclasses import asdict, dataclass, field, fields
+from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
 
 import torch
@@ -16,6 +17,15 @@ from fastvideo.v1.utils import (FlexibleArgumentParser, StoreBoolean,
 logger = init_logger(__name__)
 
 
+class STA_Mode(str, Enum):
+    """STA (Sliding Tile Attention) modes."""
+    STA_INFERENCE = "STA_inference"
+    STA_SEARCHING = "STA_searching"
+    STA_TUNING = "STA_tuning"
+    STA_TUNING_CFG = "STA_tuning_cfg"
+    NONE = None
+
+
 def preprocess_text(prompt: str) -> str:
     return prompt
 
@@ -76,7 +86,7 @@ class PipelineConfig:
 
     # STA (Sliding Tile Attention) parameters
     mask_strategy_file_path: Optional[str] = None
-    STA_mode: Optional[str] = None
+    STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
     skip_time_steps: int = 15
 
     # Compilation

fastvideo/v1/fastvideo_args.py

@@ -8,7 +8,7 @@ from contextlib import contextmanager
 from dataclasses import field
 from typing import Any, Dict, List, Optional
 
-from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.v1.configs.pipelines.base import PipelineConfig, STA_Mode
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import FlexibleArgumentParser, StoreBoolean
 
@@ -63,7 +63,7 @@ class FastVideoArgs:
 
     # STA (Sliding Tile Attention) parameters
     mask_strategy_file_path: Optional[str] = None
-    STA_mode: Optional[str] = None
+    STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
     skip_time_steps: int = 15
 
     # Compilation
@@ -74,6 +74,9 @@ class FastVideoArgs:
     # VSA parameters
     VSA_sparsity: float = 0.0  # inference/validation sparsity
 
+    # Stage verification
+    enable_stage_verification: bool = True
+
     @property
     def training_mode(self) -> bool:
         return not self.inference_mode
@@ -178,12 +181,10 @@ class FastVideoArgs:
         parser.add_argument(
             "--STA-mode",
             type=str,
-            default=FastVideoArgs.STA_mode,
-            choices=[
-                "STA_inference", "STA_searching", "STA_tuning",
-                "STA_tuning_cfg", None
-            ],
-            help="STA mode",
+            default=FastVideoArgs.STA_mode.value,
+            choices=[mode.value for mode in STA_Mode],
+            help=
+            "STA mode contains STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None",
         )
         parser.add_argument(
             "--skip-time-steps",
@@ -231,6 +232,14 @@ class FastVideoArgs:
             help="Validation sparsity for VSA",
         )
 
+        # Stage verification
+        parser.add_argument(
+            "--enable-stage-verification",
+            action=StoreBoolean,
+            default=FastVideoArgs.enable_stage_verification,
+            help="Enable input/output verification for pipeline stages",
+        )
+
         # Add pipeline configuration arguments
         PipelineConfig.add_cli_args(parser)
 

fastvideo/v1/forward_context.py

@@ -5,15 +5,16 @@ import time
 from collections import defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
-# if TYPE_CHECKING:
-from fastvideo.v1.attention import AttentionMetadata
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+
+if TYPE_CHECKING:
+    from fastvideo.v1.attention import AttentionMetadata
+    from fastvideo.v1.pipelines import ForwardBatch
 
 logger = init_logger(__name__)
 
@@ -36,13 +37,13 @@ class ForwardContext:
     # attn_layers: Dict[str, Any]
     # TODO: extend to support per-layer dynamic forward context
     attn_metadata: "AttentionMetadata"  # set dynamically for each forward pass
-    forward_batch: Optional[ForwardBatch] = None
+    forward_batch: Optional["ForwardBatch"] = None
 
 
-_forward_context: Optional[ForwardContext] = None
+_forward_context: Optional["ForwardContext"] = None
 
 
-def get_forward_context() -> ForwardContext:
+def get_forward_context() -> "ForwardContext":
     """Get the current forward context."""
     assert _forward_context is not None, (
         "Forward context is not set. "
@@ -54,7 +55,7 @@ def get_forward_context() -> ForwardContext:
 @contextmanager
 def set_forward_context(current_timestep,
                         attn_metadata,
-                        forward_batch: Optional[ForwardBatch] = None,
+                        forward_batch: Optional["ForwardBatch"] = None,
                         fastvideo_args: Optional[FastVideoArgs] = None):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.

fastvideo/v1/layers/layernorm.py

@@ -5,6 +5,7 @@ from typing import Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 from fastvideo.v1.layers.custom_op import CustomOp
 
@@ -95,6 +96,22 @@ class ScaleResidual(nn.Module):
         return residual + x * gate
 
 
+# adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py
+# NOTE(will): Needed to match behavior of diffusers and wan2.1 even while using
+# FSDP's MixedPrecisionPolicy
+class FP32LayerNorm(nn.LayerNorm):
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        origin_dtype = inputs.dtype
+        return F.layer_norm(
+            inputs.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        ).to(origin_dtype)
+
+
 class ScaleResidualLayerNormScaleShift(nn.Module):
     """
     Fused operation that combines:
@@ -112,6 +129,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
         eps: float = 1e-6,
         elementwise_affine: bool = False,
         dtype: torch.dtype = torch.float32,
+        compute_dtype: torch.dtype | None = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -121,10 +139,15 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
                                 eps=eps,
                                 dtype=dtype)
         elif norm_type == "layer":
-            self.norm = nn.LayerNorm(hidden_size,
-                                     elementwise_affine=elementwise_affine,
-                                     eps=eps,
-                                     dtype=dtype)
+            if compute_dtype == torch.float32:
+                self.norm = FP32LayerNorm(hidden_size,
+                                          elementwise_affine=elementwise_affine,
+                                          eps=eps)
+            else:
+                self.norm = nn.LayerNorm(hidden_size,
+                                         elementwise_affine=elementwise_affine,
+                                         eps=eps,
+                                         dtype=dtype)
         else:
             raise NotImplementedError(f"Norm type {norm_type} not implemented")
 
@@ -163,18 +186,25 @@ class LayerNormScaleShift(nn.Module):
         eps: float = 1e-6,
         elementwise_affine: bool = False,
         dtype: torch.dtype = torch.float32,
+        compute_dtype: torch.dtype | None = None,
         prefix: str = "",
     ):
         super().__init__()
+        self.compute_dtype = compute_dtype
         if norm_type == "rms":
             self.norm = RMSNorm(hidden_size,
                                 has_weight=elementwise_affine,
                                 eps=eps)
         elif norm_type == "layer":
-            self.norm = nn.LayerNorm(hidden_size,
-                                     elementwise_affine=elementwise_affine,
-                                     eps=eps,
-                                     dtype=dtype)
+            if self.compute_dtype == torch.float32:
+                self.norm = FP32LayerNorm(hidden_size,
+                                          elementwise_affine=elementwise_affine,
+                                          eps=eps)
+            else:
+                self.norm = nn.LayerNorm(hidden_size,
+                                         elementwise_affine=elementwise_affine,
+                                         eps=eps,
+                                         dtype=dtype)
         else:
             raise NotImplementedError(f"Norm type {norm_type} not implemented")
 
@@ -182,4 +212,7 @@ class LayerNormScaleShift(nn.Module):
                 scale: torch.Tensor) -> torch.Tensor:
         """Apply ln followed by scale and shift in a single fused operation."""
         normalized = self.norm(x)
-        return normalized * (1.0 + scale) + shift
+        if self.compute_dtype == torch.float32:
+            return (normalized.float() * (1.0 + scale) + shift).to(x.dtype)
+        else:
+            return normalized * (1.0 + scale) + shift

fastvideo/v1/models/dits/wanvideo.py

@@ -14,8 +14,8 @@ from fastvideo.v1.configs.models.dits import WanVideoConfig
 from fastvideo.v1.configs.sample.wan import WanTeaCacheParams
 from fastvideo.v1.distributed.parallel_state import get_sp_world_size
 from fastvideo.v1.forward_context import get_forward_context
-from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, RMSNorm,
-                                           ScaleResidual,
+from fastvideo.v1.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift,
+                                           RMSNorm, ScaleResidual,
                                            ScaleResidualLayerNormScaleShift)
 from fastvideo.v1.layers.linear import ReplicatedLinear
 # from torch.nn import RMSNorm
@@ -34,9 +34,9 @@ class WanImageEmbedding(torch.nn.Module):
     def __init__(self, in_features: int, out_features: int):
         super().__init__()
 
-        self.norm1 = nn.LayerNorm(in_features)
+        self.norm1 = FP32LayerNorm(in_features)
         self.ff = MLP(in_features, in_features, out_features, act_type="gelu")
-        self.norm2 = nn.LayerNorm(out_features)
+        self.norm2 = FP32LayerNorm(out_features)
 
     def forward(self,
                 encoder_hidden_states_image: torch.Tensor) -> torch.Tensor:
@@ -232,7 +232,7 @@ class WanTransformerBlock(nn.Module):
         super().__init__()
 
         # 1. Self-attention
-        self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
         self.to_q = ReplicatedLinear(dim, dim, bias=True)
         self.to_k = ReplicatedLinear(dim, dim, bias=True)
         self.to_v = ReplicatedLinear(dim, dim, bias=True)
@@ -263,7 +263,8 @@ class WanTransformerBlock(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=True,
-            dtype=torch.float32)
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
 
         # 2. Cross-attention
         if added_kv_proj_dim is not None:
@@ -283,7 +284,8 @@ class WanTransformerBlock(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=False,
-            dtype=torch.float32)
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
@@ -375,7 +377,7 @@ class WanTransformerBlock_VSA(nn.Module):
         super().__init__()
 
         # 1. Self-attention
-        self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
         self.to_q = ReplicatedLinear(dim, dim, bias=True)
         self.to_k = ReplicatedLinear(dim, dim, bias=True)
         self.to_v = ReplicatedLinear(dim, dim, bias=True)
@@ -407,7 +409,8 @@ class WanTransformerBlock_VSA(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=True,
-            dtype=torch.float32)
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
 
         # 2. Cross-attention
         if added_kv_proj_dim is not None:
@@ -427,7 +430,8 @@ class WanTransformerBlock_VSA(nn.Module):
             norm_type="layer",
             eps=eps,
             elementwise_affine=False,
-            dtype=torch.float32)
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
@@ -564,7 +568,8 @@ class WanTransformer3DModel(CachableDiT):
                                             norm_type="layer",
                                             eps=config.eps,
                                             elementwise_affine=False,
-                                            dtype=torch.float32)
+                                            dtype=torch.float32,
+                                            compute_dtype=torch.float32)
         self.proj_out = nn.Linear(
             inner_dim, config.out_channels * math.prod(config.patch_size))
         self.scale_shift_table = nn.Parameter(
@@ -572,6 +577,17 @@ class WanTransformer3DModel(CachableDiT):
 
         self.gradient_checkpointing = False
 
+        # For type checking
+        self.previous_e0_even = None
+        self.previous_e0_odd = None
+        self.previous_residual_even = None
+        self.previous_residual_odd = None
+        self.is_even = True
+        self.should_calc_even = True
+        self.should_calc_odd = True
+        self.accumulated_rel_l1_distance_even = 0
+        self.accumulated_rel_l1_distance_odd = 0
+        self.cnt = 0
         self.__post_init__()
 
     def forward(self,
@@ -658,7 +674,7 @@ class WanTransformer3DModel(CachableDiT):
         # 5. Output norm, projection & unpatchify
         shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
                                                                           dim=1)
-        hidden_states = self.norm_out(hidden_states.float(), shift, scale)
+        hidden_states = self.norm_out(hidden_states, shift, scale)
         hidden_states = self.proj_out(hidden_states)
 
         hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,

fastvideo/v1/models/utils.py

@@ -51,11 +51,6 @@ def auto_attributes(init_func):
     return wrapper
 
 
-def set_random_seed(seed: int) -> None:
-    from fastvideo.v1.platforms import current_platform
-    current_platform.seed_everything(seed)
-
-
 def set_weight_attrs(
     weight: torch.Tensor,
     weight_attrs: Optional[Dict[str, Any]],

fastvideo/v1/pipelines/__init__.py

@@ -11,7 +11,8 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
 from fastvideo.v1.pipelines.lora_pipeline import LoRAPipeline
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.pipelines.pipeline_batch_info import (ForwardBatch,
+                                                        TrainingBatch)
 from fastvideo.v1.pipelines.pipeline_registry import PipelineRegistry
 from fastvideo.v1.utils import (maybe_download_model,
                                 verify_model_config_and_directory)
@@ -63,4 +64,5 @@ __all__ = [
     "PipelineRegistry",
     "ForwardBatch",
     "LoRAPipeline",
+    "TrainingBatch",
 ]

fastvideo/v1/pipelines/composed_pipeline_base.py

@@ -298,3 +298,7 @@ class ComposedPipelineBase(ABC):
 
         # Return the output
         return batch
+
+    def train(self) -> None:
+        raise NotImplementedError(
+            "if training_mode is True, the pipeline must implement this method")

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -13,6 +13,7 @@ from typing import Any, Dict, List, Optional, Union
 
 import torch
 
+from fastvideo.v1.attention import AttentionMetadata
 from fastvideo.v1.configs.sample.teacache import (TeaCacheParams,
                                                   WanTeaCacheParams)
 
@@ -136,3 +137,30 @@ class ForwardBatch:
 
     def __str__(self):
         return pprint.pformat(asdict(self), indent=2, width=120)
+
+
+@dataclass
+class TrainingBatch:
+    current_timestep: int = 0
+    current_vsa_sparsity: float = 0.0
+
+    # Dataloader batch outputs
+    latents: Optional[torch.Tensor] = None
+    encoder_hidden_states: Optional[torch.Tensor] = None
+    encoder_attention_mask: Optional[torch.Tensor] = None
+    info: Optional[Dict[str, Any]] = None
+
+    # Transformer inputs
+    noisy_model_input: Optional[torch.Tensor] = None
+    timesteps: Optional[torch.Tensor] = None
+    sigmas: Optional[torch.Tensor] = None
+    noise: Optional[torch.Tensor] = None
+
+    attn_metadata: Optional[AttentionMetadata] = None
+
+    # Training loss
+    loss: torch.Tensor | None = None
+
+    # Training outputs
+    total_loss: float | None = None
+    grad_norm: float | None = None

fastvideo/v1/pipelines/stages/base.py

@@ -15,10 +15,16 @@ import torch
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
 
 logger = init_logger(__name__)
 
 
+class StageVerificationError(Exception):
+    """Exception raised when stage verification fails."""
+    pass
+
+
 class PipelineStage(ABC):
     """
     Abstract base class for all pipeline stages.
@@ -28,6 +34,70 @@ class PipelineStage(ABC):
     for a specific part of the process, such as prompt encoding, latent preparation, etc.
     """
 
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """
+        Verify the input for the stage.
+
+        Example:
+            from fastvideo.v1.pipelines.stages.validators import V, VerificationResult
+            
+            def verify_input(self, batch, fastvideo_args):
+                result = VerificationResult()
+                result.add_check("height", batch.height, V.positive_int_divisible(8))
+                result.add_check("width", batch.width, V.positive_int_divisible(8))
+                result.add_check("image_latent", batch.image_latent, V.is_tensor)
+                return result
+
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+
+        Returns:
+            A VerificationResult containing the verification status.
+        
+        """
+        # Default implementation - no verification
+        return VerificationResult()
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """
+        Verify the output for the stage.
+
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+
+        Returns:
+            A VerificationResult containing the verification status.
+        """
+        # Default implementation - no verification
+        return VerificationResult()
+
+    def _run_verification(self, verification_result: VerificationResult,
+                          stage_name: str, verification_type: str) -> None:
+        """
+        Run verification and raise errors if any checks fail.
+        
+        Args:
+            verification_result: Results from verify_input or verify_output
+            stage_name: Name of the current stage
+            verification_type: "input" or "output"
+        """
+        if not verification_result.is_valid():
+            failed_fields = verification_result.get_failed_fields()
+            if failed_fields:
+                # Get detailed failure information
+                detailed_summary = verification_result.get_failure_summary()
+
+                failed_fields_str = ", ".join(failed_fields)
+                error_msg = (
+                    f"{verification_type.capitalize()} verification failed for {stage_name}: "
+                    f"Failed fields: {failed_fields_str}\n"
+                    f"Details: {detailed_summary}")
+                raise StageVerificationError(error_msg)
+
     @property
     def device(self) -> torch.device:
         """Get the device for this stage."""
@@ -48,7 +118,7 @@ class PipelineStage(ABC):
         fastvideo_args: FastVideoArgs,
     ) -> ForwardBatch:
         """
-        Execute the stage's processing on the batch with optional logging.
+        Execute the stage's processing on the batch with optional verification and logging.
         Should not be overridden by subclasses.
         
         Args:
@@ -58,34 +128,56 @@ class PipelineStage(ABC):
         Returns:
             The updated batch information after this stage's processing.
         """
-        # if envs.ENABLE_STAGE_LOGGING:
+        stage_name = self.__class__.__name__
+
+        # Check if verification is enabled (simple approach for prototype)
+        enable_verification = getattr(fastvideo_args,
+                                      'enable_stage_verification', False)
+
+        if enable_verification:
+            # Pre-execution input verification
+            try:
+                input_result = self.verify_input(batch, fastvideo_args)
+                self._run_verification(input_result, stage_name, "input")
+            except Exception as e:
+                logger.error("Input verification failed for %s: %s", stage_name,
+                             str(e))
+                raise
+
+        # Execute the actual stage logic
+        # envs.ENABLE_STAGE_LOGGING
         if False:
-            self._logger.info("[%s] Starting execution", self._stage_name)
+            self._logger.info("[%s] Starting execution", stage_name)
             start_time = time.perf_counter()
 
             try:
-                # Call the actual implementation
-                result = self._call_implementation(batch, fastvideo_args)
-
+                result = self.forward(batch, fastvideo_args)
                 execution_time = time.perf_counter() - start_time
                 self._logger.info("[%s] Execution completed in %s ms",
-                                  self._stage_name, execution_time * 1000)
-
-                return result
+                                  stage_name, execution_time * 1000)
             except Exception as e:
                 execution_time = time.perf_counter() - start_time
                 self._logger.error(
-                    "[%s] Error during execution after %s ms: %s",
-                    self._stage_name, execution_time * 1000, e)
-                self._logger.error("[%s] Traceback: %s", self._stage_name,
+                    "[%s] Error during execution after %s ms: %s", stage_name,
+                    execution_time * 1000, e)
+                self._logger.error("[%s] Traceback: %s", stage_name,
                                    traceback.format_exc())
-
-                # Re-raise the exception
                 raise
         else:
-            # Just call the implementation directly if logging is disabled
-            # TODO(will): Also handle backward
-            return self.forward(batch, fastvideo_args)
+            # Direct execution (current behavior)
+            result = self.forward(batch, fastvideo_args)
+
+        if enable_verification:
+            # Post-execution output verification
+            try:
+                output_result = self.verify_output(result, fastvideo_args)
+                self._run_verification(output_result, stage_name, "output")
+            except Exception as e:
+                logger.error("Output verification failed for %s: %s",
+                             stage_name, str(e))
+                raise
+
+        return result
 
     @abstractmethod
     def forward(

fastvideo/v1/pipelines/stages/conditioning.py

@@ -9,6 +9,8 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 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
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
 
 logger = init_logger(__name__)
 
@@ -69,3 +71,24 @@ class ConditioningStage(PipelineStage):
                 [batch.negative_attention_mask_2, batch.attention_mask_2])
 
         return batch
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify conditioning stage inputs."""
+        result = VerificationResult()
+        result.add_check("do_classifier_free_guidance",
+                         batch.do_classifier_free_guidance, V.bool_value)
+        result.add_check("guidance_scale", batch.guidance_scale,
+                         V.positive_float)
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        result.add_check(
+            "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x:
+            not batch.do_classifier_free_guidance or V.list_not_empty(x))
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify conditioning stage outputs."""
+        result = VerificationResult()
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        return result

fastvideo/v1/pipelines/stages/decoding.py

@@ -11,6 +11,8 @@ from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 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
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
 from fastvideo.v1.utils import PRECISION_TO_TYPE
 
 logger = init_logger(__name__)
@@ -27,6 +29,23 @@ class DecodingStage(PipelineStage):
     def __init__(self, vae) -> None:
         self.vae: ParallelTiledVAE = vae
 
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify decoding stage inputs."""
+        result = VerificationResult()
+        # Denoised latents for VAE decoding: [batch_size, channels, frames, height_latents, width_latents]
+        result.add_check("latents", batch.latents,
+                         [V.is_tensor, V.with_dims(5)])
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify decoding stage outputs."""
+        result = VerificationResult()
+        # Decoded video/images: [batch_size, channels, frames, height, width]
+        result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)])
+        return result
+
     def forward(
         self,
         batch: ForwardBatch,

fastvideo/v1/pipelines/stages/denoising.py

@@ -3,15 +3,15 @@
 Denoising stage for diffusion pipelines.
 """
 
-import importlib.util
 import inspect
-from typing import Any, Dict, Iterable, List, Optional
+from typing import Any, Dict, Iterable, Optional
 
 import torch
 from einops import rearrange
 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_sp_parallel_rank, get_sp_world_size,
                                       get_torch_device, get_world_group)
 from fastvideo.v1.distributed.communication_op import (
@@ -21,19 +21,24 @@ from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
 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
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
 from fastvideo.v1.platforms import AttentionBackendEnum
+from fastvideo.v1.utils import dict_to_3d_list
 
-st_attn_available = False
-if importlib.util.find_spec("st_attn") is not None:
-    st_attn_available = True
+try:
     from fastvideo.v1.attention.backends.sliding_tile_attn import (
         SlidingTileAttentionBackend)
+    st_attn_available = True
+except ImportError:
+    st_attn_available = False
 
-vsa_available = False
-if importlib.util.find_spec("vsa") is not None:
-    vsa_available = True
+try:
     from fastvideo.v1.attention.backends.video_sparse_attn import (
         VideoSparseAttentionBackend)
+    vsa_available = True
+except ImportError:
+    vsa_available = False
 
 logger = init_logger(__name__)
 
@@ -117,20 +122,6 @@ class DenoisingStage(PipelineStage):
         num_warmup_steps = len(
             timesteps) - num_inference_steps * self.scheduler.order
 
-        # Create 3D list for mask strategy
-        def dict_to_3d_list(mask_strategy,
-                            t_max=50,
-                            l_max=60,
-                            h_max=24) -> List:
-            result = [[[None for _ in range(h_max)] for _ in range(l_max)]
-                      for _ in range(t_max)]
-            if mask_strategy is None:
-                return result
-            for key, value in mask_strategy.items():
-                t, layer, h = map(int, key.split('_'))
-                result[t][layer][h] = value
-            return result
-
         # Prepare image latents and embeddings for I2V generation
         image_embeds = batch.image_embeds
         if len(image_embeds) > 0:
@@ -143,7 +134,8 @@ class DenoisingStage(PipelineStage):
             self.transformer.forward,
             {
                 "encoder_hidden_states_image": image_embeds,
-                "mask_strategy": dict_to_3d_list(None)
+                "mask_strategy": dict_to_3d_list(
+                    None, t_max=50, l_max=60, h_max=24)
             },
         )
 
@@ -304,7 +296,7 @@ class DenoisingStage(PipelineStage):
         batch.latents = latents
 
         # Save STA mask search results if needed
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.STA_mode == 'STA_searching':
+        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.STA_mode == STA_Mode.STA_SEARCHING:
             self.save_sta_search_results(batch)
 
         if fastvideo_args.use_cpu_offload:
@@ -402,7 +394,7 @@ class DenoisingStage(PipelineStage):
             raise NotImplementedError(
                 "STA mask search/tuning is not supported for this resolution")
 
-        if STA_mode == "STA_searching" or STA_mode == "STA_tuning" or STA_mode == "STA_tuning_cfg":
+        if STA_mode == STA_Mode.STA_SEARCHING or STA_mode == STA_Mode.STA_TUNING or STA_mode == STA_Mode.STA_TUNING_CFG:
             size = (batch.width, batch.height)
             if size == (1280, 768):
                 # TODO: make it configurable
@@ -424,18 +416,18 @@ class DenoisingStage(PipelineStage):
             layer_num += self.transformer.config.num_single_layers
         head_num = self.transformer.config.num_attention_heads
 
-        if STA_mode == "STA_searching":
+        if STA_mode == STA_Mode.STA_SEARCHING:
             STA_param = configure_sta(
-                mode='STA_searching',
+                mode=STA_Mode.STA_SEARCHING,
                 layer_num=layer_num,
                 head_num=head_num,
                 time_step_num=timesteps_num,
                 mask_candidates=sparse_mask_candidates_searching +
                 full_mask,  # last is full mask; Can add more sparse masks while keep last one as full mask
             )
-        elif STA_mode == 'STA_tuning':
+        elif STA_mode == STA_Mode.STA_TUNING:
             STA_param = configure_sta(
-                mode='STA_tuning',
+                mode=STA_Mode.STA_TUNING,
                 layer_num=layer_num,
                 head_num=head_num,
                 time_step_num=timesteps_num,
@@ -448,9 +440,9 @@ class DenoisingStage(PipelineStage):
                 save_dir=
                 f'output/mask_search_strategy_{size[0]}x{size[1]}/',  # Custom save directory
                 timesteps=timesteps_num)
-        elif STA_mode == 'STA_tuning_cfg':
+        elif STA_mode == STA_Mode.STA_TUNING_CFG:
             STA_param = configure_sta(
-                mode='STA_tuning_cfg',
+                mode=STA_Mode.STA_TUNING_CFG,
                 layer_num=layer_num,
                 head_num=head_num,
                 time_step_num=timesteps_num,
@@ -463,12 +455,12 @@ class DenoisingStage(PipelineStage):
                 skip_time_steps=skip_time_steps,
                 save_dir=f'output/mask_search_strategy_{size[0]}x{size[1]}/',
                 timesteps=timesteps_num)
-        elif STA_mode == 'STA_inference':
+        elif STA_mode == STA_Mode.STA_INFERENCE:
             import fastvideo.v1.envs as envs
             config_file = envs.FASTVIDEO_ATTENTION_CONFIG
             if config_file is None:
                 raise ValueError("FASTVIDEO_ATTENTION_CONFIG is not set")
-            STA_param = configure_sta(mode='STA_inference',
+            STA_param = configure_sta(mode=STA_Mode.STA_INFERENCE,
                                       layer_num=layer_num,
                                       head_num=head_num,
                                       time_step_num=timesteps_num,
@@ -517,3 +509,37 @@ class DenoisingStage(PipelineStage):
                 mask_strategies=sparse_mask_candidates_searching,
                 output_dir=f'output/mask_search_result_neg_{size[0]}x{size[1]}/'
             )
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify denoising stage inputs."""
+        result = VerificationResult()
+        result.add_check("timesteps", batch.timesteps,
+                         [V.is_tensor, V.min_dims(1)])
+        result.add_check("latents", batch.latents,
+                         [V.is_tensor, V.with_dims(5)])
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        result.add_check("image_embeds", batch.image_embeds, V.is_list)
+        result.add_check("image_latent", batch.image_latent,
+                         V.none_or_tensor_with_dims(5))
+        result.add_check("num_inference_steps", batch.num_inference_steps,
+                         V.positive_int)
+        result.add_check("guidance_scale", batch.guidance_scale,
+                         V.positive_float)
+        result.add_check("eta", batch.eta, V.non_negative_float)
+        result.add_check("generator", batch.generator,
+                         V.generator_or_list_generators)
+        result.add_check("do_classifier_free_guidance",
+                         batch.do_classifier_free_guidance, V.bool_value)
+        result.add_check(
+            "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x:
+            not batch.do_classifier_free_guidance or V.list_not_empty(x))
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify denoising stage outputs."""
+        result = VerificationResult()
+        result.add_check("latents", batch.latents,
+                         [V.is_tensor, V.with_dims(5)])
+        return result

fastvideo/v1/pipelines/stages/encoding.py

@@ -16,6 +16,8 @@ from fastvideo.v1.models.vision_utils import (get_default_height_width,
                                               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
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
 from fastvideo.v1.utils import PRECISION_TO_TYPE
 
 logger = init_logger(__name__)
@@ -174,3 +176,23 @@ class EncodingStage(PipelineStage):
             image = normalize(image)
 
         return image
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify encoding stage inputs."""
+        result = VerificationResult()
+        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,
+                         V.generator_or_list_generators)
+        result.add_check("num_frames", batch.num_frames, V.positive_int)
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify encoding stage outputs."""
+        result = VerificationResult()
+        result.add_check("image_latent", batch.image_latent,
+                         [V.is_tensor, V.with_dims(5)])
+        return result

fastvideo/v1/pipelines/stages/image_encoding.py

@@ -14,6 +14,8 @@ 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
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
 
 logger = init_logger(__name__)
 
@@ -71,3 +73,19 @@ class ImageEncodingStage(PipelineStage):
             torch.cuda.empty_cache()
 
         return batch
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify image encoding stage inputs."""
+        result = VerificationResult()
+        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
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify image encoding stage outputs."""
+        result = VerificationResult()
+        result.add_check("image_embeds", batch.image_embeds,
+                         V.list_of_tensors_dims(3))
+        return result

fastvideo/v1/pipelines/stages/input_validation.py

@@ -9,9 +9,14 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 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,
+                                                      VerificationResult)
 
 logger = init_logger(__name__)
 
+# Alias for convenience
+V = StageValidators
+
 
 class InputValidationStage(PipelineStage):
     """
@@ -87,3 +92,31 @@ class InputValidationStage(PipelineStage):
             )
 
         return batch
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify input validation stage inputs."""
+        result = VerificationResult()
+        result.add_check("seed", batch.seed, [V.not_none, V.positive_int])
+        result.add_check("num_videos_per_prompt", batch.num_videos_per_prompt,
+                         V.positive_int)
+        result.add_check(
+            "prompt_or_embeds", None, lambda _: V.string_or_list_strings(
+                batch.prompt) or V.list_not_empty(batch.prompt_embeds))
+        result.add_check("height", batch.height, V.positive_int)
+        result.add_check("width", batch.width, V.positive_int)
+        result.add_check("num_inference_steps", batch.num_inference_steps,
+                         V.positive_int)
+        result.add_check(
+            "guidance_scale", batch.guidance_scale, lambda x: not batch.
+            do_classifier_free_guidance or V.positive_float(x))
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify input validation stage outputs."""
+        result = VerificationResult()
+        result.add_check("seeds", batch.seeds, V.list_not_empty)
+        result.add_check("generator", batch.generator,
+                         V.generator_or_list_generators)
+        return result

fastvideo/v1/pipelines/stages/latent_preparation.py

@@ -2,6 +2,7 @@
 """
 Latent preparation stage for diffusion pipelines.
 """
+
 from diffusers.utils.torch_utils import randn_tensor
 
 from fastvideo.v1.distributed import get_torch_device
@@ -9,6 +10,8 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 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
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
 
 logger = init_logger(__name__)
 
@@ -126,4 +129,32 @@ class LatentPreparationStage(PipelineStage):
             latent_num_frames = (video_length - 1) // temporal_scale_factor + 1
         else:  # stepvideo only
             latent_num_frames = video_length // 17 * 3
-        return latent_num_frames
+        return int(latent_num_frames)
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify latent preparation stage inputs."""
+        result = VerificationResult()
+        result.add_check(
+            "prompt_or_embeds", None, lambda _: V.string_or_list_strings(
+                batch.prompt) or V.list_not_empty(batch.prompt_embeds))
+        result.add_check("prompt_embeds", batch.prompt_embeds,
+                         V.list_of_tensors)
+        result.add_check("num_videos_per_prompt", batch.num_videos_per_prompt,
+                         V.positive_int)
+        result.add_check("generator", batch.generator,
+                         V.generator_or_list_generators)
+        result.add_check("num_frames", batch.num_frames, V.positive_int)
+        result.add_check("height", batch.height, V.positive_int)
+        result.add_check("width", batch.width, V.positive_int)
+        result.add_check("latents", batch.latents, V.none_or_tensor)
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify latent preparation stage outputs."""
+        result = VerificationResult()
+        result.add_check("latents", batch.latents,
+                         [V.is_tensor, V.with_dims(5)])
+        result.add_check("raw_latent_shape", batch.raw_latent_shape, V.is_tuple)
+        return result

fastvideo/v1/pipelines/stages/stepvideo_encoding.py

@@ -1,10 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
+
 import torch
 
+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.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
+from fastvideo.v1.pipelines.stages.validators import StageValidators as V
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
 
 logger = init_logger(__name__)
 
@@ -47,3 +51,29 @@ class StepvideoPromptEncodingStage(PipelineStage):
         batch.clip_embedding_pos = pos_clip
         batch.clip_embedding_neg = neg_clip
         return batch
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify stepvideo encoding stage inputs."""
+        result = VerificationResult()
+        result.add_check("prompt", batch.prompt, V.string_not_empty)
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify stepvideo encoding stage outputs."""
+        result = VerificationResult()
+        result.add_check("prompt_embeds", batch.prompt_embeds,
+                         [V.is_tensor, V.with_dims(3)])
+        result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds,
+                         [V.is_tensor, V.with_dims(3)])
+        result.add_check("prompt_attention_mask", batch.prompt_attention_mask,
+                         [V.is_tensor, V.with_dims(2)])
+        result.add_check("negative_attention_mask",
+                         batch.negative_attention_mask,
+                         [V.is_tensor, V.with_dims(2)])
+        result.add_check("clip_embedding_pos", batch.clip_embedding_pos,
+                         [V.is_tensor, V.with_dims(2)])
+        result.add_check("clip_embedding_neg", batch.clip_embedding_neg,
+                         [V.is_tensor, V.with_dims(2)])
+        return result

fastvideo/v1/pipelines/stages/text_encoding.py

@@ -12,6 +12,8 @@ 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
 from fastvideo.v1.pipelines.stages.base import PipelineStage
+from fastvideo.v1.pipelines.stages.validators import StageValidators as V
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
 
 logger = (__name__)
 
@@ -113,3 +115,30 @@ class TextEncodingStage(PipelineStage):
                 torch.cuda.empty_cache()
 
         return batch
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify text encoding stage inputs."""
+        result = VerificationResult()
+        result.add_check("prompt", batch.prompt, V.string_or_list_strings)
+        result.add_check(
+            "negative_prompt", batch.negative_prompt, lambda x: not batch.
+            do_classifier_free_guidance or V.string_not_empty(x))
+        result.add_check("do_classifier_free_guidance",
+                         batch.do_classifier_free_guidance, V.bool_value)
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.is_list)
+        result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds,
+                         V.none_or_list)
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify text encoding stage outputs."""
+        result = VerificationResult()
+        result.add_check("prompt_embeds", batch.prompt_embeds,
+                         V.list_of_tensors_min_dims(2))
+        result.add_check(
+            "negative_prompt_embeds", batch.negative_prompt_embeds,
+            lambda x: not batch.do_classifier_free_guidance or V.
+            list_of_tensors_with_min_dims(x, 2))
+        return result

fastvideo/v1/pipelines/stages/timestep_preparation.py

@@ -12,6 +12,8 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 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
+from fastvideo.v1.pipelines.stages.validators import VerificationResult
 
 logger = init_logger(__name__)
 
@@ -95,3 +97,22 @@ class TimestepPreparationStage(PipelineStage):
         batch.timesteps = timesteps
 
         return batch
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify timestep preparation stage inputs."""
+        result = VerificationResult()
+        result.add_check("num_inference_steps", batch.num_inference_steps,
+                         V.positive_int)
+        result.add_check("timesteps", batch.timesteps, V.none_or_tensor)
+        result.add_check("sigmas", batch.sigmas, V.none_or_list)
+        result.add_check("n_tokens", batch.n_tokens, V.none_or_positive_int)
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify timestep preparation stage outputs."""
+        result = VerificationResult()
+        result.add_check("timesteps", batch.timesteps,
+                         [V.is_tensor, V.with_dims(1)])
+        return result

fastvideo/v1/pipelines/stages/validators.py

@@ -0,0 +1,486 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Common validators for pipeline stage verification.
+
+This module provides reusable validation functions that can be used across
+all pipeline stages for input/output verification.
+"""
+
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+
+
+class StageValidators:
+    """Common validators for pipeline stages."""
+
+    @staticmethod
+    def not_none(value: Any) -> bool:
+        """Check if value is not None."""
+        return value is not None
+
+    @staticmethod
+    def positive_int(value: Any) -> bool:
+        """Check if value is a positive integer."""
+        return isinstance(value, int) and value > 0
+
+    @staticmethod
+    def positive_float(value: Any) -> bool:
+        """Check if value is a positive float."""
+        return isinstance(value, (int, float)) and value > 0
+
+    @staticmethod
+    def non_negative_float(value: Any) -> bool:
+        """Check if value is a non-negative float."""
+        return isinstance(value, (int, float)) and value >= 0
+
+    @staticmethod
+    def divisible_by(value: Any, divisor: int) -> bool:
+        """Check if value is divisible by divisor."""
+        return value is not None and isinstance(value,
+                                                int) and value % divisor == 0
+
+    @staticmethod
+    def is_tensor(value: Any) -> bool:
+        """Check if value is a torch tensor and doesn't contain NaN values."""
+        if not isinstance(value, torch.Tensor):
+            return False
+        return not torch.isnan(value).any().item()
+
+    @staticmethod
+    def tensor_with_dims(value: Any, dims: int) -> bool:
+        """Check if value is a tensor with specific dimensions and no NaN values."""
+        if not isinstance(value, torch.Tensor):
+            return False
+        if value.dim() != dims:
+            return False
+        return not torch.isnan(value).any().item()
+
+    @staticmethod
+    def tensor_min_dims(value: Any, min_dims: int) -> bool:
+        """Check if value is a tensor with at least min_dims dimensions and no NaN values."""
+        if not isinstance(value, torch.Tensor):
+            return False
+        if value.dim() < min_dims:
+            return False
+        return not torch.isnan(value).any().item()
+
+    @staticmethod
+    def tensor_shape_matches(value: Any, expected_shape: tuple) -> bool:
+        """Check if tensor shape matches expected shape (None for any size) and no NaN values."""
+        if not isinstance(value, torch.Tensor):
+            return False
+        if len(value.shape) != len(expected_shape):
+            return False
+        for actual, expected in zip(value.shape, expected_shape):
+            if expected is not None and actual != expected:
+                return False
+        return not torch.isnan(value).any().item()
+
+    @staticmethod
+    def list_not_empty(value: Any) -> bool:
+        """Check if value is a non-empty list."""
+        return isinstance(value, list) and len(value) > 0
+
+    @staticmethod
+    def list_length(value: Any, length: int) -> bool:
+        """Check if list has specific length."""
+        return isinstance(value, list) and len(value) == length
+
+    @staticmethod
+    def list_min_length(value: Any, min_length: int) -> bool:
+        """Check if list has at least min_length items."""
+        return isinstance(value, list) and len(value) >= min_length
+
+    @staticmethod
+    def string_not_empty(value: Any) -> bool:
+        """Check if value is a non-empty string."""
+        return isinstance(value, str) and len(value.strip()) > 0
+
+    @staticmethod
+    def string_or_list_strings(value: Any) -> bool:
+        """Check if value is a string or list of strings."""
+        if isinstance(value, str):
+            return True
+        if isinstance(value, list):
+            return all(isinstance(item, str) for item in value)
+        return False
+
+    @staticmethod
+    def bool_value(value: Any) -> bool:
+        """Check if value is a boolean."""
+        return isinstance(value, bool)
+
+    @staticmethod
+    def generator_or_list_generators(value: Any) -> bool:
+        """Check if value is a Generator or list of Generators."""
+        if isinstance(value, torch.Generator):
+            return True
+        if isinstance(value, list):
+            return all(isinstance(item, torch.Generator) for item in value)
+        return False
+
+    @staticmethod
+    def is_list(value: Any) -> bool:
+        """Check if value is a list (can be empty)."""
+        return isinstance(value, list)
+
+    @staticmethod
+    def is_tuple(value: Any) -> bool:
+        """Check if value is a tuple."""
+        return isinstance(value, tuple)
+
+    @staticmethod
+    def none_or_tensor(value: Any) -> bool:
+        """Check if value is None or a tensor without NaN values."""
+        if value is None:
+            return True
+        if not isinstance(value, torch.Tensor):
+            return False
+        return not torch.isnan(value).any().item()
+
+    @staticmethod
+    def list_of_tensors_with_dims(value: Any, dims: int) -> bool:
+        """Check if value is a non-empty list where all items are tensors with specific dimensions and no NaN values."""
+        if not isinstance(value, list) or len(value) == 0:
+            return False
+        for item in value:
+            if not isinstance(item, torch.Tensor):
+                return False
+            if item.dim() != dims:
+                return False
+            if torch.isnan(item).any().item():
+                return False
+        return True
+
+    @staticmethod
+    def list_of_tensors(value: Any) -> bool:
+        """Check if value is a non-empty list where all items are tensors without NaN values."""
+        if not isinstance(value, list) or len(value) == 0:
+            return False
+        for item in value:
+            if not isinstance(item, torch.Tensor):
+                return False
+            if torch.isnan(item).any().item():
+                return False
+        return True
+
+    @staticmethod
+    def list_of_tensors_with_min_dims(value: Any, min_dims: int) -> bool:
+        """Check if value is a non-empty list where all items are tensors with at least min_dims dimensions and no NaN values."""
+        if not isinstance(value, list) or len(value) == 0:
+            return False
+        for item in value:
+            if not isinstance(item, torch.Tensor):
+                return False
+            if item.dim() < min_dims:
+                return False
+            if torch.isnan(item).any().item():
+                return False
+        return True
+
+    @staticmethod
+    def none_or_tensor_with_dims(dims: int) -> Callable[[Any], bool]:
+        """Return a validator that checks if value is None or a tensor with specific dimensions and no NaN values."""
+
+        def validator(value: Any) -> bool:
+            if value is None:
+                return True
+            if not isinstance(value, torch.Tensor):
+                return False
+            if value.dim() != dims:
+                return False
+            return not torch.isnan(value).any().item()
+
+        return validator
+
+    @staticmethod
+    def none_or_list(value: Any) -> bool:
+        """Check if value is None or a list."""
+        return value is None or isinstance(value, list)
+
+    @staticmethod
+    def none_or_positive_int(value: Any) -> bool:
+        """Check if value is None or a positive integer."""
+        return value is None or (isinstance(value, int) and value > 0)
+
+    # Helper methods that return functions for common patterns
+    @staticmethod
+    def with_dims(dims: int) -> Callable[[Any], bool]:
+        """Return a validator that checks if tensor has specific dimensions and no NaN values."""
+
+        def validator(value: Any) -> bool:
+            return StageValidators.tensor_with_dims(value, dims)
+
+        return validator
+
+    @staticmethod
+    def min_dims(min_dims: int) -> Callable[[Any], bool]:
+        """Return a validator that checks if tensor has at least min_dims dimensions and no NaN values."""
+
+        def validator(value: Any) -> bool:
+            return StageValidators.tensor_min_dims(value, min_dims)
+
+        return validator
+
+    @staticmethod
+    def divisible(divisor: int) -> Callable[[Any], bool]:
+        """Return a validator that checks if value is divisible by divisor."""
+
+        def validator(value: Any) -> bool:
+            return StageValidators.divisible_by(value, divisor)
+
+        return validator
+
+    @staticmethod
+    def positive_int_divisible(divisor: int) -> Callable[[Any], bool]:
+        """Return a validator that checks if value is a positive integer divisible by divisor."""
+
+        def validator(value: Any) -> bool:
+            return (isinstance(value, int) and value > 0
+                    and StageValidators.divisible_by(value, divisor))
+
+        return validator
+
+    @staticmethod
+    def list_of_tensors_dims(dims: int) -> Callable[[Any], bool]:
+        """Return a validator that checks if value is a list of tensors with specific dimensions and no NaN values."""
+
+        def validator(value: Any) -> bool:
+            return StageValidators.list_of_tensors_with_dims(value, dims)
+
+        return validator
+
+    @staticmethod
+    def list_of_tensors_min_dims(min_dims: int) -> Callable[[Any], bool]:
+        """Return a validator that checks if value is a list of tensors with at least min_dims dimensions and no NaN values."""
+
+        def validator(value: Any) -> bool:
+            return StageValidators.list_of_tensors_with_min_dims(
+                value, min_dims)
+
+        return validator
+
+
+class ValidationFailure:
+    """Details about a specific validation failure."""
+
+    def __init__(self,
+                 validator_name: str,
+                 actual_value: Any,
+                 expected: Optional[str] = None,
+                 error_msg: Optional[str] = None):
+        self.validator_name = validator_name
+        self.actual_value = actual_value
+        self.expected = expected
+        self.error_msg = error_msg
+
+    def __str__(self) -> str:
+        parts = [f"Validator '{self.validator_name}' failed"]
+
+        if self.error_msg:
+            parts.append(f"Error: {self.error_msg}")
+
+        # Add actual value info (but limit very long representations)
+        actual_str = self._format_value(self.actual_value)
+        parts.append(f"Actual: {actual_str}")
+
+        if self.expected:
+            parts.append(f"Expected: {self.expected}")
+
+        return ". ".join(parts)
+
+    def _format_value(self, value: Any) -> str:
+        """Format a value for display in error messages."""
+        if value is None:
+            return "None"
+        elif isinstance(value, torch.Tensor):
+            return f"tensor(shape={list(value.shape)}, dtype={value.dtype})"
+        elif isinstance(value, list):
+            if len(value) == 0:
+                return "[]"
+            elif len(value) <= 3:
+                item_strs = [self._format_value(item) for item in value]
+                return f"[{', '.join(item_strs)}]"
+            else:
+                return f"list(length={len(value)}, first_item={self._format_value(value[0])})"
+        elif isinstance(value, str):
+            if len(value) > 50:
+                return f"'{value[:47]}...'"
+            else:
+                return f"'{value}'"
+        else:
+            return f"{type(value).__name__}({value})"
+
+
+class VerificationResult:
+    """Wrapper class for stage verification results."""
+
+    def __init__(self) -> None:
+        self._checks: Dict[str, bool] = {}
+        self._failures: Dict[str, List[ValidationFailure]] = {}
+
+    def add_check(
+        self, field_name: str, value: Any,
+        validators: Union[Callable[[Any], bool], List[Callable[[Any], bool]]]
+    ) -> 'VerificationResult':
+        """
+        Add a validation check for a field.
+        
+        Args:
+            field_name: Name of the field being checked
+            value: The actual value to validate
+            validators: Single validation function or list of validation functions.
+                       Each function will be called with the value as its first argument.
+            
+        Returns:
+            Self for method chaining
+            
+        Examples:
+            # Single validator
+            result.add_check("tensor", my_tensor, V.is_tensor)
+            
+            # Multiple validators (all must pass)
+            result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+            
+            # Using partial functions for parameters
+            result.add_check("height", batch.height, [V.not_none, V.divisible(8)])
+        """
+        if not isinstance(validators, list):
+            validators = [validators]
+
+        failures = []
+        all_passed = True
+
+        # Apply all validators and collect detailed failure info
+        for validator in validators:
+            try:
+                passed = validator(value)
+                if not passed:
+                    all_passed = False
+                    failure = self._create_validation_failure(validator, value)
+                    failures.append(failure)
+            except Exception as e:
+                # If any validator raises an exception, consider the check failed
+                all_passed = False
+                validator_name = getattr(validator, '__name__', str(validator))
+                failure = ValidationFailure(
+                    validator_name=validator_name,
+                    actual_value=value,
+                    error_msg=f"Exception during validation: {str(e)}")
+                failures.append(failure)
+
+        self._checks[field_name] = all_passed
+        if not all_passed:
+            self._failures[field_name] = failures
+
+        return self
+
+    def _create_validation_failure(self, validator: Callable,
+                                   value: Any) -> ValidationFailure:
+        """Create a ValidationFailure with detailed information."""
+        validator_name = getattr(validator, '__name__', str(validator))
+
+        # Try to extract meaningful expected value info based on validator type
+        expected = None
+        error_msg = None
+
+        # Handle common validator patterns
+        if hasattr(validator, '__closure__') and validator.__closure__:
+            # This is likely a closure (like our helper functions)
+            if 'dims' in validator_name or 'with_dims' in str(validator):
+                if isinstance(value, torch.Tensor):
+                    expected = f"tensor with {validator.__closure__[0].cell_contents} dimensions"
+                else:
+                    expected = "tensor with specific dimensions"
+            elif 'divisible' in str(validator):
+                expected = f"integer divisible by {validator.__closure__[0].cell_contents}"
+
+        # Handle specific validator types and check for NaN values
+        if validator_name == 'is_tensor':
+            expected = "torch.Tensor without NaN values"
+            if isinstance(value,
+                          torch.Tensor) and torch.isnan(value).any().item():
+                error_msg = f"tensor contains {torch.isnan(value).sum().item()} NaN values"
+        elif validator_name == 'positive_int':
+            expected = "positive integer"
+        elif validator_name == 'not_none':
+            expected = "non-None value"
+        elif validator_name == 'list_not_empty':
+            expected = "non-empty list"
+        elif validator_name == 'bool_value':
+            expected = "boolean value"
+        elif 'tensor_with_dims' in validator_name or 'tensor_min_dims' in validator_name:
+            if isinstance(value, torch.Tensor):
+                if torch.isnan(value).any().item():
+                    error_msg = f"tensor has {value.dim()} dimensions but contains {torch.isnan(value).sum().item()} NaN values"
+                else:
+                    error_msg = f"tensor has {value.dim()} dimensions"
+        elif validator_name == 'is_list':
+            expected = "list"
+        elif validator_name == 'none_or_tensor':
+            expected = "None or tensor without NaN values"
+            if isinstance(value,
+                          torch.Tensor) and torch.isnan(value).any().item():
+                error_msg = f"tensor contains {torch.isnan(value).sum().item()} NaN values"
+        elif validator_name == 'list_of_tensors':
+            expected = "non-empty list of tensors without NaN values"
+            if isinstance(value, list) and len(value) > 0:
+                nan_count = 0
+                for item in value:
+                    if isinstance(
+                            item,
+                            torch.Tensor) and torch.isnan(item).any().item():
+                        nan_count += torch.isnan(item).sum().item()
+                if nan_count > 0:
+                    error_msg = f"list contains tensors with total {nan_count} NaN values"
+        elif 'list_of_tensors_with_dims' in validator_name:
+            expected = "non-empty list of tensors with specific dimensions and no NaN values"
+            if isinstance(value, list) and len(value) > 0:
+                nan_count = 0
+                for item in value:
+                    if isinstance(
+                            item,
+                            torch.Tensor) and torch.isnan(item).any().item():
+                        nan_count += torch.isnan(item).sum().item()
+                if nan_count > 0:
+                    error_msg = f"list contains tensors with total {nan_count} NaN values"
+
+        return ValidationFailure(validator_name=validator_name,
+                                 actual_value=value,
+                                 expected=expected,
+                                 error_msg=error_msg)
+
+    def is_valid(self) -> bool:
+        """Check if all validations passed."""
+        return all(self._checks.values())
+
+    def get_failed_fields(self) -> List[str]:
+        """Get list of fields that failed validation."""
+        return [field for field, passed in self._checks.items() if not passed]
+
+    def get_detailed_failures(self) -> Dict[str, List[ValidationFailure]]:
+        """Get detailed failure information for each failed field."""
+        return self._failures.copy()
+
+    def get_failure_summary(self) -> str:
+        """Get a comprehensive summary of all validation failures."""
+        if self.is_valid():
+            return "All validations passed"
+
+        summary_parts = []
+        for field_name, failures in self._failures.items():
+            field_summary = f"\n  Field '{field_name}':"
+            for i, failure in enumerate(failures, 1):
+                field_summary += f"\n    {i}. {failure}"
+            summary_parts.append(field_summary)
+
+        return "Validation failures:" + "".join(summary_parts)
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for backward compatibility."""
+        return self._checks.copy()
+
+
+# Alias for convenience
+V = StageValidators

fastvideo/v1/platforms/cuda.py

@@ -123,14 +123,13 @@ class CudaPlatformBase(Platform):
                     SlidingTileAttentionBackend)
                 logger.info("Using Sliding Tile Attention backend.")
 
-                # Overwrite with the actual backend
-                envs.FASTVIDEO_ATTENTION_BACKEND = "SLIDING_TILE_ATTN"
                 return "fastvideo.v1.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend"
             except ImportError as e:
-                logger.info(e)
-                logger.info(
-                    "Sliding Tile Attention backend is not installed. Fall back to Flash Attention."
-                )
+                logger.error(
+                    "Failed to import Sliding Tile Attention backend: %s",
+                    str(e))
+                raise ImportError(
+                    "Sliding Tile Attention backend is not installed. ") from e
         elif selected_backend == AttentionBackendEnum.SAGE_ATTN:
             try:
                 from sageattention import sageattn  # noqa: F401
@@ -139,8 +138,6 @@ class CudaPlatformBase(Platform):
                     SageAttentionBackend)
                 logger.info("Using Sage Attention backend.")
 
-                # Overwrite with the actual backend
-                envs.FASTVIDEO_ATTENTION_BACKEND = "SAGE_ATTN"
                 return "fastvideo.v1.attention.backends.sage_attn.SageAttentionBackend"
             except ImportError as e:
                 logger.info(e)
@@ -155,14 +152,13 @@ class CudaPlatformBase(Platform):
                     VideoSparseAttentionBackend)
                 logger.info("Using Video Sparse Attention backend.")
 
-                # Overwrite with the actual backend
-                envs.FASTVIDEO_ATTENTION_BACKEND = "VIDEO_SPARSE_ATTN"
                 return "fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionBackend"
             except ImportError as e:
-                logger.info(e)
-                logger.info(
-                    "Video Sparse Attention backend is not installed. Fall back to Flash Attention."
-                )
+                logger.error(
+                    "Failed to import Video Sparse Attention backend: %s",
+                    str(e))
+                raise ImportError(
+                    "Video Sparse Attention backend is not installed. ") from e
         elif selected_backend == AttentionBackendEnum.TORCH_SDPA:
             logger.info("Using Torch SDPA backend.")
             return "fastvideo.v1.attention.backends.sdpa.SDPABackend"
@@ -209,14 +205,10 @@ class CudaPlatformBase(Platform):
         if target_backend == AttentionBackendEnum.TORCH_SDPA:
             logger.info("Using Torch SDPA backend.")
 
-            # Overwrite with the actual backend
-            envs.FASTVIDEO_ATTENTION_BACKEND = "TORCH_SDPA"
             return "fastvideo.v1.attention.backends.sdpa.SDPABackend"
 
         logger.info("Using Flash Attention backend.")
 
-        # Overwrite with the actual backend
-        envs.FASTVIDEO_ATTENTION_BACKEND = "FLASH_ATTN"
         return "fastvideo.v1.attention.backends.flash_attn.FlashAttentionBackend"
 
     @classmethod

fastvideo/v1/platforms/interface.py

@@ -169,6 +169,7 @@ class Platform:
             random.seed(seed)
             np.random.seed(seed)
             torch.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
 
     @classmethod
     def verify_model_arch(cls, model_arch: str) -> None:

fastvideo/v1/tests/inference/STA/test_STA_inference.py

@@ -0,0 +1,52 @@
+import os
+import sys
+import subprocess
+from pathlib import Path
+import pytest
+
+NUM_NODES = "1"
+NUM_GPUS_PER_NODE = "1"
+
+# Set environment variables
+os.environ["FASTVIDEO_ATTENTION_CONFIG"] = "assets/mask_strategy_wan.json"
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
+
+def test_inference():
+    """Test the inference functionality"""
+    # Create command as in wan_14B-STA.sh
+    cmd = [
+        "fastvideo", "generate",
+        "--model-path", "Wan-AI/Wan2.1-T2V-14B-Diffusers",
+        "--sp-size", "1",
+        "--tp-size", "1",
+        "--num-gpus", "1",
+        "--height", "768",
+        "--width", "1280",
+        "--num-frames", "69",
+        "--num-inference-steps", "2",
+        "--fps", "16",
+        "--guidance-scale", "5.0",
+        "--flow-shift", "5.0",
+        "--prompt", "A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic.",
+        "--negative-prompt", "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
+        "--seed", "1024",
+        "--output-path", "outputs_video/STA_1024/",
+    ]
+    
+    # Run the command
+    subprocess.run(cmd, check=True)
+    
+    # Verify output directory exists
+    output_dir = Path("outputs_video/STA_1024/")
+    assert output_dir.exists(), f"Output directory {output_dir} does not exist"
+    
+    # Verify that video files were generated
+    video_files = list(output_dir.glob("*.mp4"))
+    assert len(video_files) > 0, "No video files were generated"
+    
+    # Verify the video file properties
+    for video_file in video_files:
+        assert video_file.stat().st_size > 0, f"Video file {video_file} is empty"
+
+if __name__ == "__main__":
+    test_inference()

fastvideo/v1/tests/modal/pr_test.py

@@ -0,0 +1,98 @@
+import modal
+
+app = modal.App()
+
+import os
+
+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")
+    .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"})
+    .run_commands("/bin/bash -c 'source $HOME/.local/bin/env && source /opt/venv/bin/activate && cd /FastVideo && uv pip install -e .[test]'")
+)
+
+@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
+    
+    os.chdir("/FastVideo")
+    
+    command = """
+    source /opt/venv/bin/activate && 
+    pytest ./fastvideo/v1/tests/encoders -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_vae_tests():
+    """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 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 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)

fastvideo/v1/tests/ssim/test_inference_similarity.py

@@ -78,7 +78,7 @@ I2V_MODEL_TO_PARAMS = {
 
 TEST_PROMPTS = [
     "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
-    "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature."
+    # "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature."
 ]
 
 I2V_TEST_PROMPTS = [
@@ -328,5 +328,5 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     if not success:
         logger.error("Failed to write SSIM results to file")
 
-    min_acceptable_ssim = 0.97
+    min_acceptable_ssim = 0.95
     assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}"

fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json

@@ -0,0 +1 @@
+{"grad_norm":0.478515625,"_runtime":95.727033597,"_wandb":{"runtime":95},"_step":5,"validation_videos_50_steps":{"videos":[{"_type":"video-file","sha256":"42a1c311521a9d460db788713be1cbf2db767494e02619b43be5bf3eed8381d8","size":158632,"path":"media/videos/validation_videos_50_steps_0_42a1c311521a9d460db7.mp4"},{"path":"media/videos/validation_videos_50_steps_0_818505095b4b5e8b7f51.mp4","_type":"video-file","sha256":"818505095b4b5e8b7f511012d45f04d151ce3344bc058fc0f3225a414a851e4a","size":147825},{"sha256":"fc334ba9ed5e66c8527ee3b408e3be2d76167fef03588bf2840f4a0792f2fe34","size":136933,"path":"media/videos/validation_videos_50_steps_0_fc334ba9ed5e66c8527e.mp4","_type":"video-file"},{"size":201797,"path":"media/videos/validation_videos_50_steps_0_ccd98f6f907635d266a7.mp4","_type":"video-file","sha256":"ccd98f6f907635d266a74783688e7ecf1dac752d79d72d69eab9ef0e3f7413eb"},{"_type":"video-file","sha256":"ca79f40a0aed38f676f12779b349ce40e9e3fb7f36c578f49a20854c70508fb4","size":147114,"path":"media/videos/validation_videos_50_steps_0_ca79f40a0aed38f676f1.mp4"},{"size":175104,"path":"media/videos/validation_videos_50_steps_0_32c9b33ff920c17e5881.mp4","_type":"video-file","sha256":"32c9b33ff920c17e588133d7a27aa400ff3dc529b01ed4f16ac4d6bb2afa0f00"},{"sha256":"2cf520bfb93401c914e93c87ef791c2f12a4e043b95dfdc98115c930e11dfe67","size":139655,"path":"media/videos/validation_videos_50_steps_0_2cf520bfb93401c914e9.mp4","_type":"video-file"},{"_type":"video-file","sha256":"1d73aba17ce582c7aef4af4d64079e3e9d3df205634eff453446bdaf2340b214","size":149028,"path":"media/videos/validation_videos_50_steps_0_1d73aba17ce582c7aef4.mp4"}],"captions":false,"_type":"videos","count":8},"train_loss":0.08922439813613892,"_timestamp":1.750202051751466e+09,"avg_step_time":0.7536672964692116,"step_time":0.4742048177868128,"learning_rate":1e-05,"vsa_sparsity":0.05}

fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py

@@ -0,0 +1,136 @@
+import os
+import sys
+import subprocess
+from pathlib import Path
+import json
+from huggingface_hub import snapshot_download
+
+# Import the training pipeline
+sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
+from fastvideo.v1.training.wan_training_pipeline import main
+from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.v1.utils import FlexibleArgumentParser
+
+wandb_name = "test_training_loss_VSA"
+reference_wandb_summary_file = "fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json"
+
+NUM_NODES = "1"
+NUM_GPUS_PER_NODE = "1"
+
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+
+def run_worker():
+    """Worker function that will be run on each GPU"""
+    # Create and populate args
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    
+    # Set the arguments as they are in finetune_v1_test.sh
+    args = parser.parse_args([
+        "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        "--inference_mode", "False",
+        "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        "--cache_dir", "/home/.cache",
+        "--data_path", "data/mini_dataset_i2v_VSA/combined_parquet_dataset",
+        "--validation_prompt_dir", "data/mini_dataset_i2v_VSA/validation_parquet_dataset",
+        "--train_batch_size", "1",
+        "--num_latent_t", "4",
+        "--num_gpus", "1",
+        "--sp_size", "1",
+        "--tp_size", "1",
+        "--hsdp_replicate_dim", "1",
+        "--hsdp_shard_dim", "1",
+        "--train_sp_batch_size", "1",
+        "--dataloader_num_workers", "4",
+        "--gradient_accumulation_steps", "1",
+        "--max_train_steps", "5",
+        "--learning_rate", "1e-5",
+        "--mixed_precision", "bf16",
+        "--checkpointing_steps", "30",
+        "--validation_steps", "10",
+        "--validation_sampling_steps", "50",
+        "--log_validation",
+        "--checkpoints_total_limit", "3",
+        "--allow_tf32",
+        "--ema_start_step", "0",
+        "--cfg", "0.0",
+        "--output_dir", "data/wan_finetune_test_VSA",
+        "--tracker_project_name", "wan_finetune_ci_VSA",
+        "--wandb_run_name", wandb_name,
+        "--num_height", "384",
+        "--num_width", "512",
+        "--num_frames", "13",
+        "--flow_shift", "3",
+        "--validation_guidance_scale", "1.0",
+        "--num_euler_timesteps", "50",
+        "--weight_decay", "0.01",
+        "--dit_precision", "fp32",
+        "--max_grad_norm", "1.0",
+        "--VSA_decay_rate", "0.01",
+        "--VSA_decay_interval_steps", "1",
+        "--VSA_sparsity", "0.9"
+    ])
+    
+    # Call the main training function
+    main(args)
+
+def test_distributed_training():
+    """Test the distributed training setup"""
+    os.environ["WANDB_MODE"] = "online"
+
+    data_dir = Path("data/mini_dataset_i2v_VSA")
+    
+    if not data_dir.exists():
+        print(f"Downloading test dataset to {data_dir}...")
+        snapshot_download(
+            repo_id="BrianChen1129/mini_dataset_i2v_VSA",
+            local_dir=str(data_dir),
+            repo_type="dataset",
+            local_dir_use_symlinks=False
+        )
+
+    # Get the current file path
+    current_file = Path(__file__).resolve()
+    
+    # Run torchrun command
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE,
+        str(current_file)
+    ]
+    
+    process = subprocess.run(cmd, check=True)
+
+    summary_file = 'wandb/latest-run/files/wandb-summary.json'
+
+    reference_wandb_summary = json.load(open(reference_wandb_summary_file))
+    wandb_summary = json.load(open(summary_file))
+
+    fields_and_thresholds = {
+        'avg_step_time': 1.0,
+        'grad_norm': 0.1,
+        'step_time': 0.5,
+        'train_loss': 0.001
+    }
+
+    failures = []
+    for field, threshold in fields_and_thresholds.items():
+        ref_value = reference_wandb_summary[field]
+        current_value = wandb_summary[field]
+        diff = abs(ref_value - current_value)
+        print(f"INFO: {field}, diff: {diff}, threshold: {threshold}, reference: {ref_value}, current: {current_value}")
+        if diff > threshold:
+            failures.append(f"FAILED: {field} difference {diff} exceeds threshold of {threshold} (reference: {ref_value}, current: {current_value})")
+
+    if failures:
+        raise AssertionError("\n".join(failures))
+
+if __name__ == "__main__":
+    if os.environ.get("LOCAL_RANK") is not None:
+        # We're being run by torchrun
+        run_worker()
+    else:
+        # We're being run directly
+        test_distributed_training()

fastvideo/v1/tests/training/Vanilla/test_training_loss.py

@@ -18,7 +18,7 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.v1.utils import FlexibleArgumentParser
 
 wandb_name = "test_training_loss"
-reference_wandb_summary_file = "fastvideo/v1/tests/training/reference_wandb_summary.json"
+reference_wandb_summary_file = "fastvideo/v1/tests/training/Vanilla/reference_wandb_summary.json"
 
 NUM_NODES = "1"
 NUM_GPUS_PER_NODE = "4"
@@ -81,7 +81,7 @@ def run_worker():
 
 def test_distributed_training():
     """Test the distributed training setup"""
-    os.environ["WANDB_MODE"] = "offline"
+    os.environ["WANDB_MODE"] = "online"
 
     data_dir = Path("data/crush-smol_parq")
     
@@ -107,14 +107,14 @@ def test_distributed_training():
     
     process = subprocess.run(cmd, check=True)
 
-    summary_file = "fastvideo/v1/tests/training/reference_wandb_summary.json"
+    summary_file = 'wandb/latest-run/files/wandb-summary.json'
 
     reference_wandb_summary = json.load(open(reference_wandb_summary_file))
     wandb_summary = json.load(open(summary_file))
 
     fields_and_thresholds = {
         'avg_step_time': 1.0,
-        'grad_norm': 0.1,
+        'grad_norm': 0.2,
         'step_time': 0.5,
         'train_loss': 0.001
     }

fastvideo/v1/training/training_pipeline.py

@@ -2,37 +2,44 @@
 import gc
 import math
 import os
-import traceback
+import time
 from abc import ABC, abstractmethod
+from collections import deque
 from typing import Any, Dict, Iterator, List
 
 import imageio
 import numpy as np
 import torch
-import torch.distributed as dist
 import torchvision
+from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.optimization import get_scheduler
 from einops import rearrange
 from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm.auto import tqdm
 
+import fastvideo.v1.envs as envs
+from fastvideo.v1.attention.backends.video_sparse_attn import (
+    VideoSparseAttentionMetadata)
 from fastvideo.v1.configs.sample import SamplingParam
 from fastvideo.v1.dataset import build_parquet_map_style_dataloader
-from fastvideo.v1.distributed import (get_sp_group, get_torch_device,
-                                      get_world_group)
+from fastvideo.v1.distributed import (cleanup_dist_env_and_memory, get_sp_group,
+                                      get_torch_device, get_world_group)
 from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines import ComposedPipelineBase
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch,
+                                    TrainingBatch)
 from fastvideo.v1.training.training_utils import (
-    compute_density_for_timestep_sampling, get_sigmas, normalize_dit_input)
+    clip_grad_norm_while_handling_failing_dtensor_cases,
+    compute_density_for_timestep_sampling, get_sigmas, load_checkpoint,
+    normalize_dit_input, save_checkpoint, shard_latents_across_sp)
+from fastvideo.v1.utils import is_vsa_available, set_random_seed
 
 import wandb  # isort: skip
 
-logger = init_logger(__name__)
+vsa_available = is_vsa_available()
 
-# Note: if checking with float32, cannot use flash-attn.
-GRADIENT_CHECK_DTYPE = torch.bfloat16
+logger = init_logger(__name__)
 
 
 class TrainingPipeline(ComposedPipelineBase, ABC):
@@ -53,6 +60,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
 
     def initialize_training_pipeline(self, training_args: TrainingArgs):
         logger.info("Initializing training pipeline...")
+        self.training_args = training_args
         self.device = get_torch_device()
         world_group = get_world_group()
         self.world_size = world_group.world_size
@@ -130,17 +138,372 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         raise NotImplementedError(
             "Training pipelines must implement this method")
 
-    @abstractmethod
-    def train_one_step(self, transformer, model_type, optimizer, lr_scheduler,
-                       loader, noise_scheduler, noise_random_generator,
-                       gradient_accumulation_steps, sp_size,
-                       precondition_outputs, max_grad_norm, weighting_scheme,
-                       logit_mean, logit_std, mode_scale):
-        """
-        Train one step of the model.
-        """
-        raise NotImplementedError(
-            "Training pipeline must implement this method")
+    def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch:
+        self.transformer.requires_grad_(True)
+        self.transformer.train()
+        self.optimizer.zero_grad()
+        training_batch.total_loss = 0.0
+        return training_batch
+
+    def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.train_loader_iter is not None
+        assert self.train_dataloader is not None
+
+        batch = next(self.train_loader_iter, None)  # type: ignore
+        if batch is None:
+            self.current_epoch += 1
+            logger.info("Starting epoch %s", self.current_epoch)
+            # Reset iterator for next epoch
+            self.train_loader_iter = iter(self.train_dataloader)
+            # Get first batch of new epoch
+            batch = next(self.train_loader_iter)
+
+        latents, encoder_hidden_states, encoder_attention_mask, infos = batch
+        training_batch.latents = latents.to(get_torch_device(),
+                                            dtype=torch.bfloat16)
+        training_batch.encoder_hidden_states = encoder_hidden_states.to(
+            get_torch_device(), dtype=torch.bfloat16)
+        training_batch.encoder_attention_mask = encoder_attention_mask.to(
+            get_torch_device(), dtype=torch.bfloat16)
+        training_batch.info = infos
+
+        return training_batch
+
+    def _normalize_dit_input(self,
+                             training_batch: TrainingBatch) -> TrainingBatch:
+        # TODO(will): support other models
+        training_batch.latents = normalize_dit_input('wan',
+                                                     training_batch.latents)
+        return training_batch
+
+    def _prepare_dit_inputs(self,
+                            training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.training_args is not None
+        assert training_batch.latents is not None
+        assert training_batch.encoder_hidden_states is not None
+        assert training_batch.encoder_attention_mask is not None
+        assert self.noise_random_generator is not None
+
+        batch_size = training_batch.latents.shape[0]
+        noise = torch.randn_like(training_batch.latents)
+        u = compute_density_for_timestep_sampling(
+            weighting_scheme=self.training_args.weighting_scheme,
+            batch_size=batch_size,
+            generator=self.noise_random_generator,
+            logit_mean=self.training_args.logit_mean,
+            logit_std=self.training_args.logit_std,
+            mode_scale=self.training_args.mode_scale,
+        )
+        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
+        timesteps = self.noise_scheduler.timesteps[indices].to(
+            device=training_batch.latents.device)
+        if self.training_args.sp_size > 1:
+            # Make sure that the timesteps are the same across all sp processes.
+            sp_group = get_sp_group()
+            sp_group.broadcast(timesteps, src=0)
+        sigmas = get_sigmas(
+            self.noise_scheduler,
+            training_batch.latents.device,
+            timesteps,
+            n_dim=training_batch.latents.ndim,
+            dtype=training_batch.latents.dtype,
+        )
+        noisy_model_input = (1.0 -
+                             sigmas) * training_batch.latents + sigmas * noise
+
+        training_batch.noisy_model_input = noisy_model_input
+        training_batch.timesteps = timesteps
+        training_batch.sigmas = sigmas
+        training_batch.noise = noise
+
+        return training_batch
+
+    def _build_attention_metadata(
+            self, training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.training_args is not None
+        latents = training_batch.latents
+        assert latents is not None
+        assert training_batch.timesteps is not None
+        patch_size = self.training_args.pipeline_config.dit_config.patch_size
+        current_vsa_sparsity = training_batch.current_vsa_sparsity
+
+        if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
+
+            dit_seq_shape = [
+                latents.shape[2] // patch_size[0],
+                latents.shape[3] // patch_size[1],
+                latents.shape[4] // patch_size[2]
+            ]
+            training_batch.attn_metadata = VideoSparseAttentionMetadata(
+                current_timestep=training_batch.timesteps,
+                dit_seq_shape=dit_seq_shape,
+                VSA_sparsity=current_vsa_sparsity)
+        else:
+            training_batch.attn_metadata = None
+
+        return training_batch
+
+    def _transformer_forward_and_compute_loss(
+            self, training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.transformer is not None
+        assert self.training_args is not None
+        assert training_batch.noisy_model_input is not None
+        assert training_batch.encoder_hidden_states is not None
+        assert training_batch.encoder_attention_mask is not None
+        assert training_batch.timesteps is not None
+        # assert training_batch.attn_metadata is not None
+        assert training_batch.latents is not None
+        assert training_batch.noise is not None
+        assert training_batch.sigmas is not None
+
+        if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
+            assert training_batch.attn_metadata is not None
+        else:
+            assert training_batch.attn_metadata is None
+
+        input_kwargs = {
+            "hidden_states":
+            training_batch.noisy_model_input,
+            "encoder_hidden_states":
+            training_batch.encoder_hidden_states,
+            "timestep":
+            training_batch.timesteps.to(get_torch_device(),
+                                        dtype=torch.bfloat16),
+            "encoder_attention_mask":
+            training_batch.encoder_attention_mask,
+            "return_dict":
+            False,
+        }
+        # if 'hunyuan' in self.training_args.model_type:
+        #     input_kwargs["guidance"] = torch.tensor(
+        #         [1000.0],
+        #         device=training_batch.noisy_model_input.device,
+        #         dtype=torch.bfloat16)
+
+        with set_forward_context(
+                current_timestep=training_batch.current_timestep,
+                attn_metadata=training_batch.attn_metadata):
+            model_pred = self.transformer(**input_kwargs)
+        if self.training_args.precondition_outputs:
+            model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
+        target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents
+
+        loss = (torch.mean((model_pred.float() - target.float())**2) /
+                self.training_args.gradient_accumulation_steps)
+
+        loss.backward()
+        avg_loss = loss.detach().clone()
+        # logger.info(f"rank: {self.rank}, avg_loss: {avg_loss.item()}",
+        #             local_main_process_only=False)
+        world_group = get_world_group()
+        world_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG)
+        training_batch.total_loss += avg_loss.item()
+
+        return training_batch
+
+    def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.training_args is not None
+        max_grad_norm = self.training_args.max_grad_norm
+
+        # TODO(will): perhaps move this into transformer api so that we can do
+        # the following:
+        # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
+        if max_grad_norm is not None:
+            model_parts = [self.transformer]
+            grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
+                [p for m in model_parts for p in m.parameters()],
+                max_grad_norm,
+                foreach=None,
+            )
+            assert grad_norm is not float('nan') or grad_norm is not float(
+                'inf')
+            grad_norm = grad_norm.item() if grad_norm is not None else 0.0
+        else:
+            grad_norm = 0.0
+        training_batch.grad_norm = grad_norm
+        return training_batch
+
+    def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.training_args is not None
+
+        training_batch = self._prepare_training(training_batch)
+
+        for _ in range(self.training_args.gradient_accumulation_steps):
+            training_batch = self._get_next_batch(training_batch)
+
+            # Shard latents across sp groups
+            training_batch.latents = shard_latents_across_sp(
+                training_batch.latents,
+                num_latent_t=self.training_args.num_latent_t)
+
+            # Normalize DIT input
+            training_batch = self._normalize_dit_input(training_batch)
+            training_batch = self._prepare_dit_inputs(training_batch)
+            training_batch = self._build_attention_metadata(training_batch)
+            training_batch = self._transformer_forward_and_compute_loss(
+                training_batch)
+
+        training_batch = self._clip_grad_norm(training_batch)
+
+        self.optimizer.step()
+        self.lr_scheduler.step()
+
+        training_batch.total_loss = training_batch.total_loss
+        training_batch.grad_norm = training_batch.grad_norm
+        return training_batch
+
+    def _resume_from_checkpoint(self) -> None:
+        assert self.training_args is not None
+        logger.info("Loading checkpoint from %s",
+                    self.training_args.resume_from_checkpoint)
+        resumed_step = load_checkpoint(
+            self.transformer, self.global_rank,
+            self.training_args.resume_from_checkpoint, self.optimizer,
+            self.train_dataloader, self.lr_scheduler,
+            self.noise_random_generator)
+        if resumed_step > 0:
+            self.init_steps = resumed_step
+            logger.info("Successfully resumed from step %s", resumed_step)
+        else:
+            logger.warning("Failed to load checkpoint, starting from step 0")
+            self.init_steps = 0
+
+    def train(self) -> None:
+        assert self.training_args is not None
+
+        # Set random seeds for deterministic training
+        assert self.training_args.seed is not None, "seed must be set"
+        seed = self.training_args.seed
+        set_random_seed(seed)
+
+        self.noise_random_generator = torch.Generator(
+            device="cpu").manual_seed(seed)
+
+        logger.info("Initialized random seeds with seed: %s", seed)
+
+        self.noise_scheduler = FlowMatchEulerDiscreteScheduler()
+
+        if self.training_args.resume_from_checkpoint:
+            self._resume_from_checkpoint()
+
+        self.train_loader_iter = iter(self.train_dataloader)
+
+        step_times: deque[float] = deque(maxlen=100)
+
+        self._log_training_info()
+        self._log_validation(self.transformer, self.training_args, 1)
+
+        # Train!
+        progress_bar = tqdm(
+            range(0, self.training_args.max_train_steps),
+            initial=self.init_steps,
+            desc="Steps",
+            # Only show the progress bar once on each machine.
+            disable=self.local_rank > 0,
+        )
+        for step in range(self.init_steps + 1,
+                          self.training_args.max_train_steps + 1):
+            start_time = time.perf_counter()
+            if vsa_available:
+                vsa_sparsity = self.training_args.VSA_sparsity
+                vsa_decay_rate = self.training_args.VSA_decay_rate
+                vsa_decay_interval_steps = self.training_args.VSA_decay_interval_steps
+                current_decay_times = min(step // vsa_decay_interval_steps,
+                                          vsa_sparsity // vsa_decay_rate)
+                current_vsa_sparsity = current_decay_times * vsa_decay_rate
+            else:
+                current_vsa_sparsity = 0.0
+
+            training_batch = TrainingBatch()
+            training_batch.current_timestep = step
+            training_batch.current_vsa_sparsity = current_vsa_sparsity
+            training_batch = self.train_one_step(training_batch)
+
+            loss = training_batch.total_loss
+            grad_norm = training_batch.grad_norm
+
+            step_time = time.perf_counter() - start_time
+            step_times.append(step_time)
+            avg_step_time = sum(step_times) / len(step_times)
+
+            progress_bar.set_postfix({
+                "loss": f"{loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
+            })
+            progress_bar.update(1)
+            if self.global_rank == 0:
+                wandb.log(
+                    {
+                        "train_loss": loss,
+                        "learning_rate": self.lr_scheduler.get_last_lr()[0],
+                        "step_time": step_time,
+                        "avg_step_time": avg_step_time,
+                        "grad_norm": grad_norm,
+                        "vsa_sparsity": current_vsa_sparsity,
+                    },
+                    step=step,
+                )
+            if step % self.training_args.checkpointing_steps == 0:
+                save_checkpoint(self.transformer, self.global_rank,
+                                self.training_args.output_dir, step,
+                                self.optimizer, self.train_dataloader,
+                                self.lr_scheduler, self.noise_random_generator)
+                self.transformer.train()
+                self.sp_group.barrier()
+            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
+                self._log_validation(self.transformer, self.training_args, step)
+                gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+                logger.info("GPU memory usage after validation: %s MB",
+                            gpu_memory_usage)
+
+        wandb.finish()
+        save_checkpoint(self.transformer, self.global_rank,
+                        self.training_args.output_dir,
+                        self.training_args.max_train_steps, self.optimizer,
+                        self.train_dataloader, self.lr_scheduler,
+                        self.noise_random_generator)
+
+        if get_sp_group():
+            cleanup_dist_env_and_memory()
+
+    def _log_training_info(self) -> None:
+        assert self.training_args is not None
+        assert self.training_args.sp_size is not None
+        assert self.training_args.gradient_accumulation_steps is not None
+        total_batch_size = (self.world_size *
+                            self.training_args.gradient_accumulation_steps /
+                            self.training_args.sp_size *
+                            self.training_args.train_sp_batch_size)
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %s", len(self.train_dataset))
+        logger.info("  Dataloader size = %s", len(self.train_dataloader))
+        logger.info("  Num Epochs = %s", self.num_train_epochs)
+        logger.info("  Resume training from step %s",
+                    self.init_steps)  # type: ignore
+        logger.info("  Instantaneous batch size per device = %s",
+                    self.training_args.train_batch_size)
+        logger.info(
+            "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
+            total_batch_size)
+        logger.info("  Gradient Accumulation steps = %s",
+                    self.training_args.gradient_accumulation_steps)
+        logger.info("  Total optimization steps = %s",
+                    self.training_args.max_train_steps)
+        logger.info(
+            "  Total training parameters per FSDP shard = %s B",
+            sum(p.numel()
+                for p in self.transformer.parameters() if p.requires_grad) /
+            1e9)
+        # print dtype
+        logger.info("  Master weight dtype: %s",
+                    self.transformer.parameters().__next__().dtype)
+
+        gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+        logger.info("GPU memory usage before train_one_step: %s MB",
+                    gpu_memory_usage)
+        logger.info("VSA validation sparsity: %s",
+                    self.training_args.VSA_sparsity)
 
     @torch.no_grad()
     def _log_validation(self, transformer, training_args, global_step) -> None:
@@ -294,259 +657,3 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         transformer.train()
         gc.collect()
         torch.cuda.empty_cache()
-
-    def gradient_check_parameters(self,
-                                  transformer,
-                                  latents,
-                                  encoder_hidden_states,
-                                  encoder_attention_mask,
-                                  timesteps,
-                                  target,
-                                  eps=5e-2,
-                                  max_params_to_check=2000) -> float:
-        """
-        Verify gradients using finite differences for FSDP models with GRADIENT_CHECK_DTYPE.
-        Uses standard tolerances for GRADIENT_CHECK_DTYPE precision.
-        """
-        assert self.training_args is not None
-        # Move all inputs to CPU and clear GPU memory
-        inputs_cpu = {
-            'latents': latents.cpu(),
-            'encoder_hidden_states': encoder_hidden_states.cpu(),
-            'encoder_attention_mask': encoder_attention_mask.cpu(),
-            'timesteps': timesteps.cpu(),
-            'target': target.cpu()
-        }
-        del latents, encoder_hidden_states, encoder_attention_mask, timesteps, target
-        torch.cuda.empty_cache()
-
-        def compute_loss() -> torch.Tensor:
-            assert self.training_args is not None
-            # Move inputs to GPU, compute loss, cleanup
-            inputs_gpu = {
-                k:
-                v.to(get_torch_device(),
-                     dtype=GRADIENT_CHECK_DTYPE
-                     if k != 'encoder_attention_mask' else None)
-                for k, v in inputs_cpu.items()
-            }
-
-            # Use GRADIENT_CHECK_DTYPE for more accurate gradient checking
-            # with torch.autocast(enabled=False, device_type="cuda"):
-            with torch.autocast("cuda", dtype=GRADIENT_CHECK_DTYPE):
-                with set_forward_context(
-                        current_timestep=inputs_gpu['timesteps'],
-                        attn_metadata=None):
-                    model_pred = transformer(
-                        hidden_states=inputs_gpu['latents'],
-                        encoder_hidden_states=inputs_gpu[
-                            'encoder_hidden_states'],
-                        timestep=inputs_gpu['timesteps'],
-                        encoder_attention_mask=inputs_gpu[
-                            'encoder_attention_mask'],
-                        return_dict=False)[0]
-
-                if self.training_args.precondition_outputs:
-                    sigmas = get_sigmas(self.noise_scheduler,
-                                        inputs_gpu['latents'].device,
-                                        inputs_gpu['timesteps'],
-                                        n_dim=inputs_gpu['latents'].ndim,
-                                        dtype=inputs_gpu['latents'].dtype)
-                    model_pred = inputs_gpu['latents'] - model_pred * sigmas
-                    target_adjusted = inputs_gpu['target']
-                else:
-                    target_adjusted = inputs_gpu['target']
-
-                loss = torch.mean((model_pred - target_adjusted)**2)
-
-            # Cleanup and return
-            loss_cpu = loss.cpu()
-            del inputs_gpu, model_pred, target_adjusted
-            if 'sigmas' in locals():
-                del sigmas
-            torch.cuda.empty_cache()
-            return loss_cpu.to(get_torch_device())
-
-        try:
-            # Get analytical gradients
-            transformer.zero_grad()
-            analytical_loss = compute_loss()
-            analytical_loss.backward()
-
-            # Check gradients for selected parameters
-            absolute_errors: list[float] = []
-            param_count = 0
-
-            rank = dist.get_rank()
-            sp_group = get_sp_group()
-            for name, param in transformer.named_parameters():
-                sp_group.barrier()
-                # skip scale_shift_table because it is not sharded
-                if 'scale_shift_table' in name:
-                    continue
-                if isinstance(param.grad, torch.distributed.tensor.DTensor):
-                    full_grad = param.grad.full_tensor()
-                    distributed = True
-                else:
-                    full_grad = param.grad
-                    distributed = False
-                    continue
-                if not (param.requires_grad and param.grad is not None
-                        and param_count < max_params_to_check
-                        and full_grad.abs().max() > 5e-4):
-                    continue
-                if not distributed and rank != 0:
-                    continue
-
-                # Get local parameter and gradient tensors
-                local_param = param._local_tensor if hasattr(
-                    param, '_local_tensor') else param
-                local_grad = param.grad._local_tensor if hasattr(
-                    param.grad, '_local_tensor') else param.grad
-
-                # Find first significant gradient element
-                flat_param = local_param.data.view(-1)
-                flat_grad = local_grad.view(-1)
-                check_idx = next((i for i in range(min(10, flat_param.numel()))
-                                  if abs(flat_grad[i]) > 1e-4), 0)
-
-                # Store original values
-                orig_value = flat_param[check_idx].item()
-                analytical_grad = flat_grad[check_idx].item()
-
-                # Compute numerical gradient
-                for delta in [eps, -eps]:
-                    with torch.no_grad():
-                        # only have a single rank modify the parameter
-                        # because we are using FSDP
-                        if rank == 0:
-                            flat_param[check_idx] = orig_value + delta
-                        loss = compute_loss()
-                        if delta > 0:
-                            loss_plus = loss.item()
-                        else:
-                            loss_minus = loss.item()
-
-                # Restore parameter and compute error
-                with torch.no_grad():
-                    flat_param[check_idx] = orig_value
-
-                numerical_grad = (loss_plus - loss_minus) / (2 * eps)
-                abs_error = abs(analytical_grad - numerical_grad)
-                rel_error = abs_error / max(abs(analytical_grad),
-                                            abs(numerical_grad), 1e-3)
-                absolute_errors.append(abs_error)
-
-                if self.global_rank == 0:
-                    logger.info(
-                        "%s[%s]: analytical=%.5f, numerical=%.5f, abs_error=%.2e, rel_error=%.2f%%",
-                        name, check_idx, analytical_grad, numerical_grad,
-                        abs_error, rel_error * 100)
-
-                # param_count += 1
-
-            # Compute and log statistics
-            if rank == 0 and absolute_errors:
-                min_err, max_err, mean_err = min(absolute_errors), max(
-                    absolute_errors
-                ), sum(absolute_errors) / len(absolute_errors)
-                logger.info("Gradient check stats: min=%s, max=%s, mean=%s",
-                            min_err, max_err, mean_err)
-
-                wandb.log({
-                    "grad_check/min_abs_error": min_err,
-                    "grad_check/max_abs_error": max_err,
-                    "grad_check/mean_abs_error": mean_err,
-                    "grad_check/analytical_loss": analytical_loss.item(),
-                })
-                return max_err
-
-            return float('inf')
-
-        except Exception as e:
-            logger.error("Gradient check failed: %s", e)
-            traceback.print_exc()
-            return float('inf')
-
-    def setup_gradient_check(self, args, loader_iter, noise_scheduler,
-                             noise_random_generator) -> float | None:
-        """
-        Setup and perform gradient check on a fresh batch.
-        Args:
-            args: Training arguments
-            loader_iter: Data loader iterator
-            noise_scheduler: Noise scheduler for diffusion
-            noise_random_generator: Random number generator for noise
-        Returns:
-            float or None: Maximum gradient error or None if check is disabled/fails
-        """
-        assert self.training_args is not None
-
-        try:
-            # Get a fresh batch and process it exactly like train_one_step
-            check_latents, check_encoder_hidden_states, check_encoder_attention_mask, check_infos = next(
-                loader_iter)
-
-            # Process exactly like in train_one_step but use GRADIENT_CHECK_DTYPE
-            check_latents = check_latents.to(get_torch_device(),
-                                             dtype=GRADIENT_CHECK_DTYPE)
-            check_encoder_hidden_states = check_encoder_hidden_states.to(
-                get_torch_device(), dtype=GRADIENT_CHECK_DTYPE)
-            check_latents = normalize_dit_input("wan", check_latents)
-            batch_size = check_latents.shape[0]
-            check_noise = torch.randn_like(check_latents)
-
-            check_u = compute_density_for_timestep_sampling(
-                weighting_scheme=args.weighting_scheme,
-                batch_size=batch_size,
-                generator=noise_random_generator,
-                logit_mean=args.logit_mean,
-                logit_std=args.logit_std,
-                mode_scale=args.mode_scale,
-            )
-            check_indices = (check_u *
-                             noise_scheduler.config.num_train_timesteps).long()
-            check_timesteps = noise_scheduler.timesteps[check_indices].to(
-                device=check_latents.device)
-
-            check_sigmas = get_sigmas(
-                noise_scheduler,
-                check_latents.device,
-                check_timesteps,
-                n_dim=check_latents.ndim,
-                dtype=check_latents.dtype,
-            )
-            check_noisy_model_input = (
-                1.0 - check_sigmas) * check_latents + check_sigmas * check_noise
-
-            # Compute target exactly like train_one_step
-            if args.precondition_outputs:
-                check_target = check_latents
-            else:
-                check_target = check_noise - check_latents
-
-            # Perform gradient check with the exact same inputs as training
-            max_grad_error = self.gradient_check_parameters(
-                transformer=self.transformer,
-                latents=
-                check_noisy_model_input,  # Use noisy input like in training
-                encoder_hidden_states=check_encoder_hidden_states,
-                encoder_attention_mask=check_encoder_attention_mask,
-                timesteps=check_timesteps,
-                target=check_target,
-                max_params_to_check=100  # Check more parameters
-            )
-
-            if max_grad_error > 5e-2:
-                logger.error("❌ Large gradient error detected: %s",
-                             max_grad_error)
-            else:
-                logger.info("✅ Gradient check passed: max error %s",
-                            max_grad_error)
-
-            return max_grad_error
-
-        except Exception as e:
-            logger.error("Gradient check setup failed: %s", e)
-            traceback.print_exc()
-            return None

fastvideo/v1/training/wan_training_pipeline.py

@@ -1,45 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
-import importlib.util
-import random
 import sys
-import time
-from collections import deque
 from copy import deepcopy
 
-import numpy as np
-import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
-from tqdm.auto import tqdm
-
-import fastvideo.v1.envs as envs
-from fastvideo.v1.attention.backends.video_sparse_attn import (
-    VideoSparseAttentionMetadata)
-from fastvideo.v1.distributed import (cleanup_dist_env_and_memory, get_sp_group,
-                                      get_torch_device, get_world_group)
 from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import (
     FlowUniPCMultistepScheduler)
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline
 from fastvideo.v1.training.training_pipeline import TrainingPipeline
-from fastvideo.v1.training.training_utils import (
-    clip_grad_norm_while_handling_failing_dtensor_cases,
-    compute_density_for_timestep_sampling, get_sigmas, load_checkpoint,
-    normalize_dit_input, save_checkpoint, shard_latents_across_sp)
+from fastvideo.v1.utils import is_vsa_available
 
-import wandb  # isort: skip
-
-vsa_available = False
-if importlib.util.find_spec("vsa") is not None:
-    vsa_available = True
+vsa_available = is_vsa_available()
 
 logger = init_logger(__name__)
 
-# Manual gradient checking flag - set to True to enable gradient verification
-ENABLE_GRADIENT_CHECK = False
-
 
 class WanTrainingPipeline(TrainingPipeline):
     """
@@ -74,321 +48,6 @@ class WanTrainingPipeline(TrainingPipeline):
 
         self.validation_pipeline = validation_pipeline
 
-    def train_one_step(  # type: ignore[override]
-        self,
-        transformer,
-        model_type,
-        optimizer,
-        lr_scheduler,
-        loader_iter,
-        noise_scheduler,
-        noise_random_generator,
-        gradient_accumulation_steps,
-        sp_size,
-        precondition_outputs,
-        max_grad_norm,
-        weighting_scheme,
-        logit_mean,
-        logit_std,
-        mode_scale,
-        patch_size,
-        current_vsa_sparsity,
-    ) -> tuple[float, float]:
-        assert self.training_args is not None
-        self.modules["transformer"].requires_grad_(True)
-        self.modules["transformer"].train()
-
-        total_loss = 0.0
-        optimizer.zero_grad()
-
-        for _ in range(gradient_accumulation_steps):
-            # Get next batch, handling epoch boundaries gracefully
-            batch = next(self.train_loader_iter, None)  # type: ignore
-            if batch is None:
-                self.current_epoch += 1
-                logger.info("Starting epoch %s", self.current_epoch)
-                # Reset iterator for next epoch
-                self.train_loader_iter = iter(self.train_dataloader)
-                # Get first batch of new epoch
-                batch = next(self.train_loader_iter)
-
-            latents, encoder_hidden_states, encoder_attention_mask, _ = batch
-
-            latents = latents.to(get_torch_device(), dtype=torch.bfloat16)
-            encoder_hidden_states = encoder_hidden_states.to(
-                get_torch_device(), dtype=torch.bfloat16)
-            latents = shard_latents_across_sp(
-                latents, num_latent_t=self.training_args.num_latent_t)
-
-            dit_seq_shape = [
-                latents.shape[2] // patch_size[0],
-                latents.shape[3] // patch_size[1],
-                latents.shape[4] // patch_size[2]
-            ]
-
-            latents = normalize_dit_input(model_type, latents)
-            batch_size = latents.shape[0]
-            noise = torch.randn_like(latents)
-            u = compute_density_for_timestep_sampling(
-                weighting_scheme=weighting_scheme,
-                batch_size=batch_size,
-                generator=noise_random_generator,
-                logit_mean=logit_mean,
-                logit_std=logit_std,
-                mode_scale=mode_scale,
-            )
-            indices = (u * noise_scheduler.config.num_train_timesteps).long()
-            timesteps = noise_scheduler.timesteps[indices].to(
-                device=latents.device)
-            if sp_size > 1:
-                # Make sure that the timesteps are the same across all sp processes.
-                sp_group = get_sp_group()
-                sp_group.broadcast(timesteps, src=0)
-            sigmas = get_sigmas(
-                noise_scheduler,
-                latents.device,
-                timesteps,
-                n_dim=latents.ndim,
-                dtype=latents.dtype,
-            )
-            noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                input_kwargs = {
-                    "hidden_states": noisy_model_input,
-                    "encoder_hidden_states": encoder_hidden_states,
-                    "timestep": timesteps,
-                    "encoder_attention_mask": encoder_attention_mask,  # B, L
-                    "return_dict": False,
-                }
-                if 'hunyuan' in model_type:
-                    input_kwargs["guidance"] = torch.tensor(
-                        [1000.0],
-                        device=noisy_model_input.device,
-                        dtype=torch.bfloat16)
-
-                if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
-                    attn_metadata = VideoSparseAttentionMetadata(
-                        current_timestep=timesteps,
-                        dit_seq_shape=dit_seq_shape,
-                        VSA_sparsity=current_vsa_sparsity)
-                else:
-                    attn_metadata = None
-
-                with set_forward_context(current_timestep=timesteps,
-                                         attn_metadata=attn_metadata):
-                    model_pred = transformer(**input_kwargs)
-
-                if precondition_outputs:
-                    model_pred = noisy_model_input - model_pred * sigmas
-                target = latents if precondition_outputs else noise - latents
-
-                loss = (torch.mean((model_pred.float() - target.float())**2) /
-                        gradient_accumulation_steps)
-
-            loss.backward()
-
-            avg_loss = loss.detach().clone()
-            # logger.info(f"rank: {self.rank}, avg_loss: {avg_loss.item()}",
-            #             local_main_process_only=False)
-            world_group = get_world_group()
-            world_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG)
-            total_loss += avg_loss.item()
-
-        # TODO(will): perhaps move this into transformer api so that we can do
-        # the following:
-        # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
-        if max_grad_norm is not None:
-            model_parts = [self.transformer]
-            grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
-                [p for m in model_parts for p in m.parameters()],
-                max_grad_norm,
-                foreach=None,
-            )
-            grad_norm = grad_norm.item() if grad_norm is not None else 0.0
-        else:
-            grad_norm = 0.0
-
-        optimizer.step()
-        lr_scheduler.step()
-        return total_loss, grad_norm
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ):
-        assert self.training_args is not None
-
-        # Set random seeds for deterministic training
-        seed = self.training_args.seed if self.training_args.seed is not None else 42
-
-        random.seed(seed)
-        np.random.seed(seed)
-        torch.manual_seed(seed)
-        torch.cuda.manual_seed_all(seed)
-
-        noise_random_generator = torch.Generator(device="cpu").manual_seed(seed)
-
-        logger.info("Initialized random seeds with seed: %s", seed)
-
-        noise_scheduler = FlowMatchEulerDiscreteScheduler()
-
-        # Train!
-        assert self.training_args.sp_size is not None
-        assert self.training_args.gradient_accumulation_steps is not None
-        total_batch_size = (self.world_size *
-                            self.training_args.gradient_accumulation_steps /
-                            self.training_args.sp_size *
-                            self.training_args.train_sp_batch_size)
-        logger.info("***** Running training *****")
-        logger.info("  Num examples = %s", len(self.train_dataset))
-        logger.info("  Dataloader size = %s", len(self.train_dataloader))
-        logger.info("  Num Epochs = %s", self.num_train_epochs)
-        logger.info("  Resume training from step %s",
-                    self.init_steps)  # type: ignore
-        logger.info("  Instantaneous batch size per device = %s",
-                    self.training_args.train_batch_size)
-        logger.info(
-            "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
-            total_batch_size)
-        logger.info("  Gradient Accumulation steps = %s",
-                    self.training_args.gradient_accumulation_steps)
-        logger.info("  Total optimization steps = %s",
-                    self.training_args.max_train_steps)
-        logger.info(
-            "  Total training parameters per FSDP shard = %s B",
-            sum(p.numel()
-                for p in self.transformer.parameters() if p.requires_grad) /
-            1e9)
-        # print dtype
-        logger.info("  Master weight dtype: %s",
-                    self.transformer.parameters().__next__().dtype)
-
-        if self.training_args.resume_from_checkpoint:
-            logger.info("Loading checkpoint from %s",
-                        self.training_args.resume_from_checkpoint)
-            resumed_step = load_checkpoint(
-                self.transformer, self.global_rank,
-                self.training_args.resume_from_checkpoint, self.optimizer,
-                self.train_dataloader, self.lr_scheduler,
-                noise_random_generator)
-            if resumed_step > 0:
-                self.init_steps = resumed_step
-                logger.info("Successfully resumed from step %s", resumed_step)
-            else:
-                logger.warning(
-                    "Failed to load checkpoint, starting from step 0")
-                self.init_steps = 0
-
-        progress_bar = tqdm(
-            range(0, self.training_args.max_train_steps),
-            initial=self.init_steps,
-            desc="Steps",
-            # Only show the progress bar once on each machine.
-            disable=self.local_rank > 0,
-        )
-
-        self.train_loader_iter = iter(self.train_dataloader)
-
-        step_times: deque[float] = deque(maxlen=100)
-
-        # TODO(will): fix this
-        # for i in range(self.init_steps):
-        #     next(loader_iter)
-        # get gpu memory usage
-        gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
-        logger.info("GPU memory usage before train_one_step: %s MB",
-                    gpu_memory_usage)
-        logger.info("VSA validation sparsity: %s",
-                    self.training_args.VSA_sparsity)
-        self._log_validation(self.transformer, self.training_args, 1)
-        if vsa_available:
-            vsa_sparsity = self.training_args.VSA_sparsity
-            vsa_decay_rate = self.training_args.VSA_decay_rate
-            vsa_decay_interval_steps = self.training_args.VSA_decay_interval_steps
-
-        for step in range(self.init_steps + 1,
-                          self.training_args.max_train_steps + 1):
-            start_time = time.perf_counter()
-            if vsa_available:
-                current_decay_times = min(step // vsa_decay_interval_steps,
-                                          vsa_sparsity // vsa_decay_rate)
-                current_vsa_sparsity = current_decay_times * vsa_decay_rate
-            else:
-                current_vsa_sparsity = 0.0
-            loss, grad_norm = self.train_one_step(
-                self.transformer,
-                # args.model_type,
-                "wan",
-                self.optimizer,
-                self.lr_scheduler,
-                self.train_loader_iter,
-                noise_scheduler,
-                noise_random_generator,
-                self.training_args.gradient_accumulation_steps,
-                self.training_args.sp_size,
-                self.training_args.precondition_outputs,
-                self.training_args.max_grad_norm,
-                self.training_args.weighting_scheme,
-                self.training_args.logit_mean,
-                self.training_args.logit_std,
-                self.training_args.mode_scale,
-                self.training_args.pipeline_config.dit_config.patch_size,
-                current_vsa_sparsity,
-            )
-
-            step_time = time.perf_counter() - start_time
-            step_times.append(step_time)
-            avg_step_time = sum(step_times) / len(step_times)
-
-            # Manual gradient checking - only at first step
-            if step == 1 and ENABLE_GRADIENT_CHECK:
-                logger.info("Performing gradient check at step %s", step)
-                self.setup_gradient_check(args, self.train_loader_iter,
-                                          noise_scheduler,
-                                          noise_random_generator)
-
-            progress_bar.set_postfix({
-                "loss": f"{loss:.4f}",
-                "step_time": f"{step_time:.2f}s",
-                "grad_norm": grad_norm,
-            })
-            progress_bar.update(1)
-            if self.global_rank == 0:
-                wandb.log(
-                    {
-                        "train_loss": loss,
-                        "learning_rate": self.lr_scheduler.get_last_lr()[0],
-                        "step_time": step_time,
-                        "avg_step_time": avg_step_time,
-                        "grad_norm": grad_norm,
-                        "vsa_sparsity": current_vsa_sparsity,
-                    },
-                    step=step,
-                )
-            if step % self.training_args.checkpointing_steps == 0:
-                save_checkpoint(self.transformer, self.global_rank,
-                                self.training_args.output_dir, step,
-                                self.optimizer, self.train_dataloader,
-                                self.lr_scheduler, noise_random_generator)
-                self.transformer.train()
-                self.sp_group.barrier()
-            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
-                self._log_validation(self.transformer, self.training_args, step)
-                gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
-                logger.info("GPU memory usage after validation: %s MB",
-                            gpu_memory_usage)
-
-        wandb.finish()
-        save_checkpoint(self.transformer, self.global_rank,
-                        self.training_args.output_dir,
-                        self.training_args.max_train_steps, self.optimizer,
-                        self.train_dataloader, self.lr_scheduler,
-                        noise_random_generator)
-
-        if get_sp_group():
-            cleanup_dist_env_and_memory()
-
 
 def main(args) -> None:
     logger.info("Starting training pipeline...")
@@ -396,7 +55,7 @@ def main(args) -> None:
     pipeline = WanTrainingPipeline.from_pretrained(
         args.pretrained_model_name_or_path, args=args)
     args = pipeline.training_args
-    pipeline.forward(None, args)
+    pipeline.train()
     logger.info("Training pipeline done")
 
 

fastvideo/v1/utils.py

@@ -5,6 +5,7 @@ import argparse
 import ctypes
 import hashlib
 import importlib
+import importlib.util
 import inspect
 import json
 import math
@@ -16,7 +17,7 @@ import tempfile
 import threading
 import traceback
 from dataclasses import dataclass, fields, is_dataclass
-from functools import partial, wraps
+from functools import lru_cache, partial, wraps
 from typing import (Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar,
                     Union, cast)
 
@@ -732,3 +733,66 @@ def get_compute_dtype() -> torch.dtype:
     else:
         state = get_mixed_precision_state()
         return state.param_dtype
+
+
+def dict_to_3d_list(
+    mask_strategy: Optional[Dict[str, Any]] = None,
+    t_max: Optional[int] = None,
+    l_max: Optional[int] = None,
+    h_max: Optional[int] = None,
+) -> List[List[List[Optional[torch.Tensor]]]]:
+    """
+    Convert a dictionary of mask indices to a 3D list of tensors.
+    Args:
+        mask_strategy: keys are "t_l_h", values are torch.Tensor masks.
+        t_max, l_max, h_max: if provided (all three), force the output shape to (t_max, l_max, h_max).
+                            If all three are None, infer shape from the data.
+    """
+    # Case 1: no data, but fixed shape requested
+    if mask_strategy is None:
+        assert t_max is not None and l_max is not None and h_max is not None, (
+            "If mask_strategy is None, you must provide t_max, l_max, and h_max"
+        )
+        return [[[None for _ in range(h_max)] for _ in range(l_max)]
+                for _ in range(t_max)]
+
+    # Parse all keys into integer tuples
+    indices = [tuple(map(int, key.split("_"))) for key in mask_strategy]
+
+    # Decide on dimensions
+    if t_max is None and l_max is None and h_max is None:
+        # fully dynamic: infer from data
+        max_timesteps_idx = max(t for t, _, _ in indices) + 1
+        max_layer_idx = max(l for _, l, _ in indices) + 1  # noqa: E741
+        max_head_idx = max(h for _, _, h in indices) + 1
+    else:
+        # require all three to be provided
+        assert t_max is not None and l_max is not None and h_max is not None, (
+            "Either supply none of (t_max, l_max, h_max) to infer dimensions, "
+            "or supply all three to fix the shape.")
+        max_timesteps_idx = t_max
+        max_layer_idx = l_max
+        max_head_idx = h_max
+
+    # Preallocate
+    result = [[[None for _ in range(max_head_idx)]
+               for _ in range(max_layer_idx)] for _ in range(max_timesteps_idx)]
+
+    # Fill in, skipping any out-of-bounds entries
+    for key, value in mask_strategy.items():
+        t, l, h = map(int, key.split("_"))  # noqa: E741
+        if 0 <= t < max_timesteps_idx and 0 <= l < max_layer_idx and 0 <= h < max_head_idx:
+            result[t][l][h] = value
+        # else: silently ignore any key that doesn’t fit
+
+    return result
+
+
+def set_random_seed(seed: int) -> None:
+    from fastvideo.v1.platforms import current_platform
+    current_platform.seed_everything(seed)
+
+
+@lru_cache(maxsize=1)
+def is_vsa_available() -> bool:
+    return importlib.util.find_spec("vsa") is not None

pyproject.toml

@@ -20,7 +20,7 @@ dependencies = [
     # Machine Learning & Transformers
     "transformers>=4.46.1", "tokenizers>=0.20.1", "sentencepiece==0.2.0",
     "timm==1.0.11", "peft==0.13.2", "diffusers>=0.33.1", "bitsandbytes",
-    "torch==2.6.0", "torchvision",
+    "torch==2.7.1", "torchvision",
 
     # Acceleration & Optimization
     "accelerate==1.0.1",
@@ -49,6 +49,9 @@ dependencies = [
     "av",
 ]
 
+[tool.uv]
+extra-index-url = ["https://download.pytorch.org/whl/cu128"]
+
 [project.optional-dependencies]
 
 # flash-attn: pip install flash-attn==2.7.4.post1 --no-cache-dir --no-build-isolation 
@@ -135,4 +138,4 @@ column_limit = 80
 [tool.isort]
 line_length = 80
 use_parentheses = true
-skip_gitignore = true
+skip_gitignore = true

scripts/inference/v1_inference_wan_STA.sh

@@ -20,5 +20,5 @@ fastvideo generate \
     --guidance-scale 5.0 \
     --prompt "A beautiful woman in a red dress walking down a street" \
     --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
-    --seed 12345 \
+    --seed 1024 \
     --output-path outputs_video/