ckpt

.buildkite/pipeline.yml

@@ -2,25 +2,26 @@ env:
   IMAGE_VERSION: "py3.12-latest"
 
 steps:
-  - block: "Start Build"
+  - label: "pre-commit"
-    blocked_state: "running"
+    command: ".buildkite/scripts/pre_commit.sh"
-    prompt: "Approve build?"
+    agents:
+      queue: "default"
+    env:
+      - BUILDKITE_CLEAN_CHECKOUT=true
+
+  - wait
 
   - 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/models/loader/**"
                 - "fastvideo/v1/tests/encoders/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
               config:
                 command: "timeout 30m .buildkite/scripts/pr_test.sh"
                 label: "Encoder Tests"
@@ -31,8 +32,10 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/v1/models/vaes/**"
-                - "fastvideo/v1/models/loaders/**"
+                - "fastvideo/v1/models/loader/**"
                 - "fastvideo/v1/tests/vaes/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
               config:
                 command: "timeout 30m .buildkite/scripts/pr_test.sh"
                 label: "VAE Tests"
@@ -43,10 +46,12 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/v1/models/dits/**"
-                - "fastvideo/v1/models/loaders/**"
+                - "fastvideo/v1/models/loader/**"
                 - "fastvideo/v1/tests/transformers/**"
                 - "fastvideo/v1/layers/**"
                 - "fastvideo/v1/attention/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
               config:
                 command: "timeout 30m .buildkite/scripts/pr_test.sh"
                 label: "Transformer Tests"
@@ -55,7 +60,8 @@ steps:
                   - TEST_TYPE=transformer
                 agents:
                   queue: "default"
-            - path: "fastvideo/v1/**/*.py"
+            - path: 
+                - "fastvideo/v1/**/*.py"
               config:
                 command: "timeout 60m .buildkite/scripts/pr_test.sh"
                 label: "SSIM Tests"
@@ -64,3 +70,79 @@ steps:
                   - TEST_TYPE=ssim
                 agents:
                   queue: "default"
+            - path:
+                - "fastvideo/v1/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                label: "Training Tests"
+                env:
+                  - BUILDKITE_CLEAN_CHECKOUT=true
+                  - TEST_TYPE=training
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/**"
+                - "csrc/attn/vsa/**"
+                - "csrc/attn/tk/**"
+                - "csrc/attn/setup_vsa.py"
+                - "csrc/attn/config_vsa.py"
+                - "csrc/attn/vsa.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                label: "Training Tests VSA"
+                env:
+                  - BUILDKITE_CLEAN_CHECKOUT=true
+                  - TEST_TYPE=training_vsa
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/**"
+                - "csrc/attn/st_attn/**"
+                - "csrc/attn/setup_sta.py"
+                - "csrc/attn/config_sta.py"
+                - "csrc/attn/st_attn.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                label: "Inference Tests STA"
+                env:
+                  - BUILDKITE_CLEAN_CHECKOUT=true
+                  - TEST_TYPE=inference_sta
+                agents:
+                  queue: "default"
+            - path:
+                - "csrc/attn/st_attn/**"
+                - "csrc/attn/setup_sta.py"
+                - "csrc/attn/config_sta.py"
+                - "csrc/attn/st_attn.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                label: "Precision Tests STA"
+                env:
+                  - BUILDKITE_CLEAN_CHECKOUT=true
+                  - TEST_TYPE=precision_sta
+                agents:
+                  queue: "default"
+            - path:
+                - "csrc/attn/vsa/**"
+                - "csrc/attn/tk/**"
+                - "csrc/attn/setup_vsa.py"
+                - "csrc/attn/config_vsa.py"
+                - "csrc/attn/vsa.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                label: "Precision Tests VSA"
+                env:
+                  - BUILDKITE_CLEAN_CHECKOUT=true
+                  - TEST_TYPE=precision_vsa
+                agents:
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -31,6 +31,10 @@ log "Setting up Modal authentication from Buildkite secrets..."
 MODAL_TOKEN_ID=$(buildkite-agent secret get modal_token_id)
 MODAL_TOKEN_SECRET=$(buildkite-agent secret get modal_token_secret)
 
+WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
+
+WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
+
 if [ -n "$MODAL_TOKEN_ID" ] && [ -n "$MODAL_TOKEN_SECRET" ]; then
     log "Retrieved Modal credentials from Buildkite secrets"
     python3 -m modal token set --token-id "$MODAL_TOKEN_ID" --token-secret "$MODAL_TOKEN_SECRET" --profile buildkite-ci --activate --verify
@@ -54,22 +58,44 @@ if [ -z "${TEST_TYPE:-}" ]; then
 fi
 log "Test type: $TEST_TYPE"
 
+MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT IMAGE_VERSION=$IMAGE_VERSION"
+
 case "$TEST_TYPE" in
     "encoder")
         log "Running encoder tests..."
-        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
+        MODAL_COMMAND="$MODAL_ENV 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"
+        MODAL_COMMAND="$MODAL_ENV 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"
+        MODAL_COMMAND="$MODAL_ENV 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"
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_ssim_tests"
+        ;;
+    "training")
+        log "Running training tests..."
+        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests"
+        ;;
+    "training_vsa")
+        log "Running training VSA tests..."
+        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests_VSA"
+        ;;
+    "inference_sta")
+        log "Running inference STA tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_STA"
+        ;;
+    "precision_sta")
+        log "Running precision STA tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_STA"
+        ;;
+    "precision_vsa")
+        log "Running precision VSA tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_VSA"
         ;;
     *)
         log "Error: Unknown test type: $TEST_TYPE"

.buildkite/scripts/pre_commit.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+set -uo pipefail
+
+log() {
+  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
+}
+
+log "=== Starting pre-commit checks ==="
+
+cd "$(dirname "$0")/../.."
+PROJECT_ROOT=$(pwd)
+log "Project root: $PROJECT_ROOT"
+
+if ! python3 -m pre_commit --version &> /dev/null; then
+    log "pre-commit not found, installing..."
+    python3 -m pip install --user pre-commit==4.0.1
+    
+    if ! python3 -m pre_commit --version &> /dev/null; then
+        log "Error: Failed to install pre-commit."
+        exit 1
+    fi
+fi
+
+log "Pre-commit version: $(python3 -m pre_commit --version)"
+
+log "Installing/updating pre-commit hooks..."
+python3 -m pre_commit install --install-hooks
+
+log "Running pre-commit checks on all files..."
+python3 -m pre_commit run --all-files
+PRE_COMMIT_EXIT_CODE=$?
+
+if [ $PRE_COMMIT_EXIT_CODE -eq 0 ]; then
+    log "Pre-commit checks completed successfully"
+else
+    log "Error: Pre-commit checks failed with exit code: $PRE_COMMIT_EXIT_CODE"
+fi
+
+log "=== Pre-commit checks completed with exit code: $PRE_COMMIT_EXIT_CODE ==="
+exit $PRE_COMMIT_EXIT_CODE 

.github/workflows/pr-test.yml

@@ -122,17 +122,17 @@ jobs:
             # Actual tests
             encoder-test:
               - 'fastvideo/v1/models/encoders/**'
-              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/models/loader/**'
               - 'fastvideo/v1/tests/encoders/**'
               - *common-paths
             vae-test:
               - 'fastvideo/v1/models/vaes/**'
-              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/models/loader/**'
               - 'fastvideo/v1/tests/vaes/**'
               - *common-paths
             transformer-test:
               - 'fastvideo/v1/models/dits/**'
-              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/models/loader/**'
               - 'fastvideo/v1/tests/transformers/**'
               - 'fastvideo/v1/layers/**'
               - 'fastvideo/v1/attention/**'

examples/inference/basic/basic.py

@@ -10,7 +10,7 @@ def main():
     # attempt to identify the optimal arguments.
     generator = VideoGenerator.from_pretrained(
         "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        # if num_gpus > 1, FastVideo will automatically handle distributed setup
+        # FastVideo will automatically handle distributed setup
         num_gpus=2,
         use_fsdp_inference=True,
         use_cpu_offload=False

examples/training/finetune/wan_i2v_14b_480p/crush_smol/README.md

@@ -0,0 +1,10 @@
+This directory contain e2e examples scripts for finetuning Wan2.1 I2V. 
+
+Execute the following commands from `FastVideo/` to run training:
+
+- Download crush-smol dataset:
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
+- Preprocess the videos and captions into latents:
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
+- Edit the following file and run finetuning:
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`

examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh

@@ -1,7 +1,10 @@
+#!/bin/bash
+
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 
+MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
 VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
 NUM_GPUS=8
@@ -12,7 +15,7 @@ NUM_GPUS=8
 training_args=(
   --tracker_project_name "wan_i2v_finetune"
   --output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
-  --max_train_steps 5000
+  --max_train_steps 2000
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 1
@@ -33,8 +36,8 @@ parallel_args=(
 
 # Model arguments
 model_args=(
-  --model_path "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+  --pretrained_model_name_or_path $MODEL_PATH
 )
 
 # Dataset arguments
@@ -56,7 +59,7 @@ validation_args=(
 optimizer_args=(
   --learning_rate 1e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 6000
+  --checkpointing_steps 1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )
@@ -66,7 +69,7 @@ miscellaneous_args=(
   --inference_mode False
   --allow_tf32
   --checkpoints_total_limit 3
-  --cfg 0.0
+  --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"
   --not_apply_cfg_solver
   --dit_precision "fp32"

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm

@@ -1,5 +1,5 @@
 #!/bin/bash
-#SBATCH --job-name=FV_2N_14B
+#SBATCH --job-name=i2v
 #SBATCH --partition=main
 #SBATCH --qos=hao
 #SBATCH --nodes=4
@@ -9,8 +9,8 @@
 #SBATCH --cpus-per-task=128
 #SBATCH --nodelist=fs-mbz-gpu-[100-850]
 #SBATCH --mem=1440G
-#SBATCH --output=4n_i2v/4n_i2v_%j.out
+#SBATCH --output=i2v_output/i2v_%j.out
-#SBATCH --error=4n_i2v/4n_i2v_%j.err
+#SBATCH --error=i2v_output/i2v_%j.err
 #SBATCH --exclusive
 set -e -x
 
@@ -30,7 +30,9 @@ nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
 export MASTER_ADDR=${nodes[0]}
 export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
 export TOKENIZERS_PARALLELISM=false
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 
 echo "MASTER_ADDR: $MASTER_ADDR"
 echo "NODE_RANK: $NODE_RANK"
@@ -38,60 +40,91 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 
-export WANDB_BASE_URL="https://api.wandb.ai"
+MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-export WANDB_MODE=online
+DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
-
-DATA_DIR=data/crush-smol_processed_i2v/combined_parquet_dataset
-VALIDATION_DIR=data/crush-smol_processed_i2v/validation_parquet_dataset
-NUM_GPUS=8
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
 # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+
+# Training arguments
+training_args=(
+  --tracker_project_name wan_i2v_finetune
+  --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"
+  --max_train_steps=2000
+  --train_batch_size=2
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps=1
+  --num_latent_t 8
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size $NUM_GPUS
+  --tp_size $NUM_GPUS
+  --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
+  --hsdp_shard_dim $NUM_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 10
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_preprocessed_path "$VALIDATION_DIR"
+  --validation_steps 100
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "1.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate=1e-5
+  --mixed_precision="bf16"
+  --checkpointing_steps=1000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --allow_tf32
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
 srun torchrun \
 --nnodes $SLURM_JOB_NUM_NODES \
 --nproc_per_node $NUM_GPUS \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py\
+    fastvideo/v1/training/wan_i2v_training_pipeline.py \
-    --model_path Wan-AI/Wan2.1-I2V-14B-480P-Diffusers \
+    "${parallel_args[@]}" \
-    --inference_mode False\
+    "${model_args[@]}" \
-    --pretrained_model_name_or_path Wan-AI/Wan2.1-I2V-14B-480P-Diffusers \
+    "${dataset_args[@]}" \
-    --cache_dir "/home/ray/.cache"\
+    "${training_args[@]}" \
-    --data_path "$DATA_DIR"\
+    "${optimizer_args[@]}" \
-    --validation_preprocessed_path "$VALIDATION_DIR"\
+    "${validation_args[@]}" \
-    --train_batch_size=2\
+    "${miscellaneous_args[@]}"
-    --num_latent_t 8 \
-    --num_gpus $NUM_GPUS \
-    --sp_size $NUM_GPUS \
-    --tp_size $NUM_GPUS \
-    --hsdp_replicate_dim $SLURM_JOB_NUM_NODES \
-    --hsdp_shard_dim $NUM_GPUS \
-    --train_sp_batch_size 1\
-    --dataloader_num_workers 10\
-    --gradient_accumulation_steps=1\
-    --max_train_steps=10000 \
-    --learning_rate=5e-5\
-    --mixed_precision="bf16"\
-    --checkpointing_steps=11000 \
-    --validation_steps 10\
-    --validation_sampling_steps "40" \
-    --log_validation \
-    --checkpoints_total_limit 3\
-    --allow_tf32\
-    --ema_start_step 0\
-    --cfg 0.0\
-    --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"\
-    --tracker_project_name wan_i2v_finetune \
-    --num_height 480 \
-    --num_width 832 \
-    --num_frames 77 \
-    --validation_guidance_scale "1.0" \
-    --num_euler_timesteps 50 \
-    --multi_phased_distill_schedule "4000-1" \
-    --weight_decay 1e-4 \
-    --not_apply_cfg_solver \
-    --dit_precision "fp32" \
-    --max_grad_norm 1.0

examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh

@@ -1,4 +1,5 @@
-# export WANDB_MODE="offline"
+#!/bin/bash
+
 GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 MODEL_TYPE="wan"

examples/training/finetune/wan_t2v_1_3b/crush_smol/README.md

@@ -0,0 +1,10 @@
+This directory contain e2e examples scripts for finetuning Wan2.1 T2v. 
+
+Execute the following commands from `FastVideo/` to run training:
+
+- Download crush-smol dataset:
+`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/download_dataset.sh`
+- Preprocess the videos and captions into latents:
+`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/preprocess_wan_data_t2v.sh`
+- Edit the following file and run finetuning:
+`bash examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh`

examples/training/finetune/wan_t2v_1_3b/crush_smol/download_dataset.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh

@@ -1,3 +1,5 @@
+#!/bin/bash
+
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
@@ -25,11 +27,11 @@ training_args=(
 
 # Parallel arguments
 parallel_args=(
-  --num_gpus $NUM_GPUS \
+  --num_gpus $NUM_GPUS 
-  --sp_size $NUM_GPUS \
+  --sp_size $NUM_GPUS 
-  --tp_size $NUM_GPUS \
+  --tp_size $NUM_GPUS
-  --hsdp_replicate_dim 1 \
+  --hsdp_replicate_dim 1
-  --hsdp_shard_dim $NUM_GPUS \
+  --hsdp_shard_dim $NUM_GPUS
 )
 
 # Model arguments
@@ -67,7 +69,7 @@ miscellaneous_args=(
   --inference_mode False
   --allow_tf32
   --checkpoints_total_limit 3
-  --cfg 0.0
+  --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"
   --not_apply_cfg_solver
   --dit_precision "fp32"

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm

@@ -7,10 +7,10 @@
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
-#SBATCH --nodelist=fs-mbz-gpu-[400-550]
+#SBATCH --nodelist=fs-mbz-gpu-[100-850]
 #SBATCH --mem=1440G
-#SBATCH --output=1n_t2v/1n_t2v_%j.out
+#SBATCH --output=t2v_output/t2v_%j.out
-#SBATCH --error=1n_t2v/1n_t2v_%j.err
+#SBATCH --error=t2v_output/t2v_%j.err
 #SBATCH --exclusive
 set -e -x
 
@@ -30,69 +30,98 @@ nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
 export MASTER_ADDR=${nodes[0]}
 export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
 export TOKENIZERS_PARALLELISM=false
-export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 
 echo "MASTER_ADDR: $MASTER_ADDR"
 echo "NODE_RANK: $NODE_RANK"
 
 # Configs
 NUM_GPUS=8
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
 VALIDATION_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/"
-NUM_GPUS=8
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+# Training arguments
+training_args=(
+  --tracker_project_name wan_t2v_finetune
+  --output_dir="outputs/wan_t2v_finetune"
+  --max_train_steps=1000
+  --train_batch_size=4
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps=1
+  --num_latent_t 8
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 4
+  --tp_size 4
+  --hsdp_replicate_dim 2
+  --hsdp_shard_dim 4
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 10
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_preprocessed_path "$VALIDATION_DIR"
+  --validation_steps 100
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "1.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate=5e-5
+  --mixed_precision="bf16"
+  --checkpointing_steps=500
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --allow_tf32
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
 srun torchrun \
 --nnodes $SLURM_JOB_NUM_NODES \
 --nproc_per_node $NUM_GPUS \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_training_pipeline.py\
+    fastvideo/v1/training/wan_training_pipeline.py \
-    --model_path $MODEL_PATH \
+    "${parallel_args[@]}" \
-    --inference_mode False\
+    "${model_args[@]}" \
-    --pretrained_model_name_or_path $MODEL_PATH \
+    "${dataset_args[@]}" \
-    --cache_dir "/home/ray/.cache"\
+    "${training_args[@]}" \
-    --data_path "$DATA_DIR"\
+    "${optimizer_args[@]}" \
-    --validation_preprocessed_path "$VALIDATION_DIR"\
+    "${validation_args[@]}" \
-    --train_batch_size=4\
+    "${miscellaneous_args[@]}"
-    --num_latent_t 8 \
-    --num_gpus $NUM_GPUS \
-    --sp_size 4 \
-    --tp_size 4 \
-    --hsdp_replicate_dim 2 \
-    --hsdp_shard_dim 4 \
-    --train_sp_batch_size 1\
-    --dataloader_num_workers 10\
-    --gradient_accumulation_steps=1\
-    --max_train_steps=10000 \
-    --learning_rate=5e-5\
-    --mixed_precision="bf16"\
-    --checkpointing_steps=11000 \
-    --validation_steps 50\
-    --validation_sampling_steps "50" \
-    --log_validation \
-    --checkpoints_total_limit 3\
-    --allow_tf32\
-    --ema_start_step 0\
-    --cfg 0.0\
-    --output_dir="$DATA_DIR/outputs/wan_t2v_finetune_2n"\
-    --tracker_project_name wan_t2v_finetune \
-    --num_height 480 \
-    --num_width 832 \
-    --num_frames 77 \
-    --validation_guidance_scale "1.0" \
-    --num_euler_timesteps 50 \
-    --multi_phased_distill_schedule "4000-1" \
-    --weight_decay 1e-4 \
-    --not_apply_cfg_solver \
-    --dit_precision "fp32" \
-    --max_grad_norm 1.0

examples/training/finetune/wan_t2v_1_3b/crush_smol/overfit.json

@@ -1,13 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": "examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}

examples/training/finetune/wan_t2v_1_3b/crush_smol/preprocess_wan_data_t2v.sh

@@ -1,4 +1,5 @@
-# export WANDB_MODE="offline"
+#!/bin/bash
+
 GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"

fastvideo/v1/configs/models/base.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field, fields
-from typing import Any, Dict
+from typing import Any, Dict, List, Tuple
 
 from fastvideo.v1.logger import init_logger
 
@@ -12,7 +12,9 @@ logger = init_logger(__name__)
 # 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
 @dataclass
 class ArchConfig:
-    pass
+    stacked_params_mapping: List[Tuple[str, str, str]] = field(
+        default_factory=list
+    )  # mapping from huggingface weight names to custom names
 
 
 @dataclass

fastvideo/v1/configs/models/dits/stepvideo.py

@@ -5,13 +5,11 @@ from typing import List, Optional, Tuple, Union
 from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
 
 
-def is_blocks(n: str, m) -> bool:
-    return "blocks" in n and str.isdigit(n.split(".")[-1])
-
-
 @dataclass
 class StepVideoArchConfig(DiTArchConfig):
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda:
+        [lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit()])
 
     _param_names_mapping: dict = field(
         default_factory=lambda: {

fastvideo/v1/configs/models/encoders/base.py

@@ -32,8 +32,11 @@ class TextEncoderArchConfig(EncoderArchConfig):
     output_past: bool = True
     scalable_attention: bool = True
     tie_word_embeddings: bool = False
-
+    stacked_params_mapping: List[Tuple[str, str, str]] = field(
+        default_factory=list
+    )  # mapping from huggingface weight names to custom names
     tokenizer_kwargs: Dict[str, Any] = field(default_factory=dict)
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [])
 
     def __post_init__(self) -> None:
         self.tokenizer_kwargs = {

fastvideo/v1/configs/models/encoders/clip.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import List, Optional, Tuple
 
 from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                        ImageEncoderConfig,
@@ -8,6 +8,14 @@ from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                        TextEncoderConfig)
 
 
+def _is_transformer_layer(n: str, m) -> bool:
+    return "layers" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m) -> bool:
+    return n.endswith("embeddings")
+
+
 @dataclass
 class CLIPTextArchConfig(TextEncoderArchConfig):
     vocab_size: int = 49408
@@ -27,6 +35,15 @@ class CLIPTextArchConfig(TextEncoderArchConfig):
     bos_token_id: int = 49406
     eos_token_id: int = 49407
     text_len: int = 77
+    stacked_params_mapping: List[Tuple[str, str,
+                                       str]] = field(default_factory=lambda: [
+                                           # (param_name, shard_name, shard_id)
+                                           ("qkv_proj", "q_proj", "q"),
+                                           ("qkv_proj", "k_proj", "k"),
+                                           ("qkv_proj", "v_proj", "v"),
+                                       ])
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings])
 
 
 @dataclass

fastvideo/v1/configs/models/encoders/llama.py

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

fastvideo/v1/configs/models/encoders/t5.py

@@ -1,11 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import List, Optional, Tuple
 
 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                        TextEncoderConfig)
 
 
+def _is_transformer_layer(n: str, m) -> bool:
+    return "block" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m) -> bool:
+    return n.endswith("shared")
+
+
+def _is_final_layernorm(n: str, m) -> bool:
+    return n.endswith("final_layer_norm")
+
+
 @dataclass
 class T5ArchConfig(TextEncoderArchConfig):
     vocab_size: int = 32128
@@ -29,6 +41,16 @@ class T5ArchConfig(TextEncoderArchConfig):
     eos_token_id: int = 1
     classifier_dropout: float = 0.0
     text_len: int = 512
+    stacked_params_mapping: List[Tuple[str, str,
+                                       str]] = field(default_factory=lambda: [
+                                           # (param_name, shard_name, shard_id)
+                                           (".qkv_proj", ".q", "q"),
+                                           (".qkv_proj", ".k", "k"),
+                                           (".qkv_proj", ".v", "v"),
+                                       ])
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda:
+        [_is_transformer_layer, _is_embeddings, _is_final_layernorm])
 
     # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py
     def __post_init__(self):

fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py

@@ -11,7 +11,7 @@ from fastvideo.v1.dataset.parquet_dataset_iterable_style import (
     build_parquet_iterable_style_dataloader)
 from fastvideo.v1.distributed import get_world_rank
 from fastvideo.v1.distributed.parallel_state import (
-    cleanup_dist_env_and_memory, get_torch_device,
+    cleanup_dist_env_and_memory, get_local_torch_device,
     maybe_init_distributed_environment_and_model_parallel)
 from fastvideo.v1.logger import init_logger
 
@@ -148,8 +148,8 @@ def main() -> None:
                 break
 
             # Move data to device
-            latents = latents.to(get_torch_device())
+            latents = latents.to(get_local_torch_device())
-            embeddings = embeddings.to(get_torch_device())
+            embeddings = embeddings.to(get_local_torch_device())
 
             # Calculate actual batch size
             batch_size = latents.size(0)

fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_map_style.py

@@ -8,11 +8,12 @@ import torch
 import torch.distributed as dist
 import torch.distributed.checkpoint as dist_cp
 
+from fastvideo.v1.dataset.dataloader.schema import pyarrow_schema_t2v
 from fastvideo.v1.dataset.parquet_dataset_map_style import (
     build_parquet_map_style_dataloader)
 from fastvideo.v1.distributed import get_world_rank
 from fastvideo.v1.distributed.parallel_state import (
-    cleanup_dist_env_and_memory, get_torch_device,
+    cleanup_dist_env_and_memory, get_local_torch_device,
     maybe_init_distributed_environment_and_model_parallel)
 from fastvideo.v1.logger import init_logger
 
@@ -67,14 +68,18 @@ def main() -> None:
 
     # Create DataLoader with proper settings
     dataset, dataloader = build_parquet_map_style_dataloader(
-        args.path, args.batch_size, args.num_data_workers)
+        args.path,
+        args.batch_size,
+        parquet_schema=pyarrow_schema_t2v,
+        num_data_workers=args.num_data_workers)
     logger.info("Initialized dataloader with %d batches", len(dataloader))
 
     if args.verify_resume:
         # First pass - record latent sums
         first_pass_sums = []
-        for i, (latents, embeddings, masks,
+        for i, batch in enumerate(dataloader):
-                caption_text) in enumerate(dataloader):
+            latents = batch['vae_latent']
+            embeddings = batch['text_embedding']
             latent_sum = latents.sum().item()
             first_pass_sums.append(latent_sum)
             logger.info("Batch %d latent sum: %f", i, latent_sum)
@@ -100,14 +105,18 @@ def main() -> None:
 
         # Recreate dataloader and load state
         dataset, dataloader = build_parquet_map_style_dataloader(
-            args.path, args.batch_size, args.num_data_workers)
+            args.path,
+            args.batch_size,
+            parquet_schema=pyarrow_schema_t2v,
+            num_data_workers=args.num_data_workers)
         load_states = {"dataloader": dataloader}
         dist_cp.load(load_states, checkpoint_id=checkpoint_dir.as_posix())
         logger.info("Rank %d: Loaded dataloader state from %s",
                     get_world_rank(), checkpoint_dir)
 
-        for i, (latents, embeddings, masks,
+        for i, batch in enumerate(dataloader):
-                caption_text) in enumerate(dataloader):
+            latents = batch['vae_latent']
+            embeddings = batch['text_embedding']
             latent_sum = latents.sum().item()
             first_pass_sums.append(latent_sum)
             logger.info("Batch %d latent sum: %f",
@@ -116,11 +125,16 @@ def main() -> None:
                 break
 
         dataset, dataloader = build_parquet_map_style_dataloader(
-            args.path, args.batch_size, args.num_data_workers)
+            args.path,
+            args.batch_size,
+            parquet_schema=pyarrow_schema_t2v,
+            num_data_workers=args.num_data_workers)
 
         # Second pass - verify latent sums match
         second_pass_sums = []
-        for i, (latents, embeddings, masks) in enumerate(dataloader):
+        for i, batch in enumerate(dataloader):
+            latents = batch['vae_latent']
+            embeddings = batch['text_embedding']
             latent_sum = latents.sum().item()
             second_pass_sums.append(latent_sum)
             logger.info("Batch %d latent sum: %f (should match first pass: %f)",
@@ -144,14 +158,15 @@ def main() -> None:
     total_samples = 0
     total_batches = 0
     for _ in range(args.num_epoch):
-        for i, (latents, embeddings, masks,
+        for i, batch in enumerate(dataloader):
-                caption_text) in enumerate(dataloader):
+            latents = batch['vae_latent']
+            embeddings = batch['text_embedding']
             if i >= args.num_batches_per_epoch:
                 break
 
             # Move data to device
-            latents = latents.to(get_torch_device())
+            latents = latents.to(get_local_torch_device())
-            embeddings = embeddings.to(get_torch_device())
+            embeddings = embeddings.to(get_local_torch_device())
 
             # Calculate actual batch size
             batch_size = latents.size(0)

fastvideo/v1/dataset/parquet_dataset_map_style.py

@@ -185,9 +185,6 @@ class LatentsParquetMapStyleDataset(Dataset):
     Note: 
     Using parquet for map style dataset is not efficient, we mainly keep it for backward compatibility and debugging.
     """
-    # Modify this in the future if we want to add more keys, for example, in image to video.
-    keys = [("vae_latent", "latent"), "text_embedding", "clip_feature",
-            "first_frame_latent", "pil_image"]
 
     def __init__(
         self,
@@ -204,10 +201,6 @@ class LatentsParquetMapStyleDataset(Dataset):
         self.path = path
         self.cfg_rate = cfg_rate
         self.parquet_schema = parquet_schema
-        if cfg_rate > 0.0:
-            raise ValueError(
-                "cfg_rate > 0.0 is not supported for now because it will trigger bug when num_data_workers > 0"
-            )
         logger.info("Initializing LatentsParquetMapStyleDataset with path: %s",
                     path)
         self.parquet_files, self.lengths = get_parquet_files_and_length(path)
@@ -243,7 +236,8 @@ class LatentsParquetMapStyleDataset(Dataset):
 
         batch = collate_rows_from_parquet_schema([row_dict],
                                                  self.parquet_schema,
-                                                 self.text_padding_length)
+                                                 self.text_padding_length,
+                                                 cfg_rate=0.0)
         negative_prompt = batch['info_list'][0]['prompt']
         negative_prompt_embedding = batch['text_embedding']
         negative_prompt_attention_mask = batch['text_attention_mask']
@@ -265,11 +259,10 @@ class LatentsParquetMapStyleDataset(Dataset):
             for idx in indices
         ]
 
-        # all_latents, all_embs, all_masks, caption_text, all_extra_latents, all_infos = collate_latents_embs_masks(
+        batch = collate_rows_from_parquet_schema(rows,
-        #     rows, self.text_padding_length, self.keys)
+                                                 self.parquet_schema,
-        # return all_latents, all_embs, all_masks, caption_text, all_extra_latents, all_infos
+                                                 self.text_padding_length,
-        batch = collate_rows_from_parquet_schema(rows, self.parquet_schema,
+                                                 cfg_rate=self.cfg_rate)
-                                                 self.text_padding_length)
         return batch
 
     def __len__(self):

fastvideo/v1/dataset/preprocessing_datasets.py

@@ -484,7 +484,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
         self.text_encoding_stage = TextEncodingStage(
             tokenizer=tokenizer,
             text_max_length=args.text_max_length,
-            cfg_rate=args.cfg)
+            cfg_rate=args.training_cfg_rate)
 
     def _load_raw_data(self) -> List[Dict]:
         """Load raw data from JSON files."""
@@ -493,10 +493,13 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
             folder_anno_pairs = [
                 line.strip().split(",") for line in f if line.strip()
             ]
-        assert len(folder_anno_pairs) == 1, "Only support one folder-annotation pair"
+        assert len(
-        assert len(folder_anno_pairs[0]) == 2, "Folder-annotation pair should have two elements"
+            folder_anno_pairs) == 1, "Only support one folder-annotation pair"
+        assert len(folder_anno_pairs[0]
+                   ) == 2, "Folder-annotation pair should have two elements"
         folder, annotation_file = folder_anno_pairs[0]
 
+        data_items: List[Dict] = []
         with open(annotation_file) as f:
             data_items = json.load(f)
 

fastvideo/v1/dataset/utils.py

@@ -1,12 +1,9 @@
-from typing import Any, Dict, List
+import random
+from typing import Any, Dict, List, cast
 
 import numpy as np
 import torch
 
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
 
 def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
     """
@@ -24,7 +21,7 @@ def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
         return t[:padding_length], torch.ones(padding_length)
 
 
-def get_torch_tensors_from_row_dict(row_dict, keys) -> Dict[str, Any]:
+def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> Dict[str, Any]:
     """
     Get the latents and prompts from a row dictionary.
     """
@@ -42,70 +39,45 @@ def get_torch_tensors_from_row_dict(row_dict, keys) -> Dict[str, Any]:
             if shape is None or bytes is None:
                 raise ValueError(f"Key {key} not found in row_dict")
         else:
-            try:
+            shape = row_dict[f"{key}_shape"]
-                shape = row_dict[f"{key}_shape"]
+            bytes = row_dict[f"{key}_bytes"]
-                bytes = row_dict[f"{key}_bytes"]
-            except KeyError:
-                continue
 
         # TODO (peiyuan): read precision
-        if len(bytes) == 0:
+        if key == 'text_embedding' and random.random() < cfg_rate:
-            return_dict[key] = torch.zeros(0, dtype=torch.bfloat16)
+            data = np.zeros((512, 4096), dtype=np.float32)
         else:
             data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy()
-            data = torch.from_numpy(data)
+        data = torch.from_numpy(data)
-            if len(data.shape) == 3:
+        if len(data.shape) == 3:
-                B, L, D = data.shape
+            B, L, D = data.shape
-                assert B == 1, "Batch size must be 1"
+            assert B == 1, "Batch size must be 1"
-                data = data.squeeze(0)
+            data = data.squeeze(0)
-            return_dict[key] = data
+        return_dict[key] = data
     return return_dict
 
 
 def collate_latents_embs_masks(
-    batch_to_process, text_padding_length, keys
+        batch_to_process,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str], Dict[str, Any],
+        text_padding_length,
-           List[Dict[str, Any]]]:
+        keys,
+        cfg_rate=0.0
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str]]:
     # Initialize tensors to hold padded embeddings and masks
     all_latents = []
     all_embs = []
     all_masks = []
-    all_clip_features = []
-    all_first_frame_latents = []
-    all_pil_images = []
-    all_infos = []
     caption_text = []
     # Process each row individually
     for i, row in enumerate(batch_to_process):
-        # Get info from row
-        info_keys = [
-            "caption", "file_name", "media_type", "width", "height",
-            "num_frames", "duration_sec", "fps"
-        ]
-        info = {}
-        for key in info_keys:
-            if key in row:
-                info[key] = row[key]
-            else:
-                info[key] = ""
-        info["prompt"] = info["caption"]
-
         # Get tensors from row
-        data = get_torch_tensors_from_row_dict(row, keys)
+        data = get_torch_tensors_from_row_dict(row, keys, cfg_rate)
         latents, emb = data["vae_latent"], data["text_embedding"]
-        clip_feature = data.get("clip_feature", None)
-        first_frame_latent = data.get("first_frame_latent", None)
-        pil_image = data.get("pil_image", None)
 
         padded_emb, mask = pad(emb, text_padding_length)
         # Store in batch tensors
         all_latents.append(latents)
         all_embs.append(padded_emb)
         all_masks.append(mask)
-        all_clip_features.append(clip_feature)
-        all_first_frame_latents.append(first_frame_latent)
-        all_pil_images.append(pil_image)
-        all_infos.append(info)
         # TODO(py): remove this once we fix preprocess
         try:
             caption_text.append(row["prompt"])
@@ -116,17 +88,14 @@ def collate_latents_embs_masks(
     all_latents = torch.stack(all_latents)
     all_embs = torch.stack(all_embs)
     all_masks = torch.stack(all_masks)
-    all_extra_latents = {
-        "clip_feature": torch.stack(all_clip_features),
-        "first_frame_latent": torch.stack(all_first_frame_latents),
-        "pil_image": all_pil_images,
-    }
 
-    return all_latents, all_embs, all_masks, caption_text, all_extra_latents, all_infos
+    return all_latents, all_embs, all_masks, caption_text
 
 
-def collate_rows_from_parquet_schema(rows, parquet_schema,
+def collate_rows_from_parquet_schema(rows,
-                                     text_padding_length) -> Dict[str, Any]:
+                                     parquet_schema,
+                                     text_padding_length,
+                                     cfg_rate=0.0) -> Dict[str, Any]:
     """
     Collate rows from parquet files based on the provided schema.
     Dynamically processes tensor fields based on schema and returns batched data.
@@ -139,10 +108,10 @@ def collate_rows_from_parquet_schema(rows, parquet_schema,
         Dict containing batched tensors and metadata
     """
     if not rows:
-        return {}
+        return cast(Dict[str, Any], {})
 
     # Initialize containers for different data types
-    batch_data = {}
+    batch_data: Dict[str, Any] = {}
 
     # Get tensor and metadata field names from schema (fields ending with '_bytes')
     tensor_fields = []
@@ -159,7 +128,7 @@ def collate_rows_from_parquet_schema(rows, parquet_schema,
             # Only add actual metadata fields, not the shape/dtype helper fields
             metadata_fields.append(field)
 
-    # Process each tensor field efficiently
+    # Process each tensor field
     for tensor_name in tensor_fields:
         tensor_list = []
 
@@ -169,9 +138,6 @@ def collate_rows_from_parquet_schema(rows, parquet_schema,
             bytes_key = f"{tensor_name}_bytes"
 
             if shape_key in row and bytes_key in row:
-                # logger.info("row: %s", row)
-                # logger.info("shape_key: %s", shape_key)
-                # logger.info("bytes_key: %s", bytes_key)
                 shape = row[shape_key]
                 bytes_data = row[bytes_key]
 
@@ -179,11 +145,12 @@ def collate_rows_from_parquet_schema(rows, parquet_schema,
                     tensor = torch.zeros(0, dtype=torch.bfloat16)
                 else:
                     # Convert bytes to tensor using float32 as default
-                    # logger.info("len(bytes_data): %s", len(bytes_data))
+                    if tensor_name == 'text_embedding' and random.random(
-                    # logger.info("shape: %s", shape)
+                    ) < cfg_rate:
-
+                        data = np.zeros((512, 4096), dtype=np.float32)
-                    data = np.frombuffer(
+                    else:
-                        bytes_data, dtype=np.float32).reshape(shape).copy()
+                        data = np.frombuffer(
+                            bytes_data, dtype=np.float32).reshape(shape).copy()
                     tensor = torch.from_numpy(data)
                     # if len(data.shape) == 3:
                     #     B, L, D = tensor.shape
@@ -196,38 +163,44 @@ def collate_rows_from_parquet_schema(rows, parquet_schema,
                 tensor_list.append(torch.zeros(0, dtype=torch.bfloat16))
 
         # Stack tensors with special handling for text embeddings
-        if tensor_list:
+        if tensor_name == 'text_embedding':
-            if tensor_name == 'text_embedding':
+            # Handle text embeddings with padding
-                # Handle text embeddings with padding
+            padded_tensors = []
-                padded_tensors = []
+            attention_masks = []
-                attention_masks = []
 
-                for tensor in tensor_list:
+            for tensor in tensor_list:
-                    if tensor.numel() > 0:
+                if tensor.numel() > 0:
-                        padded_tensor, mask = pad(tensor, text_padding_length)
+                    padded_tensor, mask = pad(tensor, text_padding_length)
-                        padded_tensors.append(padded_tensor)
+                    padded_tensors.append(padded_tensor)
-                        attention_masks.append(mask)
+                    attention_masks.append(mask)
-                    else:
+                else:
-                        # Handle empty embeddings - assume default embedding dimension
+                    # Handle empty embeddings - assume default embedding dimension
-                        padded_tensors.append(
+                    padded_tensors.append(
-                            torch.zeros(text_padding_length,
+                        torch.zeros(text_padding_length,
-                                        768,
+                                    768,
-                                        dtype=torch.bfloat16))
+                                    dtype=torch.bfloat16))
-                        attention_masks.append(torch.zeros(text_padding_length))
+                    attention_masks.append(torch.zeros(text_padding_length))
 
-                batch_data[tensor_name] = torch.stack(padded_tensors)
+            batch_data[tensor_name] = torch.stack(padded_tensors)
-                batch_data['text_attention_mask'] = torch.stack(attention_masks)
+            batch_data['text_attention_mask'] = torch.stack(attention_masks)
-            else:
+        else:
-                # Stack other tensors directly, handling None values
+            # Stack all tensors to preserve batch consistency
-                valid_tensors = [
+            # Don't filter out None or empty tensors as this breaks batch sizing
-                    t for t in tensor_list if t is not None and t.numel() > 0
+            try:
+                batch_data[tensor_name] = torch.stack(tensor_list)
+            except ValueError as e:
+                shapes = [
+                    t.shape
+                    if t is not None and hasattr(t, 'shape') else 'None/Invalid'
+                    for t in tensor_list
                 ]
-                if valid_tensors:
+                raise ValueError(
-                    batch_data[tensor_name] = torch.stack(valid_tensors)
+                    f"Failed to stack tensors for field '{tensor_name}'. "
-                elif tensor_list:  # All tensors are empty but exist
+                    f"Tensor shapes: {shapes}. "
-                    batch_data[tensor_name] = torch.stack(tensor_list)
+                    f"All tensors in a batch must have compatible shapes. "
+                    f"Original error: {e}") from e
 
-    # Process metadata fields efficiently into info_list
+    # Process metadata fields into info_list
     info_list = []
     for row in rows:
         info = {}

fastvideo/v1/dataset/validation_dataset.py

@@ -19,7 +19,6 @@ class ValidationDataset(torch.utils.data.IterableDataset):
 
         self.filename = pathlib.Path(filename)
         # get directory of filename
-        # TODO(will)
         self.dir = os.path.abspath(self.filename.parent)
 
         if not self.filename.exists():

fastvideo/v1/distributed/__init__.py

@@ -3,10 +3,10 @@
 from fastvideo.v1.distributed.communication_op import *
 from fastvideo.v1.distributed.parallel_state import (
     cleanup_dist_env_and_memory, get_dp_group, get_dp_rank, get_dp_world_size,
-    get_sp_group, get_sp_parallel_rank, get_sp_world_size, get_torch_device,
+    get_local_torch_device, get_sp_group, get_sp_parallel_rank,
-    get_tp_group, get_tp_rank, get_tp_world_size, get_world_group,
+    get_sp_world_size, get_tp_group, get_tp_rank, get_tp_world_size,
-    get_world_rank, get_world_size, init_distributed_environment,
+    get_world_group, get_world_rank, get_world_size,
-    initialize_model_parallel,
+    init_distributed_environment, initialize_model_parallel,
     maybe_init_distributed_environment_and_model_parallel,
     model_parallel_is_initialized)
 from fastvideo.v1.distributed.utils import *
@@ -40,5 +40,5 @@ __all__ = [
     "get_tp_world_size",
 
     # Get torch device
-    "get_torch_device",
+    "get_local_torch_device",
 ]

fastvideo/v1/distributed/parallel_state.py

@@ -36,6 +36,7 @@ from unittest.mock import patch
 
 import torch
 import torch.distributed
+import torch.distributed as dist
 from torch.distributed import Backend, ProcessGroup, ReduceOp
 
 import fastvideo.v1.envs as envs
@@ -692,6 +693,7 @@ class GroupCoordinator:
 
 
 _WORLD: Optional[GroupCoordinator] = None
+_NODE: Optional[GroupCoordinator] = None
 
 
 def get_world_group() -> GroupCoordinator:
@@ -699,6 +701,11 @@ def get_world_group() -> GroupCoordinator:
     return _WORLD
 
 
+def get_node_group() -> GroupCoordinator:
+    assert _NODE is not None, ("node group is not initialized")
+    return _NODE
+
+
 def init_world_group(ranks: List[int], local_rank: int,
                      backend: str) -> GroupCoordinator:
     return GroupCoordinator(
@@ -710,6 +717,18 @@ def init_world_group(ranks: List[int], local_rank: int,
     )
 
 
+def init_node_group(local_rank: int, backend: str):
+    cpu_group = get_world_group().cpu_group
+    node_ranks = same_node_ranks(cpu_group)
+    node_size = len(node_ranks)
+    all_node_ranks = [
+        list(range(i * node_size, (i + 1) * node_size))
+        for i in range(dist.get_world_size() // node_size)
+    ]
+    global _NODE
+    _NODE = init_model_parallel_group(all_node_ranks, local_rank, backend)
+
+
 def init_model_parallel_group(
     group_ranks: List[List[int]],
     local_rank: int,
@@ -782,6 +801,8 @@ def init_distributed_environment(
     else:
         assert _WORLD.world_size == torch.distributed.get_world_size(), (
             "world group already initialized with a different world size")
+    # Init a group for each node
+    init_node_group(local_rank, backend)
 
 
 _SP: Optional[GroupCoordinator] = None
@@ -904,7 +925,7 @@ def get_dp_rank() -> int:
     return get_dp_group().rank_in_group
 
 
-def get_torch_device() -> torch.device:
+def get_local_torch_device() -> torch.device:
     """Return the torch device for the current rank."""
     return torch.device(f"cuda:{envs.LOCAL_RANK}")
 
@@ -1021,17 +1042,22 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
             "torch._C._host_emptyCache() only available in Pytorch >=2.5")
 
 
-def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
+def same_node_ranks(pg: Union[ProcessGroup, StatelessProcessGroup],
-                        source_rank: int = 0) -> List[bool]:
+                    source_rank: int = 0) -> List[int]:
     """
-    This is a collective operation that returns if each rank is in the same node
+    This is a collective operation that returns ranks that are in the same node
     as the source rank. It tests if processes are attached to the same
     memory system (shared access to shared memory).
+    Args:
+        pg: the global process group to test
+        source_rank: the rank to test against
+    Returns:
+        A list of ranks that are in the same node as the source rank.
     """
     if isinstance(pg, ProcessGroup):
         assert torch.distributed.get_backend(
             pg) != torch.distributed.Backend.NCCL, (
-                "in_the_same_node_as should be tested with a non-NCCL group.")
+                "same_node_ranks should be tested with a non-NCCL group.")
         # local rank inside the group
         rank = torch.distributed.get_rank(group=pg)
         world_size = torch.distributed.get_world_size(group=pg)
@@ -1103,7 +1129,7 @@ def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
             rank_data = pg.broadcast_obj(is_in_the_same_node, src=i)
             aggregated_data += rank_data
 
-    return [x == 1 for x in aggregated_data.tolist()]
+    return [i for i, x in enumerate(aggregated_data.tolist()) if x == 1]
 
 
 def initialize_tensor_parallel_group(

fastvideo/v1/fastvideo_args.py

@@ -58,8 +58,10 @@ class FastVideoArgs:
 
     output_type: str = "pil"
 
-    use_cpu_offload: bool = True
+    use_cpu_offload: bool = True  # For DiT
     use_fsdp_inference: bool = True
+    text_encoder_offload: bool = True
+    pin_cpu_memory: bool = True
 
     # STA (Sliding Tile Attention) parameters
     mask_strategy_file_path: Optional[str] = None
@@ -208,7 +210,7 @@ class FastVideoArgs:
             "--use-cpu-offload",
             action=StoreBoolean,
             help=
-            "Use CPU offload for model inference. Enable if run out of memory with FSDP.",
+            "Use CPU offload for DiT inference. Enable if run out of memory with FSDP.",
         )
         parser.add_argument(
             "--use-fsdp-inference",
@@ -216,7 +218,19 @@ class FastVideoArgs:
             help=
             "Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.",
         )
-
+        parser.add_argument(
+            "--text-encoder-cpu-offload",
+            action=StoreBoolean,
+            help=
+            "Use CPU offload for text encoder. Enable if run out of memory.",
+        )
+        parser.add_argument(
+            "--pin-cpu-memory",
+            action=StoreBoolean,
+            help=
+            "Pin memory for CPU offload. Only added as a temp workaround if it throws \"CUDA error: invalid argument\". "
+            "Should be enabled in almost all cases",
+        )
         parser.add_argument(
             "--disable-autocast",
             action=StoreBoolean,
@@ -384,7 +398,7 @@ class TrainingArgs(FastVideoArgs):
     # diffusion setting
     ema_decay: float = 0.0
     ema_start_step: int = 0
-    cfg: float = 0.0
+    training_cfg_rate: float = 0.0
     precondition_outputs: bool = False
 
     # validation & logs
@@ -528,7 +542,7 @@ class TrainingArgs(FastVideoArgs):
                             type=int,
                             default=0,
                             help="Step to start EMA")
-        parser.add_argument("--cfg",
+        parser.add_argument("--training-cfg-rate",
                             type=float,
                             help="Classifier-free guidance scale")
         parser.add_argument(

fastvideo/v1/layers/layernorm.py

@@ -6,6 +6,7 @@ from typing import Optional, Tuple, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.distributed.tensor import DTensor
 
 from fastvideo.v1.layers.custom_op import CustomOp
 
@@ -70,7 +71,12 @@ class RMSNorm(CustomOp):
         x = x * torch.rsqrt(variance + self.variance_epsilon)
         x = x.to(orig_dtype)
         if self.has_weight:
-            x = x * self.weight
+            # TODO(wenxuan): When using CPU offload, FSDP has a bug that doesn't unwrap DTensor in final_layer_norm.
+            # Report this
+            if isinstance(self.weight, DTensor):
+                x = x * self.weight.to(x.device).full_tensor()
+            else:
+                x = x * self.weight
         if residual is None:
             return x
         else:

fastvideo/v1/models/dits/stepvideo.py

@@ -455,10 +455,7 @@ class StepVideoTransformerBlock(nn.Module):
 
 class StepVideoModel(BaseDiT):
     # (Optional) Keep the same attribute for compatibility with splitting, etc.
-    _fsdp_shard_conditions = [
+    _fsdp_shard_conditions = StepVideoConfig()._fsdp_shard_conditions
-        lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit(),
-        # lambda n, m: "pos_embed" in n  # If needed for the patch embedding.
-    ]
     _param_names_mapping = StepVideoConfig()._param_names_mapping
     _reverse_param_names_mapping = StepVideoConfig(
     )._reverse_param_names_mapping

fastvideo/v1/models/dits/wanvideo.py

@@ -25,12 +25,9 @@ from fastvideo.v1.layers.rotary_embedding import (_apply_rotary_emb,
                                                   get_rotary_pos_embed)
 from fastvideo.v1.layers.visual_embedding import (ModulateProjection,
                                                   PatchEmbed, TimestepEmbedder)
-from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.dits.base import CachableDiT
 from fastvideo.v1.platforms import AttentionBackendEnum
 
-logger = init_logger(__name__)
-
 
 class WanImageEmbedding(torch.nn.Module):
 

fastvideo/v1/models/encoders/base.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from abc import ABC, abstractmethod
-from typing import Optional, Tuple
+from dataclasses import field
+from typing import List, Optional, Tuple
 
 import torch
 from torch import nn
@@ -12,6 +13,9 @@ from fastvideo.v1.platforms import AttentionBackendEnum
 
 
 class TextEncoder(nn.Module, ABC):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [])
+    _stacked_params_mapping: List[Tuple[str, str,
+                                        str]] = field(default_factory=list)
     _supported_attention_backends: Tuple[
         AttentionBackendEnum,
         ...] = TextEncoderConfig()._supported_attention_backends
@@ -19,6 +23,8 @@ class TextEncoder(nn.Module, ABC):
     def __init__(self, config: TextEncoderConfig) -> None:
         super().__init__()
         self.config = config
+        self._fsdp_shard_conditions = config._fsdp_shard_conditions
+        self._stacked_params_mapping = config.arch_config.stacked_params_mapping
         if not self.supported_attention_backends:
             raise ValueError(
                 f"Subclass {self.__class__.__name__} must define _supported_attention_backends"

fastvideo/v1/models/encoders/clip.py

@@ -596,12 +596,7 @@ class CLIPVisionModel(ImageEncoder):
     #        ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986
     def load_weights(self, weights: Iterable[Tuple[str,
                                                    torch.Tensor]]) -> Set[str]:
-        stacked_params_mapping = [
+
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
         params_dict = dict(self.named_parameters())
         loaded_params: Set[str] = set()
         layer_count = len(self.vision_model.encoder.layers)
@@ -620,7 +615,8 @@ class CLIPVisionModel(ImageEncoder):
                 if layer_idx >= layer_count:
                     continue
 
-            for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            for (param_name, weight_name,
+                 shard_id) in self.config.arch_config.stacked_params_mapping:
                 if weight_name not in name:
                     continue
                 name = name.replace(weight_name, param_name)

fastvideo/v1/models/encoders/llama.py

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

fastvideo/v1/models/encoders/t5.py

@@ -494,7 +494,7 @@ class T5Stack(nn.Module):
                 attention_mask=attention_mask,
                 attn_metadata=attn_metadata,
             )
-        hidden_states = self.final_layer_norm.forward_native(hidden_states)
+        hidden_states = self.final_layer_norm.forward(hidden_states)
         return hidden_states
 
 
@@ -631,19 +631,13 @@ class UMT5EncoderModel(TextEncoder):
 
     def load_weights(self, weights: Iterable[Tuple[str,
                                                    torch.Tensor]]) -> Set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".qkv_proj", ".q", "q"),
-            (".qkv_proj", ".k", "k"),
-            (".qkv_proj", ".v", "v"),
-        ]
         params_dict = dict(self.named_parameters())
         loaded_params: Set[str] = set()
         for name, loaded_weight in weights:
             loaded = False
             if "decoder" in name or "lm_head" in name:
                 continue
-            for param_name, weight_name, shard_id in stacked_params_mapping:
+            for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
                 if weight_name not in name:
                     continue
                 name = name.replace(weight_name, param_name)

fastvideo/v1/models/loader/component_loader.py

@@ -10,17 +10,20 @@ from copy import deepcopy
 from typing import Any, Generator, Iterable, List, Optional, Tuple, cast
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 from safetensors.torch import load_file as safetensors_load_file
 from transformers import AutoImageProcessor, AutoTokenizer
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
 
 from fastvideo.v1.configs.models import EncoderConfig
-from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.hf_transformer_utils import get_diffusers_config
-from fastvideo.v1.models.loader.fsdp_load import maybe_load_fsdp_model
+from fastvideo.v1.models.loader.fsdp_load import (init_device_mesh,
+                                                  maybe_load_fsdp_model,
+                                                  shard_model)
 from fastvideo.v1.models.loader.utils import set_default_torch_dtype
 from fastvideo.v1.models.loader.weight_utils import (
     filter_duplicate_safetensors_files, filter_files_not_needed_for_inference,
@@ -163,16 +166,19 @@ class TextEncoderLoader(ComponentLoader):
         return hf_folder, hf_weights_files, use_safetensors
 
     def _get_weights_iterator(
-            self, source: "Source"
+            self,
+            source: "Source",
+            to_cpu: bool = True
     ) -> Generator[Tuple[str, torch.Tensor], None, None]:
         """Get an iterator for the model weights based on the load format."""
         hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
             source.model_or_path, source.fall_back_to_pt,
             source.allow_patterns_overrides)
         if use_safetensors:
-            weights_iterator = safetensors_weights_iterator(hf_weights_files)
+            weights_iterator = safetensors_weights_iterator(
+                hf_weights_files, to_cpu)
         else:
-            weights_iterator = pt_weights_iterator(hf_weights_files)
+            weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu)
 
         if self.counter_before_loading_weights == 0.0:
             self.counter_before_loading_weights = time.perf_counter()
@@ -181,10 +187,11 @@ class TextEncoderLoader(ComponentLoader):
                 for (name, tensor) in weights_iterator)
 
     def _get_all_weights(
-        self,
+            self,
-        model_config: Any,
+            model_config: Any,
-        model: nn.Module,
+            model: nn.Module,
-        model_path: str,
+            model_path: str,
+            to_cpu: bool = True
     ) -> Generator[Tuple[str, torch.Tensor], None, None]:
         primary_weights = TextEncoderLoader.Source(
             model_path,
@@ -193,14 +200,14 @@ class TextEncoderLoader(ComponentLoader):
             allow_patterns_overrides=getattr(model, "allow_patterns_overrides",
                                              None),
         )
-        yield from self._get_weights_iterator(primary_weights)
+        yield from self._get_weights_iterator(primary_weights, to_cpu)
 
         secondary_weights = cast(
             Iterable[TextEncoderLoader.Source],
             getattr(model, "secondary_weights", ()),
         )
         for source in secondary_weights:
-            yield from self._get_weights_iterator(source)
+            yield from self._get_weights_iterator(source, to_cpu)
 
     def load(self, model_path: str, architecture: str,
              fastvideo_args: FastVideoArgs):
@@ -233,16 +240,22 @@ class TextEncoderLoader(ComponentLoader):
             encoder_precision = fastvideo_args.pipeline_config.text_encoder_precisions[
                 1]
 
-        target_device = get_torch_device()
+        target_device = get_local_torch_device()
         # TODO(will): add support for other dtypes
         return self.load_model(model_path, encoder_config, target_device,
-                               encoder_precision)
+                               fastvideo_args, encoder_precision)
 
     def load_model(self,
                    model_path: str,
                    model_config: EncoderConfig,
                    target_device: torch.device,
+                   fastvideo_args: FastVideoArgs,
                    dtype: str = "fp16"):
+        use_cpu_offload = fastvideo_args.text_encoder_offload and len(
+            getattr(model_config, "_fsdp_shard_conditions", [])) > 0
+
+        if fastvideo_args.text_encoder_offload:
+            target_device = torch.device("cpu")
         with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]):
             with target_device:
                 architectures = getattr(model_config, "architectures", [])
@@ -251,12 +264,26 @@ class TextEncoderLoader(ComponentLoader):
 
             weights_to_load = {name for name, _ in model.named_parameters()}
             loaded_weights = model.load_weights(
-                self._get_all_weights(model_config, model, model_path))
+                self._get_all_weights(model_config, model, model_path,
+                                      use_cpu_offload))
             self.counter_after_loading_weights = time.perf_counter()
             logger.info(
                 "Loading weights took %.2f seconds",
                 self.counter_after_loading_weights -
                 self.counter_before_loading_weights)
+
+            if use_cpu_offload:
+                mesh = init_device_mesh(
+                    "cuda",
+                    mesh_shape=(1, dist.get_world_size()),
+                    mesh_dim_names=("offload", "replicate"),
+                )
+                shard_model(model,
+                            cpu_offload=True,
+                            reshard_after_forward=True,
+                            mesh=mesh["offload"],
+                            fsdp_shard_conditions=model._fsdp_shard_conditions,
+                            pin_cpu_memory=fastvideo_args.pin_cpu_memory)
             # We only enable strict check for non-quantized models
             # that have loaded weights tracking currently.
             # if loaded_weights is not None:
@@ -290,10 +317,10 @@ class ImageEncoderLoader(TextEncoderLoader):
         encoder_config = fastvideo_args.pipeline_config.image_encoder_config
         encoder_config.update_model_arch(model_config)
 
-        target_device = get_torch_device()
+        target_device = get_local_torch_device()
         # TODO(will): add support for other dtypes
         return self.load_model(
-            model_path, encoder_config, target_device,
+            model_path, encoder_config, target_device, fastvideo_args,
             fastvideo_args.pipeline_config.image_encoder_precision)
 
 
@@ -346,7 +373,7 @@ class VAELoader(ComponentLoader):
         with set_default_torch_dtype(PRECISION_TO_TYPE[
                 fastvideo_args.pipeline_config.vae_precision]):
             vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-            vae = vae_cls(vae_config).to(get_torch_device())
+            vae = vae_cls(vae_config).to(get_local_torch_device())
 
         # Find all safetensors files
         safetensors_list = glob.glob(
@@ -405,7 +432,7 @@ class TransformerLoader(ComponentLoader):
                 "hf_config": hf_config
             },
             weight_dir_list=safetensors_list,
-            device=get_torch_device(),
+            device=get_local_torch_device(),
             hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
             hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
             cpu_offload=fastvideo_args.use_cpu_offload,

fastvideo/v1/models/loader/fsdp_load.py

@@ -69,6 +69,7 @@ def maybe_load_fsdp_model(
     fsdp_inference: bool = False,
     output_dtype: Optional[torch.dtype] = None,
     training_mode: bool = True,
+    pin_cpu_memory: bool = True,
 ) -> torch.nn.Module:
     """
     Load the model with FSDP if is training, else load the model without FSDP.
@@ -101,9 +102,12 @@ def maybe_load_fsdp_model(
                 cpu_offload=cpu_offload,
                 reshard_after_forward=True,
                 mp_policy=mp_policy,
-                mesh=device_mesh)
+                mesh=device_mesh,
+                fsdp_shard_conditions=model._fsdp_shard_conditions,
+                pin_cpu_memory=pin_cpu_memory)
 
-    weight_iterator = safetensors_weights_iterator(weight_dir_list)
+    weight_iterator = safetensors_weights_iterator(
+        weight_dir_list, to_cpu=cpu_offload, async_broadcast=not cpu_offload)
     param_names_mapping_fn = get_param_names_mapping(model._param_names_mapping)
     load_model_from_full_model_state_dict(
         model,
@@ -126,12 +130,13 @@ def maybe_load_fsdp_model(
 
 def shard_model(
     model,
-    *,
     cpu_offload: bool,
     reshard_after_forward: bool = True,
-    mp_policy: Optional[MixedPrecisionPolicy] = None,
+    mp_policy: Optional[MixedPrecisionPolicy] = MixedPrecisionPolicy(),  # noqa
-    dp_mesh: Optional[DeviceMesh] = None,
     mesh: Optional[DeviceMesh] = None,
+    fsdp_shard_conditions: Optional[List[Callable[[str, nn.Module],
+                                                  bool]]] = None,
+    pin_cpu_memory: bool = True,
 ) -> None:
     """
     Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API.
@@ -150,19 +155,28 @@ def shard_model(
         reshard_after_forward (bool): Whether to reshard parameters and buffers after
             the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy
             from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy.
-        dp_mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism.
+        mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism.
             Default to None.
+        fsdp_shard_conditions (Optional[List[Callable[[str, nn.Module], bool]]]): A list of functions to determine
+            which modules to shard with FSDP. 
 
     Raises:
         ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered.
     """
+    if fsdp_shard_conditions is None or len(fsdp_shard_conditions) == 0:
+        logger.warning(
+            "The FSDP shard condition list is empty or None. No modules will be sharded in %s",
+            type(model).__name__)
+        return
+
     fsdp_kwargs = {
         "reshard_after_forward": reshard_after_forward,
         "mesh": mesh,
         "mp_policy": mp_policy,
     }
     if cpu_offload:
-        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy()
+        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(
+            pin_memory=pin_cpu_memory)
 
     # iterating in reverse to start with
     # lowest-level modules first
@@ -172,7 +186,7 @@ def shard_model(
     for n, m in reversed(list(model.named_modules())):
         if any([
                 shard_condition(n, m)
-                for shard_condition in model._fsdp_shard_conditions
+                for shard_condition in fsdp_shard_conditions
         ]):
             fully_shard(m, **fsdp_kwargs)
             num_layers_sharded += 1
@@ -181,7 +195,6 @@ def shard_model(
         raise ValueError(
             "No layer modules were sharded. Please check if shard conditions are working as expected."
         )
-
     # Finally shard the entire model to account for any stragglers
     fully_shard(model, **fsdp_kwargs)
 
@@ -224,6 +237,9 @@ def load_model_from_full_model_state_dict(
     to_merge_params: DefaultDict[str, Dict[Any, Any]] = defaultdict(dict)
     reverse_param_names_mapping = {}
     assert param_names_mapping is not None
+
+    # iterate over all the weights to sync broadcast before use
+    full_sd_iterator = list(full_sd_iterator)  # type: ignore
     for source_param_name, full_tensor in full_sd_iterator:
         target_param_name, merge_index, num_params_to_merge = param_names_mapping(
             source_param_name)

fastvideo/v1/models/loader/weight_utils.py

@@ -11,9 +11,11 @@ from typing import Generator, List, Optional, Tuple, Union
 import filelock
 import huggingface_hub.constants
 import torch
+import torch.distributed as dist
 from safetensors.torch import safe_open
 from tqdm.auto import tqdm
 
+from fastvideo.v1.distributed.parallel_state import get_node_group
 from fastvideo.v1.logger import init_logger
 
 logger = init_logger(__name__)
@@ -118,36 +120,77 @@ _BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elap
 
 
 def safetensors_weights_iterator(
-    hf_weights_files: List[str]
+    hf_weights_files: List[str],
+    to_cpu: bool = False,
+    async_broadcast: bool = False
 ) -> Generator[Tuple[str, torch.Tensor], None, None]:
-    """Iterate over the weights in the model safetensor files."""
+    """Iterate over the weights in the model safetensor files.
-    enable_tqdm = not torch.distributed.is_initialized(
+    Args:
-    ) or torch.distributed.get_rank() == 0
+        hf_weights_files: List of safetensor files to load.
+        to_cpu: Whether to load the weights to CPU. If False, will load to the GPU device bound to the current process.
+        async_broadcast: Whether to overlap loading from disk and broadcasting to other ranks. If True,
+            must iterate over all the weights before use. Only use if to_cpu is False.
+    """
+    local_rank = get_node_group().rank
+    device = f"cuda:{local_rank}" if not to_cpu else "cpu"
+    enable_tqdm = not torch.distributed.is_initialized() or get_node_group(
+    ).rank == 0
+    assert not (async_broadcast
+                and to_cpu), "Cannot broadcast weights when loading to CPU"
+
+    handles = []
     for st_file in tqdm(
             hf_weights_files,
             desc="Loading safetensors checkpoint shards",
             disable=not enable_tqdm,
             bar_format=_BAR_FORMAT,
     ):
-        with safe_open(st_file, framework="pt") as f:
+        with safe_open(st_file, framework="pt", device=device) as f:
             for name in f.keys():  # noqa: SIM118
-                param = f.get_tensor(name)
+                if to_cpu:
+                    param = f.get_tensor(name)
+                else:
+                    if local_rank == 0:
+                        param = f.get_tensor(name)
+                    else:
+                        shape = f.get_slice(name).get_shape()
+                        param = torch.empty(shape, device=device)
+                    # broadcast to local ranks
+                    # TODO(Wenxuan): scatter instead of broadcast
+                    if get_node_group().world_size > 1:
+                        group = get_node_group().device_group
+                        if async_broadcast:
+                            handle = dist.broadcast(param,
+                                                    src=dist.get_global_rank(
+                                                        group, 0),
+                                                    async_op=True)
+                            handles.append(handle)
+                        else:
+                            dist.broadcast(param,
+                                           src=dist.get_global_rank(group, 0))
                 yield name, param
 
+        if async_broadcast:
+            for handle in handles:
+                handle.wait()
+
 
 def pt_weights_iterator(
-    hf_weights_files: List[str]
+    hf_weights_files: List[str],
+    to_cpu: bool = True  # default to CPU for text encoder
 ) -> Generator[Tuple[str, torch.Tensor], None, None]:
     """Iterate over the weights in the model bin/pt files."""
-    enable_tqdm = not torch.distributed.is_initialized(
+    local_rank = get_node_group().rank
-    ) or torch.distributed.get_rank() == 0
+    device = f"cuda:{local_rank}" if not to_cpu else "cpu"
+    enable_tqdm = not torch.distributed.is_initialized() or get_node_group(
+    ).rank == 0
     for bin_file in tqdm(
             hf_weights_files,
             desc="Loading pt checkpoint shards",
             disable=not enable_tqdm,
             bar_format=_BAR_FORMAT,
     ):
-        state = torch.load(bin_file, map_location="cpu", weights_only=True)
+        state = torch.load(bin_file, map_location=device, weights_only=True)
         yield from state.items()
         del state
 

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -149,11 +149,9 @@ class TrainingBatch:
 
     # Dataloader batch outputs
     latents: Optional[torch.Tensor] = None
-    # original_latents: Optional[torch.Tensor] = None
     encoder_hidden_states: Optional[torch.Tensor] = None
     encoder_attention_mask: Optional[torch.Tensor] = None
     # i2v
-    # extra_latents: Optional[Dict[str, Any]] = None
     preprocessed_image: Optional[torch.Tensor] = None
     image_embeds: Optional[torch.Tensor] = None
     image_latents: Optional[torch.Tensor] = None

fastvideo/v1/pipelines/preprocess/preprocess_pipeline_base.py

@@ -18,7 +18,7 @@ from tqdm import tqdm
 from fastvideo.v1.configs.sample import SamplingParam
 from fastvideo.v1.dataset import ValidationDataset, getdataset
 from fastvideo.v1.dataset.preprocessing_datasets import PreprocessBatch
-from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
@@ -104,7 +104,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                         if strict:
                             raise ValueError(
                                 f"Failed to convert tensor {tensor_name} to bytes: {e}"
-                            )
+                            ) from e
                         record[field] = b''  # Empty bytes for missing data
                 else:
                     record[field] = b''  # Empty bytes for missing data
@@ -139,7 +139,8 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                     except (ValueError, TypeError) as e:
                         if strict:
                             raise ValueError(
-                                f"Failed to convert field {field} to int: {e}")
+                                f"Failed to convert field {field} to int: {e}"
+                            ) from e
                         record[field] = 0
                 else:
                     record[field] = 0
@@ -157,7 +158,7 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                         if strict:
                             raise ValueError(
                                 f"Failed to convert field {field} to float: {e}"
-                            )
+                            ) from e
                         record[field] = 0.0
                 else:
                     record[field] = 0.0
@@ -211,8 +212,8 @@ class BasePreprocessPipeline(ComposedPipelineBase):
         # Log unfilled fields as warning if not in strict mode
         if unfilled_fields:
             logger.warning(
-                f"Some fields were not filled and got default values: {unfilled_fields}"
+                "Some fields were not filled and got default values: %s",
-            )
+                unfilled_fields)
 
         return record
 
@@ -221,7 +222,6 @@ class BasePreprocessPipeline(ComposedPipelineBase):
             video_name: str,
             vae_latent: np.ndarray,
             text_embedding: np.ndarray,
-            # text_attention_mask: np.ndarray,
             valid_data: Dict[str, Any],
             idx: int,
             extra_features: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
@@ -328,7 +328,8 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                 # VAE
                 with torch.autocast("cuda", dtype=torch.float32):
                     latents = self.get_module("vae").encode(
-                        valid_data["pixel_values"].to(get_torch_device())).mean
+                        valid_data["pixel_values"].to(
+                            get_local_torch_device())).mean
 
                 # Get extra features if needed
                 extra_features = self.get_extra_features(
@@ -380,8 +381,6 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                 # Convert tensors to numpy arrays
                 vae_latent = latent.cpu().numpy()
                 text_embedding = prompt_embeds[idx].cpu().numpy()
-                # text_attention_mask = prompt_attention_mask[idx].cpu().numpy(
-                # ).astype(np.uint8)
 
                 # Get extra features for this sample if needed
                 sample_extra_features = {}
@@ -398,7 +397,6 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                     video_name=video_name,
                     vae_latent=vae_latent,
                     text_embedding=text_embedding,
-                    # text_attention_mask=text_attention_mask,
                     valid_data=valid_data,
                     idx=idx,
                     extra_features=sample_extra_features)
@@ -543,14 +541,13 @@ class BasePreprocessPipeline(ComposedPipelineBase):
             valid_data["text"] = [prompt]
 
             # Create record for Parquet dataset
-            record = self.create_record(
+            record = self.create_record(video_name=file_name,
-                video_name=file_name,
+                                        vae_latent=np.array([],
-                vae_latent=np.array([], dtype=np.float32),
+                                                            dtype=np.float32),
-                text_embedding=text_embedding,
+                                        text_embedding=text_embedding,
-                # text_attention_mask=text_attention_mask,
+                                        valid_data=valid_data,
-                valid_data=valid_data,
+                                        idx=0,
-                idx=0,
+                                        extra_features=sample_extra_features)
-                extra_features=sample_extra_features)
             batch_data.append(record)
 
             logger.info("Saved validation sample: %s", file_name)

fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py

@@ -13,7 +13,7 @@ import torch
 from PIL import Image
 
 from fastvideo.v1.dataset.dataloader.schema import pyarrow_schema_i2v
-from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.models.vision_utils import (get_default_height_width,
@@ -58,7 +58,6 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
         result_batch = self.image_encoding_stage(batch, fastvideo_args)
         clip_features = result_batch.image_embeds[0]
 
-        # image = self.pil_to_tensor(image)
         image = self.preprocess(
             image,
             vae_scale_factor=self.get_module("vae").spatial_compression_ratio,
@@ -83,9 +82,8 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
                            fastvideo_args: FastVideoArgs) -> Dict[str, Any]:
 
         # TODO(will): move these to cpu at some point
-        self.get_module("image_encoder").to(get_torch_device())
+        self.get_module("image_encoder").to(get_local_torch_device())
-        # self.get_module("image_processor").to(get_torch_device())
+        self.get_module("vae").to(get_local_torch_device())
-        self.get_module("vae").to(get_torch_device())
 
         features = {}
         """Get CLIP features from the first frame of each video."""
@@ -109,7 +107,7 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
         # Get CLIP features
         pixel_values = torch.cat(
             [img['pixel_values'] for img in processed_images],
-            dim=0).to(get_torch_device())
+            dim=0).to(get_local_torch_device())
         with torch.no_grad():
             image_inputs = {'pixel_values': pixel_values}
             with set_forward_context(current_timestep=0, attn_metadata=None):
@@ -131,8 +129,8 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
                                         height, width)
             ],
                                         dim=2)
-            video_condition = video_condition.to(device=get_torch_device(),
+            video_condition = video_condition.to(
-                                                 dtype=torch.float32)
+                device=get_local_torch_device(), dtype=torch.float32)
             video_conditions.append(video_condition)
 
         video_conditions = torch.cat(video_conditions, dim=0)
@@ -187,19 +185,16 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
             video_name: str,
             vae_latent: np.ndarray,
             text_embedding: np.ndarray,
-            # text_attention_mask: np.ndarray,
+            valid_data: Dict[str, Any],
-            valid_data: Optional[Dict[str, Any]],
             idx: int,
             extra_features: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
         """Create a record for the Parquet dataset with CLIP features."""
-        record = super().create_record(
+        record = super().create_record(video_name=video_name,
-            video_name=video_name,
+                                       vae_latent=vae_latent,
-            vae_latent=vae_latent,
+                                       text_embedding=text_embedding,
-            text_embedding=text_embedding,
+                                       valid_data=valid_data,
-            #    text_attention_mask=text_attention_mask,
+                                       idx=idx,
-            valid_data=valid_data,
+                                       extra_features=extra_features)
-            idx=idx,
-            extra_features=extra_features)
 
         if extra_features and "clip_feature" in extra_features:
             clip_feature = extra_features["clip_feature"]

fastvideo/v1/pipelines/preprocess/v1_preprocess.py

@@ -84,7 +84,7 @@ if __name__ == "__main__":
                         type=str,
                         default="google/t5-v1_1-xxl")
     parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-    parser.add_argument("--cfg", type=float, default=0.0)
+    parser.add_argument("--training_cfg_rate", type=float, default=0.0)
     parser.add_argument(
         "--output_dir",
         type=str,

fastvideo/v1/pipelines/stages/decoding.py

@@ -5,7 +5,7 @@ Decoding stage for diffusion pipelines.
 
 import torch
 
-from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.vaes.common import ParallelTiledVAE
@@ -61,7 +61,7 @@ class DecodingStage(PipelineStage):
         Returns:
             The batch with decoded outputs.
         """
-        self.vae = self.vae.to(get_torch_device())
+        self.vae = self.vae.to(get_local_torch_device())
 
         latents = batch.latents
         # TODO(will): remove this once we add input/output validation for stages

fastvideo/v1/pipelines/stages/denoising.py

@@ -12,8 +12,9 @@ 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,
+from fastvideo.v1.distributed import (get_local_torch_device,
-                                      get_torch_device, get_world_group)
+                                      get_sp_parallel_rank, get_sp_world_size,
+                                      get_world_group)
 from fastvideo.v1.distributed.communication_op import (
     sequence_model_parallel_all_gather)
 from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -192,7 +193,7 @@ class DenoisingStage(PipelineStage):
                         [fastvideo_args.pipeline_config.embedded_cfg_scale] *
                         latent_model_input.shape[0],
                         dtype=torch.float32,
-                        device=get_torch_device(),
+                        device=get_local_torch_device(),
                     ).to(target_dtype) *
                     1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale
                     is not None else None)

fastvideo/v1/pipelines/stages/encoding.py

@@ -7,7 +7,7 @@ from typing import Optional
 import PIL.Image
 import torch
 
-from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.vaes.common import ParallelTiledVAE
@@ -49,7 +49,7 @@ class EncodingStage(PipelineStage):
         Returns:
             The batch with encoded outputs.
         """
-        self.vae = self.vae.to(get_torch_device())
+        self.vae = self.vae.to(get_local_torch_device())
 
         assert batch.height is not None
         assert batch.width is not None
@@ -65,18 +65,21 @@ class EncodingStage(PipelineStage):
                 image,
                 vae_scale_factor=self.vae.spatial_compression_ratio,
                 height=batch.height,
-                width=batch.width).to(get_torch_device(), dtype=torch.float32)
+                width=batch.width).to(get_local_torch_device(),
+                                      dtype=torch.float32)
 
             image = image.unsqueeze(2)
+        else:
+            # assumes image is loaded from parquet file and used for validation
+            image = image.transpose(1, 2)
         logger.info("image: %s", image.shape)
-        image = image.transpose(1, 2)
         video_condition = torch.cat([
             image,
             image.new_zeros(image.shape[0], image.shape[1],
                             batch.num_frames - 1, batch.height, batch.width)
         ],
                                     dim=2)
-        video_condition = video_condition.to(device=get_torch_device(),
+        video_condition = video_condition.to(device=get_local_torch_device(),
                                              dtype=torch.float32)
 
         # Setup VAE precision
@@ -100,6 +103,7 @@ class EncodingStage(PipelineStage):
         generator = batch.generator
         if generator is None:
             raise ValueError("Generator must be provided")
+        # latent_condition = self.retrieve_latents(encoder_output, generator, sample_mode="argmax")
         latent_condition = self.retrieve_latents(encoder_output, generator)
 
         # Apply shifting if needed

fastvideo/v1/pipelines/stages/image_encoding.py

@@ -7,7 +7,7 @@ This module contains implementations of image encoding stages for diffusion pipe
 
 import torch
 
-from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
@@ -55,12 +55,12 @@ class ImageEncodingStage(PipelineStage):
             The batch with encoded prompt embeddings.
         """
         if fastvideo_args.use_cpu_offload:
-            self.image_encoder = self.image_encoder.to(get_torch_device())
+            self.image_encoder = self.image_encoder.to(get_local_torch_device())
 
         image = batch.pil_image
 
         image_inputs = self.image_processor(
-            images=image, return_tensors="pt").to(get_torch_device())
+            images=image, return_tensors="pt").to(get_local_torch_device())
         with set_forward_context(current_timestep=0, attn_metadata=None):
             outputs = self.image_encoder(**image_inputs)
             image_embeds = outputs.last_hidden_state

fastvideo/v1/pipelines/stages/latent_preparation.py

@@ -5,7 +5,7 @@ Latent preparation stage for diffusion pipelines.
 
 from diffusers.utils.torch_utils import randn_tensor
 
-from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
@@ -62,7 +62,7 @@ class LatentPreparationStage(PipelineStage):
 
         # Get required parameters
         dtype = batch.prompt_embeds[0].dtype
-        device = get_torch_device()
+        device = get_local_torch_device()
         generator = batch.generator
         latents = batch.latents
         num_frames = latent_num_frames if latent_num_frames is not None else batch.num_frames

fastvideo/v1/pipelines/stages/text_encoding.py

@@ -5,9 +5,7 @@ Prompt encoding stages for diffusion pipelines.
 This module contains implementations of prompt encoding stages for diffusion pipelines.
 """
 
-import torch
+from fastvideo.v1.distributed import get_local_torch_device
-
-from fastvideo.v1.distributed import get_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
@@ -62,8 +60,6 @@ class TextEncodingStage(PipelineStage):
                 fastvideo_args.pipeline_config.text_encoder_configs,
                 fastvideo_args.pipeline_config.preprocess_text_funcs,
                 fastvideo_args.pipeline_config.postprocess_text_funcs):
-            if fastvideo_args.use_cpu_offload:
-                text_encoder = text_encoder.to(get_torch_device())
 
             assert isinstance(batch.prompt, (str, list))
             if isinstance(batch.prompt, str):
@@ -71,8 +67,9 @@ class TextEncodingStage(PipelineStage):
             texts = []
             for prompt_str in batch.prompt:
                 texts.append(preprocess_func(prompt_str))
-            text_inputs = tokenizer(
+            text_inputs = tokenizer(texts,
-                texts, **encoder_config.tokenizer_kwargs).to(get_torch_device())
+                                    **encoder_config.tokenizer_kwargs).to(
+                                        get_local_torch_device())
             input_ids = text_inputs["input_ids"]
             attention_mask = text_inputs["attention_mask"]
             with set_forward_context(current_timestep=0, attn_metadata=None):
@@ -91,8 +88,8 @@ class TextEncodingStage(PipelineStage):
                 assert isinstance(batch.negative_prompt, str)
                 negative_text = preprocess_func(batch.negative_prompt)
                 negative_text_inputs = tokenizer(
-                    negative_text,
+                    negative_text, **encoder_config.tokenizer_kwargs).to(
-                    **encoder_config.tokenizer_kwargs).to(get_torch_device())
+                        get_local_torch_device())
                 negative_input_ids = negative_text_inputs["input_ids"]
                 negative_attention_mask = negative_text_inputs["attention_mask"]
                 with set_forward_context(current_timestep=0,
@@ -110,9 +107,8 @@ class TextEncodingStage(PipelineStage):
                     batch.negative_attention_mask.append(
                         negative_attention_mask)
 
-            if fastvideo_args.use_cpu_offload:
+            if fastvideo_args.text_encoder_offload:
                 text_encoder.to('cpu')
-                torch.cuda.empty_cache()
 
         return batch
 

fastvideo/v1/pipelines/stages/timestep_preparation.py

@@ -7,7 +7,7 @@ This module contains implementations of timestep preparation stages for diffusio
 
 import inspect
 
-from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
@@ -45,7 +45,7 @@ class TimestepPreparationStage(PipelineStage):
             The batch with prepared timesteps.
         """
         scheduler = self.scheduler
-        device = get_torch_device()
+        device = get_local_torch_device()
         num_inference_steps = batch.num_inference_steps
         timesteps = batch.timesteps
         sigmas = batch.sigmas

fastvideo/v1/pipelines/stepvideo/stepvideo_pipeline.py

@@ -14,7 +14,7 @@ from typing import Any, Dict
 import torch
 from huggingface_hub import hf_hub_download
 
-from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.encoders.bert import HunyuanClip  # type: ignore
@@ -78,7 +78,7 @@ class StepVideoPipeline(LoRAPipeline, ComposedPipelineBase):
         """
         Initialize the pipeline.
         """
-        target_device = get_torch_device()
+        target_device = get_local_torch_device()
         llm_dir = os.path.join(self.model_path, "step_llm")
         clip_dir = os.path.join(self.model_path, "hunyuan_clip")
         text_enc = self.build_llm(llm_dir, target_device)

fastvideo/v1/tests/encoders/test_clip_encoder.py

@@ -6,7 +6,7 @@ import numpy as np
 import pytest
 import torch
 from transformers import AutoConfig
-
+import gc
 from fastvideo.models.hunyuan.text_encoder import (load_text_encoder,
                                                    load_tokenizer)
 # from fastvideo.v1.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer
@@ -16,6 +16,8 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import maybe_download_model
 from fastvideo.v1.configs.models.encoders import CLIPTextConfig
+from torch.distributed.tensor import DTensor
+from torch.testing import assert_close
 
 logger = init_logger(__name__)
 
@@ -66,7 +68,6 @@ def test_clip_encoder():
     # Load the HuggingFace implementation directly
     # model2 = CLIPTextModel(hf_config)
     # model2 = model2.to(torch.float16)
-    model2 = model2.to(device)
     model2.eval()
 
     # Sanity check weights between the two models
@@ -78,19 +79,20 @@ def test_clip_encoder():
     logger.info("Model1 has %d parameters", len(params1))
     logger.info("Model2 has %d parameters", len(params2))
 
-    # Compare a few key parameters
+    for name1, param1 in sorted(params1.items()):
-
+        name2 = name1
-    # weight_diffs = []
+        skip = False
-    # for (name1, param1), (name2, param2) in zip(
+        for param_name, weight_name, shard_id in model2.config.arch_config.stacked_params_mapping:
-    #     sorted(params1.items()), sorted(params2.items())
+            if weight_name not in name1:
-    # ):
+                skip = True
-    #     # if len(weight_diffs) < 5:  # Just check a few parameters
+        # stacked params are more troublesome
-    #     max_diff = torch.max(torch.abs(param1 - param2)).item()
+        if skip:
-    #     mean_diff = torch.mean(torch.abs(param1 - param2)).item()
+            continue
-    #     weight_diffs.append((name1, name2, max_diff, mean_diff))
+        param2 = params2[name2]
-    #     logger.info(f"Parameter: {name1} vs {name2}")
+        param2 = param2.to_local().to(device) if isinstance(param2, DTensor) else param2.to(device)
-    #     logger.info(f"  Max diff: {max_diff}, Mean diff: {mean_diff}")
+        assert_close(param1, param2, atol=1e-4, rtol=1e-4)
-
+    gc.collect()
+    torch.cuda.empty_cache()
     # Load tokenizer
     tokenizer, _ = load_tokenizer(tokenizer_type="clipL",
                                   tokenizer_path=args.model_path,
@@ -168,5 +170,5 @@ def test_clip_encoder():
                 f"Pooler outputs differ significantly: mean diff = {mean_diff_pooler.item()}"
             assert max_diff_hidden < 1e-1, \
                 f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}"
-            assert max_diff_pooler < 1e-2, \
+            assert max_diff_pooler < 2e-2, \
                 f"Pooler outputs differ significantly: max diff = {max_diff_pooler.item()}"

fastvideo/v1/tests/encoders/test_llama_encoder.py

@@ -5,7 +5,7 @@ import numpy as np
 import pytest
 import torch
 from transformers import AutoConfig
-
+import gc
 from fastvideo.models.hunyuan.text_encoder import (load_text_encoder,
                                                    load_tokenizer)
 from fastvideo.v1.configs.pipelines import PipelineConfig
@@ -15,7 +15,8 @@ from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.loader.component_loader import TextEncoderLoader
 from fastvideo.v1.utils import maybe_download_model
 from fastvideo.v1.configs.models.encoders import LlamaConfig
-
+from torch.distributed.tensor import DTensor
+from torch.testing import assert_close
 logger = init_logger(__name__)
 
 os.environ["MASTER_ADDR"] = "localhost"
@@ -62,7 +63,6 @@ def test_llama_encoder():
 
     # Convert to float16 and move to device
     # model2 = model2.to(torch.float16)
-    model2 = model2.to(device)
     model2.eval()
 
     # Sanity check weights between the two models
@@ -77,34 +77,28 @@ def test_llama_encoder():
     # Compare a few key parameters
     weight_diffs = []
     # check if embed_tokens are the same
-    print(model1.embed_tokens.weight.shape, model2.embed_tokens.weight.shape)
+    device = model1.embed_tokens.weight.device
     assert torch.allclose(model1.embed_tokens.weight,
-                          model2.embed_tokens.weight)
+                          model2.embed_tokens.weight.to_local().to(device) if isinstance(model2.embed_tokens.weight, DTensor) else model2.embed_tokens.weight.to(device))
     weights = [
         "layers.{}.input_layernorm.weight",
         "layers.{}.post_attention_layernorm.weight"
     ]
-    # for (name1, param1), (name2, param2) in zip(
+
-    #     sorted(params1.items()), sorted(params2.items())
+    for name1, param1 in sorted(params1.items()):
-    # ):
+        name2 = name1
-    for layer_idx in range(hf_config.num_hidden_layers):
+        skip = False
-        for w in weights:
+        for param_name, weight_name, shard_id in model2.config.arch_config.stacked_params_mapping:
-            name1 = w.format(layer_idx)
+            if weight_name not in name1:
-            name2 = w.format(layer_idx)
+                skip = True
-            p1 = params1[name1]
+        # stacked params are more troublesome
-            p2 = params2[name2]
+        if skip:
-            # print(type(p2))
+            continue
-            if "gate_up" in name2:
+        param2 = params2[name2]
-                # print("skipping gate_up")
+        param2 = param2.to_local().to(device) if isinstance(param2, DTensor) else param2.to(device)
-                continue
+        assert_close(param1, param2, atol=1e-4, rtol=1e-4)
-            try:
+    gc.collect()
-                # logger.info(f"Parameter: {name1} vs {name2}")
+    torch.cuda.empty_cache()
-                max_diff = torch.max(torch.abs(p1 - p2)).item()
-                mean_diff = torch.mean(torch.abs(p1 - p2)).item()
-                weight_diffs.append((name1, name2, max_diff, mean_diff))
-                # logger.info(f"  Max diff: {max_diff}, Mean diff: {mean_diff}")
-            except Exception as e:
-                logger.info("Error comparing %s and %s: %s", name1, name2, e)
 
     tokenizer, _ = load_tokenizer(tokenizer_type="llm",
                                   tokenizer_path=TOKENIZER_PATH,

fastvideo/v1/tests/encoders/test_t5_encoder.py

@@ -4,6 +4,8 @@ import os
 import numpy as np
 import pytest
 import torch
+from torch.distributed.tensor import DTensor
+from torch.testing import assert_close
 from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel
 
 from fastvideo.v1.configs.pipelines import PipelineConfig
@@ -41,13 +43,13 @@ def test_t5_encoder():
     tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
 
 
-    args = FastVideoArgs(model_path=TEXT_ENCODER_PATH, pipeline_config=PipelineConfig(text_encoder_configs=(T5Config(),), text_encoder_precisions=(precision_str,)))
+    args = FastVideoArgs(model_path=TEXT_ENCODER_PATH,
+                        pipeline_config=PipelineConfig(text_encoder_configs=(T5Config(),),
+                        text_encoder_precisions=(precision_str,)),
+                        pin_cpu_memory=False)
     loader = TextEncoderLoader()
     model2 = loader.load(TEXT_ENCODER_PATH, "", args)
-
+    model2 = model2.to(precision)
-    # Convert to float16 and move to device
-    # model2 = model2.to(precision)
-    model2 = model2.to(device)
     model2.eval()
 
     # Sanity check weights between the two models
@@ -64,23 +66,17 @@ def test_t5_encoder():
     weights = ["encoder.block.{}.layer.0.layer_norm.weight", "encoder.block.{}.layer.0.SelfAttention.relative_attention_bias.weight", \
                "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\
                 "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \
-                "encoder.block.{}.layer.1.layer_norm.weight", "encoder.final_layer_norm.weight", "shared.weight"]
+                "encoder.block.{}.layer.1.layer_norm.weight", "encoder.final_layer_norm.weight"]
+    
     for idx in range(hf_config.num_hidden_layers):
         for w in weights:
             name1 = w.format(idx)
             name2 = w.format(idx)
             p1 = params1[name1]
             p2 = params2[name2]
-            assert p1.dtype == p2.dtype
+            p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1)
-            try:
+            assert_close(p1, p2, atol=1e-4, rtol=1e-4)
-                logger.info("Parameter: %s vs %s", name1, name2)
+    
-                max_diff = torch.max(torch.abs(p1 - p2)).item()
-                mean_diff = torch.mean(torch.abs(p1 - p2)).item()
-                weight_diffs.append((name1, name2, max_diff, mean_diff))
-                logger.info("  Max diff: %s, Mean diff: %s", max_diff,
-                            mean_diff)
-            except Exception as e:
-                logger.info("Error comparing %s and %s: %s", name1, name2, e)
 
     # Test with some sample prompts
     prompts = [

fastvideo/v1/tests/modal/pr_test.py

@@ -4,95 +4,83 @@ app = modal.App()
 
 import os
 
-image_version = os.getenv("IMAGE_VERSION", "latest")
+image_version = os.getenv("IMAGE_VERSION")
 image_tag = f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{image_version}"
 print(f"Using image: {image_tag}")
 
 image = (
     modal.Image.from_registry(image_tag, add_python="3.12")
+    .run_commands("rm -rf /FastVideo")
     .apt_install("cmake", "pkg-config", "build-essential", "curl", "libssl-dev")
     .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable")
     .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
-    .env({"PATH": "/root/.cargo/bin:$PATH"})
+    .env({
-    .run_commands("/bin/bash -c 'source $HOME/.local/bin/env && source /opt/venv/bin/activate && cd /FastVideo && uv pip install -e .[test]'")
+        "PATH": "/root/.cargo/bin:$PATH",
+        "BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""),
+        "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
+    })
 )
 
+def run_test(pytest_command: str):
+    """Helper function to run a test suite with custom pytest command"""
+    import subprocess
+    import sys
+    import os
+    
+    git_repo = os.environ.get("BUILDKITE_REPO")
+    git_commit = os.environ.get("BUILDKITE_COMMIT")
+    
+    print(f"Cloning repository: {git_repo}")
+    print(f"Checking out commit: {git_commit}")
+    
+    command = f"""
+    source $HOME/.local/bin/env &&
+    source /opt/venv/bin/activate &&
+    git clone {git_repo} /FastVideo &&
+    cd /FastVideo &&
+    git checkout {git_commit} &&
+    uv pip install -e .[test] &&
+    {pytest_command}
+    """
+    
+    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_encoder_tests():
-    """Run encoder tests on L40S GPU"""
+    run_test("pytest ./fastvideo/v1/tests/encoders -vs")
-    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"""
+    run_test("pytest ./fastvideo/v1/tests/vaes -vs")
-    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"""
+    run_test("pytest ./fastvideo/v1/tests/transformers -vs")
-    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"""
+    run_test("pytest ./fastvideo/v1/tests/ssim -vs")
-    import subprocess
+
-    import sys
+@app.function(gpu="L40S:4", image=image, timeout=1800, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
-    import os
+def run_training_tests():
-    
+    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/Vanilla -srP")
-    os.chdir("/FastVideo")
+
-    
+@app.function(gpu="H100:1", image=image, timeout=1800, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
-    command = """
+def run_training_tests_VSA():
-    source /opt/venv/bin/activate && 
+    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/VSA -srP")
-    pytest ./fastvideo/v1/tests/ssim -vs
+
-    """
+@app.function(gpu="H100:1", image=image, timeout=1800)
-    
+def run_inference_tests_STA():
-    result = subprocess.run([
+    run_test("pytest ./fastvideo/v1/tests/inference/STA -srP")
-        "/bin/bash", "-c", command
+
-    ], stdout=sys.stdout, stderr=sys.stderr, check=False)
+@app.function(gpu="H100:1", image=image, timeout=1800)
-    
+def run_precision_tests_STA():
-    sys.exit(result.returncode)
+    run_test("python csrc/attn/tests/test_sta.py")
+
+@app.function(gpu="H100:1", image=image, timeout=1800)
+def run_precision_tests_VSA():
+    run_test("python csrc/attn/tests/test_block_sparse.py")

fastvideo/v1/tests/nightly/test_e2e_i2v_overfit_single_sample.py

@@ -122,7 +122,7 @@ def run_training():
         "--checkpoints_total_limit", "3",
         "--allow_tf32",
         "--ema_start_step", "0",
-        "--cfg", "0.0",
+        "--training_cfg_rate", "0.1",
         "--output_dir", LOCAL_OUTPUT_DIR,
         "--tracker_project_name", "wan_i2v_finetune_overfit_ci",
         "--num_height", "480",

fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py

@@ -31,12 +31,9 @@ LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs"))
 def download_data():
     # create the data dir if it doesn't exist
     data_dir = Path(DATA_DIR)
-    if data_dir.exists():
-        print(f"Removing existing data directory at {data_dir}")
-        shutil.rmtree(data_dir)
 
     print(f"Creating data directory at {data_dir}")
-    os.makedirs(data_dir)
+    os.makedirs(data_dir, exist_ok=True)
 
     print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
     try:
@@ -122,7 +119,7 @@ def run_training():
         "--checkpoints_total_limit", "3",
         "--allow_tf32",
         "--ema_start_step", "0",
-        "--cfg", "0.0",
+        "--training_cfg_rate", "0.0",
         "--output_dir", LOCAL_OUTPUT_DIR,
         "--tracker_project_name", "wan_finetune_overfit_ci",
         "--num_height", "480",

fastvideo/v1/tests/ssim/README.md

@@ -1,13 +1,14 @@
-The reference videos in the `reference_videos` directory are used as part of an e2e test to ensure consistency in video generation quality across code changes. `test_inference_similarity.py` compares newly generated videos against these references using Structural Similarity Index (SSIM) metrics to detect any regressions in visual quality across code changes.
+The reference videos in the `*_reference_videos` directory are used as part of an e2e test to ensure consistency in video generation quality across code changes. `test_inference_similarity.py` compares newly generated videos against these references using Structural Similarity Index (SSIM) metrics to detect any regressions in visual quality across code changes.
 
-`reference_videos/FastHunyuan-diffusers/FLASH_ATTN/` videos were generated on commit `66107fd5b8469fed25972feb632cd48887dac451`.
+`A40_reference_videos` are generated on A40s and so on.
-`reference_videos/FastHunyuan-diffusers/TORCH_SDPA/` videos were generated on commit `4ea008b8a16d7f5678a44b187ebdd7d9d0416ff1`.
+
-`reference_videos/Wan2.1-T2V-1.3B-Diffusers` videos were generated on commit `d085770a70988c7b26632a0c3123c24a57f7ca77`.
+run `bash update_reference_videos.sh` from inside the `fastvideo/v1/tests/ssim/` directory after running `test_inference_similarity.py` to update reference videos. Note: make sure to update the path to the corresponding device.
-`reference_videos/Wan2.1-I2V-14B-480P-Diffusers` videos were generated on commit `d085770a70988c7b26632a0c3123c24a57f7ca77`.
+
+all reference videos are were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
 
 ## Generation Details
 
-2 x NVIDIA A40 GPUs
+2 x NVIDIA L40S GPUs
 
 ## Generation Parameters
 

fastvideo/v1/tests/ssim/test_inference_similarity.py

@@ -2,6 +2,7 @@
 import json
 import os
 
+import torch
 import pytest
 
 from fastvideo import VideoGenerator
@@ -11,6 +12,14 @@ from fastvideo.v1.worker.multiproc_executor import MultiprocExecutor
 
 logger = init_logger(__name__)
 
+device_name = torch.cuda.get_device_name()
+device_reference_folder_suffix = '_reference_videos'
+
+if "A40" in device_name:
+    device_reference_folder = "A40" + device_reference_folder_suffix
+elif "L40S" in device_name:
+    device_reference_folder = "L40S" + device_reference_folder_suffix
+
 # Base parameters from the shell script
 HUNYUAN_PARAMS = {
     "num_gpus": 2,
@@ -188,8 +197,8 @@ def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     assert os.path.exists(
         output_dir), f"Output video was not generated at {output_dir}"
 
-    reference_folder = os.path.join(script_dir, 'reference_videos', model_id, ATTENTION_BACKEND)
+    reference_folder = os.path.join(script_dir, device_reference_folder, model_id, ATTENTION_BACKEND)
-    
+
     if not os.path.exists(reference_folder):
         logger.error("Reference folder missing")
         raise FileNotFoundError(
@@ -288,8 +297,8 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     assert os.path.exists(
         output_dir), f"Output video was not generated at {output_dir}"
 
-    reference_folder = os.path.join(script_dir, 'reference_videos', model_id, ATTENTION_BACKEND)
+    reference_folder = os.path.join(script_dir, device_reference_folder, model_id, ATTENTION_BACKEND)
-    
+
     if not os.path.exists(reference_folder):
         logger.error("Reference folder missing")
         raise FileNotFoundError(

fastvideo/v1/tests/ssim/update_reference_videos.sh

@@ -0,0 +1,63 @@
+#!/bin/bash
+
+# Script to update reference videos using videos from generated_videos directory
+# Both directories should exist in the same directory as this script
+
+set -e  # Exit on any error
+
+# Define directory paths
+GENERATED_DIR="generated_videos"
+REFERENCE_DIR="set_me_to_correct_path"
+
+# Colors for output
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+NC='\033[0m' # No Color
+
+echo -e "${YELLOW}Starting reference video update...${NC}"
+
+# Check if generated_videos directory exists
+if [ ! -d "$GENERATED_DIR" ]; then
+    echo -e "${RED}Error: $GENERATED_DIR directory not found!${NC}"
+    exit 1
+fi
+
+# Check if reference_videos directory exists
+if [ ! -d "$REFERENCE_DIR" ]; then
+    echo -e "${RED}Error: $REFERENCE_DIR directory not found!${NC}"
+    exit 1
+fi
+
+# Function to copy videos recursively
+copy_videos() {
+    local src_dir="$1"
+    local dst_dir="$2"
+    
+    # Find all video files in the source directory
+    find "$src_dir" -type f \( -name "*.mp4" -o -name "*.avi" -o -name "*.mov" -o -name "*.mkv" -o -name "*.webm" -o -name "*.flv" \) | while read -r video_file; do
+        # Get relative path from source directory
+        relative_path="${video_file#$src_dir/}"
+        
+        # Construct destination path
+        dst_file="$dst_dir/$relative_path"
+        
+        # Create destination directory if it doesn't exist
+        dst_file_dir=$(dirname "$dst_file")
+        mkdir -p "$dst_file_dir"
+        
+        # Copy the video file
+        echo -e "${GREEN}Copying: $relative_path${NC}"
+        cp "$video_file" "$dst_file"
+    done
+}
+
+# Perform the copy operation
+echo -e "${YELLOW}Copying videos from $GENERATED_DIR to $REFERENCE_DIR...${NC}"
+copy_videos "$GENERATED_DIR" "$REFERENCE_DIR"
+
+echo -e "${GREEN}Reference videos updated successfully!${NC}"
+
+# Show summary
+video_count=$(find "$GENERATED_DIR" -type f \( -name "*.mp4" -o -name "*.avi" -o -name "*.mov" -o -name "*.mkv" -o -name "*.webm" -o -name "*.flv" \) | wc -l)
+echo -e "${YELLOW}Total videos processed: $video_count${NC}"

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

@@ -31,7 +31,6 @@ def run_worker():
         "--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_preprocessed_path", "data/mini_dataset_i2v_VSA/validation_parquet_dataset",
         "--train_batch_size", "1",
@@ -54,7 +53,7 @@ def run_worker():
         "--checkpoints_total_limit", "3",
         "--allow_tf32",
         "--ema_start_step", "0",
-        "--cfg", "0.0",
+        "--training_cfg_rate", "0.0",
         "--output_dir", "data/wan_finetune_test_VSA",
         "--tracker_project_name", "wan_finetune_ci_VSA",
         "--wandb_run_name", wandb_name,

fastvideo/v1/tests/training/Vanilla/a40_reference_wandb_summary.json

@@ -1 +1 @@
-{"_timestamp":1.7496170016478686e+09,"validation_videos_8_steps":{"_type":"videos","count":5,"videos":[{"sha256":"d81aa715df0c3ba4db8b242bc10442844453f3684a45dc55b89ecd27b2414fc7","size":429642,"path":"media/videos/validation_videos_8_steps_0_d81aa715df0c3ba4db8b.mp4","_type":"video-file"},{"sha256":"cd72a3d513eca6b41b03b80e6fa044ce7219c35e969d2ca20b9cb48c91e585c6","size":477837,"path":"media/videos/validation_videos_8_steps_0_cd72a3d513eca6b41b03.mp4","_type":"video-file"},{"_type":"video-file","sha256":"43d47c211a69bf0be3544738e76e7d8fa58bb108c00cb21dc34eaad3c6ce7cc3","size":409419,"path":"media/videos/validation_videos_8_steps_0_43d47c211a69bf0be354.mp4"},{"_type":"video-file","sha256":"ea674ec9e200bc97563c9d87d9dc07110c3f42c0ab7277dd237ae96dd8f90a10","size":333966,"path":"media/videos/validation_videos_8_steps_0_ea674ec9e200bc97563c.mp4"},{"sha256":"d81aa715df0c3ba4db8b242bc10442844453f3684a45dc55b89ecd27b2414fc7","size":429642,"path":"media/videos/validation_videos_8_steps_0_d81aa715df0c3ba4db8b.mp4","_type":"video-file"}],"captions":false},"step_time":2.5065076276659966,"_wandb":{"runtime":53},"learning_rate":1e-06,"_step":5,"_runtime":53.172758961,"grad_norm":5.65625,"train_loss":0.3915919363498688,"avg_step_time":2.8116052336990833}
+{"_timestamp":1.7496170016478686e+09,"validation_videos_8_steps":{"_type":"videos","count":5,"videos":[{"sha256":"d81aa715df0c3ba4db8b242bc10442844453f3684a45dc55b89ecd27b2414fc7","size":429642,"path":"media/videos/validation_videos_8_steps_0_d81aa715df0c3ba4db8b.mp4","_type":"video-file"},{"sha256":"cd72a3d513eca6b41b03b80e6fa044ce7219c35e969d2ca20b9cb48c91e585c6","size":477837,"path":"media/videos/validation_videos_8_steps_0_cd72a3d513eca6b41b03.mp4","_type":"video-file"},{"_type":"video-file","sha256":"43d47c211a69bf0be3544738e76e7d8fa58bb108c00cb21dc34eaad3c6ce7cc3","size":409419,"path":"media/videos/validation_videos_8_steps_0_43d47c211a69bf0be354.mp4"},{"_type":"video-file","sha256":"ea674ec9e200bc97563c9d87d9dc07110c3f42c0ab7277dd237ae96dd8f90a10","size":333966,"path":"media/videos/validation_videos_8_steps_0_ea674ec9e200bc97563c.mp4"},{"sha256":"d81aa715df0c3ba4db8b242bc10442844453f3684a45dc55b89ecd27b2414fc7","size":429642,"path":"media/videos/validation_videos_8_steps_0_d81aa715df0c3ba4db8b.mp4","_type":"video-file"}],"captions":false},"step_time":2.5065076276659966,"_wandb":{"runtime":53},"learning_rate":1e-06,"_step":5,"_runtime":53.172758961,"grad_norm":0.408203125,"train_loss":0.07883700542151928,"avg_step_time":2.8116052336990833}

fastvideo/v1/tests/training/Vanilla/l40s_reference_wandb_summary.json

@@ -0,0 +1 @@
+{"step_time":5.501357046999999,"grad_norm":0.384765625,"train_loss":0.07890288904309273,"avg_step_time":5.831571423200001}

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

@@ -18,7 +18,8 @@ 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/Vanilla/reference_wandb_summary.json"
+a40_reference_wandb_summary_file = "fastvideo/v1/tests/training/Vanilla/a40_reference_wandb_summary.json"
+l40s_reference_wandb_summary_file = "fastvideo/v1/tests/training/Vanilla/l40s_reference_wandb_summary.json"
 
 NUM_NODES = "1"
 NUM_GPUS_PER_NODE = "4"
@@ -36,9 +37,8 @@ def run_worker():
         "--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/crush-smol_processed_t2v/combined_parquet_dataset",
-        "--data_path", "data/crush-smol_parq/combined_parquet_dataset",
+        "--validation_preprocessed_path", "data/crush-smol_processed_t2v/validation_parquet_dataset",
-        "--validation_preprocessed_path", "data/crush-smol_parq/validation_parquet_dataset",
         "--train_batch_size", "2",
         "--num_latent_t", "4",
         "--num_gpus", "4",
@@ -59,7 +59,7 @@ def run_worker():
         "--checkpoints_total_limit", "3",
         "--allow_tf32",
         "--ema_start_step", "0",
-        "--cfg", "0.0",
+        "--training_cfg_rate", "0.0",
         "--output_dir", "data/wan_finetune_test",
         "--tracker_project_name", "wan_finetune_ci",
         "--wandb_run_name", wandb_name,
@@ -83,12 +83,12 @@ def test_distributed_training():
     """Test the distributed training setup"""
     os.environ["WANDB_MODE"] = "online"
 
-    data_dir = Path("data/crush-smol_parq")
+    data_dir = Path("data/crush-smol_processed_t2v")
     
     if not data_dir.exists():
         print(f"Downloading test dataset to {data_dir}...")
         snapshot_download(
-            repo_id="PY007/crush-smol",
+            repo_id="wlsaidhi/crush-smol_processed_t2v",
             local_dir=str(data_dir),
             repo_type="dataset",
             local_dir_use_symlinks=False
@@ -109,12 +109,20 @@ def test_distributed_training():
 
     summary_file = 'wandb/latest-run/files/wandb-summary.json'
 
+    device_name = torch.cuda.get_device_name()
+    if "A40" in device_name:
+        reference_wandb_summary_file = a40_reference_wandb_summary_file
+    elif "L40S" in device_name:
+        reference_wandb_summary_file = l40s_reference_wandb_summary_file
+    else:
+        raise ValueError(f"Unknown device: {device_name}")
+
     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.2,
+        'grad_norm': 0.3,
         'step_time': 0.5,
         'train_loss': 0.0025
     }

fastvideo/v1/tests/transformers/test_hunyuanvideo.py

@@ -80,7 +80,10 @@ def test_hunyuanvideo_distributed():
 
     # Initialize with identical weights
     model = initialize_identical_weights(model, seed=42)
-    shard_model(model, cpu_offload=False, reshard_after_forward=True)
+    shard_model(model, cpu_offload=True,
+                reshard_after_forward=True,
+                fsdp_shard_conditions=model._fsdp_shard_conditions
+                )
     for n, p in chain(model.named_parameters(), model.named_buffers()):
         if p.is_meta:
             raise RuntimeError(
@@ -92,7 +95,7 @@ def test_hunyuanvideo_distributed():
 
     # Move to GPU based on local rank (0 or 1 for 2 GPUs)
     device = torch.device(f"cuda:0")
-    model = model.to(device)
+    model = model
 
     batch_size = 1
     seq_len = 3

fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py

@@ -62,7 +62,7 @@ def test_hunyuanvideo_distributed():
 
     precision_str = "bf16"
     args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         use_cpu_offload=False,
+                         use_cpu_offload=True,
                          pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str))
     args.device = torch.device(f"cuda:{LOCAL_RANK}")
 

fastvideo/v1/tests/transformers/test_wanvideo.py

@@ -34,12 +34,12 @@ def test_wan_transformer():
     precision = torch.bfloat16
     precision_str = "bf16"
     args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         use_cpu_offload=False,
+                         use_cpu_offload=True,
                          pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str))
     args.device = device
 
     loader = TransformerLoader()
-    model2 = loader.load(TRANSFORMER_PATH, "", args).to(device, dtype=precision)
+    model2 = loader.load(TRANSFORMER_PATH, "", args).to(dtype=precision)
 
     model1 = WanTransformer3DModel.from_pretrained(
         TRANSFORMER_PATH, device=device,

fastvideo/v1/tests/vaes/test_hunyuan_vae.py

@@ -28,8 +28,12 @@ MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
 VAE_PATH = os.path.join(MODEL_PATH, "vae")
 CONFIG_PATH = os.path.join(VAE_PATH, "config.json")
 
-# Latent generated on commit 250f0b916cebb18a1c15c4aae1a0b480604d066a with 1 x A40
+# Latent generated on commit d71a4ebffc2034922fc379568b6a6aa722f3744c with 1 x A40
-REFERENCE_LATENT = -105.51324462890625
+# torch 2.7.1
+A40_REFERENCE_LATENT = -106.22467041015625
+
+# Latent generated on commit 2b54068960c41d42221e8b8719a374b499855029 with 1 x L40S
+L40S_REFERENCE_LATENT = -158.32318115234375
 
 
 @pytest.mark.usefixtures("distributed_setup")
@@ -66,6 +70,14 @@ def test_hunyuan_vae():
         latent = model.encode(input_tensor).mean.double().sum().item()
 
     # Check if latents are similar
+    device_name = torch.cuda.get_device_name()
+    if "A40" in device_name:
+        REFERENCE_LATENT = A40_REFERENCE_LATENT
+    elif "L40S" in device_name:
+        REFERENCE_LATENT = L40S_REFERENCE_LATENT
+    else:
+        raise ValueError(f"Unknown device: {device_name}")
+
     diff_encoded_latents = abs(REFERENCE_LATENT - latent)
     logger.info(
         f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}"

fastvideo/v1/training/training_pipeline.py

@@ -24,9 +24,9 @@ from fastvideo.v1.configs.sample import SamplingParam
 from fastvideo.v1.dataset import build_parquet_map_style_dataloader
 from fastvideo.v1.dataset.dataloader.schema import (
     pyarrow_schema_t2v, pyarrow_schema_t2v_validation)
-from fastvideo.v1.distributed import (cleanup_dist_env_and_memory, get_sp_group,
+from fastvideo.v1.distributed import (cleanup_dist_env_and_memory,
-                                      get_torch_device, get_world_group,
+                                      get_local_torch_device, get_sp_group,
-                                      sequence_model_parallel_all_gather)
+                                      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
@@ -60,14 +60,14 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         raise RuntimeError(
             "create_pipeline_stages should not be called for training pipeline")
 
-    def set_schemas(self):
+    def set_schemas(self) -> None:
         self.train_dataset_schema = pyarrow_schema_t2v
         self.validation_dataset_schema = pyarrow_schema_t2v_validation
 
     def initialize_training_pipeline(self, training_args: TrainingArgs):
         logger.info("Initializing training pipeline...")
+        self.device = get_local_torch_device()
         self.training_args = training_args
-        self.device = get_torch_device()
         world_group = get_world_group()
         self.world_size = world_group.world_size
         self.global_rank = world_group.rank
@@ -80,8 +80,6 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         self.seed = training_args.seed
         assert self.transformer is not None
         self.set_schemas()
-        # self.train_dataset_schema = pyarrow_schema_t2v
-        # self.validation_dataset_schema = pyarrow_schema_t2v_validation
 
         self.transformer.requires_grad_(True)
         self.transformer.train()
@@ -117,6 +115,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             training_args.train_batch_size,
             parquet_schema=self.train_dataset_schema,
             num_data_workers=training_args.dataloader_num_workers,
+            cfg_rate=training_args.training_cfg_rate,
             drop_last=True,
             text_padding_length=training_args.pipeline_config.
             text_encoder_configs[0].arch_config.
@@ -158,6 +157,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         return training_batch
 
     def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.training_args is not None
         assert self.train_loader_iter is not None
         assert self.train_dataloader is not None
 
@@ -170,27 +170,19 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             # Get first batch of new epoch
             batch = next(self.train_loader_iter)
 
-        # latents, encoder_hidden_states, encoder_attention_mask, caption_text, extra_latents, infos = batch
+        # latents, encoder_hidden_states, encoder_attention_mask, infos = batch
-        # for key, value in batch.items():
-        #     if isinstance(value, torch.Tensor):
-        #         logger.info("key: %s, shape: %s", key, value.shape)
-        #     else:
-        #         logger.info("key: %s, value: %s", key, value)
-        # print("--------------------------------")
-        # logger.info("batch: %s", batch)
         latents = batch['vae_latent']
+        latents = latents[:, :, :self.training_args.num_latent_t]
         encoder_hidden_states = batch['text_embedding']
         encoder_attention_mask = batch['text_attention_mask']
-        # extra_latents = batch['extra_latents']
         infos = batch['info_list']
 
-        training_batch.latents = latents.to(get_torch_device(),
+        training_batch.latents = latents.to(get_local_torch_device(),
                                             dtype=torch.bfloat16)
         training_batch.encoder_hidden_states = encoder_hidden_states.to(
-            get_torch_device(), dtype=torch.bfloat16)
+            get_local_torch_device(), dtype=torch.bfloat16)
         training_batch.encoder_attention_mask = encoder_attention_mask.to(
-            get_torch_device(), dtype=torch.bfloat16)
+            get_local_torch_device(), dtype=torch.bfloat16)
-        # training_batch.extra_latents = extra_latents
         training_batch.infos = infos
 
         return training_batch
@@ -283,7 +275,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             "encoder_hidden_states":
             training_batch.encoder_hidden_states,
             "timestep":
-            training_batch.timesteps.to(get_torch_device(),
+            training_batch.timesteps.to(get_local_torch_device(),
                                         dtype=torch.bfloat16),
             "encoder_attention_mask":
             training_batch.encoder_attention_mask,
@@ -296,8 +288,10 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             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.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
@@ -321,9 +315,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             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
 
-        # if self.training_args.sp_size > 1:
+        # make sure no implicit broadcasting happens
-        #     model_pred = sequence_model_parallel_all_gather(model_pred, dim=2)
-
         assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
         loss = (torch.mean((model_pred.float() - target.float())**2) /
                 self.training_args.gradient_accumulation_steps)
@@ -368,20 +360,12 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         for _ in range(self.training_args.gradient_accumulation_steps):
             training_batch = self._get_next_batch(training_batch)
 
-            training_batch.latents = training_batch.latents[:, :, :self.
-                                                            training_args.
-                                                            num_latent_t]
-
             # Normalize DIT input
             training_batch = self._normalize_dit_input(training_batch)
-
+            # Create noisy model input
-            # Save original latents
-            # training_batch.original_latents = training_batch.latents.clone()
-
-
             training_batch = self._prepare_dit_inputs(training_batch)
 
-            # # Shard latents across sp groups
+            # Shard latents across sp groups
             training_batch.latents = shard_latents_across_sp(
                 training_batch.latents,
                 num_latent_t=self.training_args.num_latent_t)
@@ -389,10 +373,11 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             training_batch.noisy_model_input = shard_latents_across_sp(
                 training_batch.noisy_model_input,
                 num_latent_t=self.training_args.num_latent_t)
-            # CRITICAL FIX: Also shard noise to match latents
+            # shard noise to match latents
             training_batch.noise = shard_latents_across_sp(
                 training_batch.noise,
                 num_latent_t=self.training_args.num_latent_t)
+
             training_batch = self._build_attention_metadata(training_batch)
             training_batch = self._build_input_kwargs(training_batch)
             training_batch = self._transformer_forward_and_compute_loss(
@@ -563,14 +548,13 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             negative_prompt_attention_mask: torch.Tensor | None
     ) -> ForwardBatch:
 
-        # logger.info("validation_batch: %s", validation_batch)
-        # latents, embeddings, masks, caption_text, extra_latents, infos = validation_batch
         prompt = validation_batch['info_list'][0]['prompt']
         prompt_embeds = validation_batch['text_embedding']
         prompt_attention_mask = validation_batch['text_attention_mask']
 
-        prompt_embeds = prompt_embeds.to(get_torch_device())
+        prompt_embeds = prompt_embeds.to(get_local_torch_device())
-        prompt_attention_mask = prompt_attention_mask.to(get_torch_device())
+        prompt_attention_mask = prompt_attention_mask.to(
+            get_local_torch_device())
 
         # Calculate sizes
         latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
@@ -621,7 +605,6 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
 
         # Set deterministic seed for validation
         set_random_seed(self.seed)
-
         logger.info("Using validation seed: %s", self.seed)
 
         # Prepare validation prompts
@@ -632,13 +615,13 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             batch_size=1,
             parquet_schema=self.validation_dataset_schema,
             num_data_workers=0,
+            cfg_rate=0.0,
             drop_last=False,
-            drop_first_row=sampling_param.negative_prompt is not None,
+            drop_first_row=sampling_param.negative_prompt is not None)
-            cfg_rate=training_args.cfg)
         if sampling_param.negative_prompt:
             negative_prompt_embeds, negative_prompt_attention_mask, negative_prompt = validation_dataset.get_validation_negative_prompt(
             )
-            logger.info("negative_prompt: %s", negative_prompt)
+            logger.info("Using negative_prompt: %s", negative_prompt)
 
         transformer.eval()
 
@@ -651,13 +634,14 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             step_videos: List[np.ndarray] = []
             step_captions: List[str | None] = []
 
-            # for _, embeddings, masks, caption_text, extra_latents, infos in validation_dataloader:
             for validation_batch in validation_dataloader:
                 batch = self._prepare_validation_inputs(
                     sampling_param, training_args, validation_batch,
                     num_inference_steps, negative_prompt_embeds,
                     negative_prompt_attention_mask)
+
                 step_captions.extend([None])  # TODO(peiyuan): add caption
+
                 # Run validation inference
                 with torch.no_grad(), torch.autocast("cuda",
                                                      dtype=torch.bfloat16):

fastvideo/v1/training/wan_i2v_training_pipeline.py

@@ -8,7 +8,7 @@ import torch
 from fastvideo.v1.configs.sample import SamplingParam
 from fastvideo.v1.dataset.dataloader.schema import (
     pyarrow_schema_i2v, pyarrow_schema_i2v_validation)
-from fastvideo.v1.distributed import get_torch_device
+from fastvideo.v1.distributed import get_local_torch_device
 from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import (
@@ -18,7 +18,6 @@ from fastvideo.v1.pipelines.pipeline_batch_info import (ForwardBatch,
 from fastvideo.v1.pipelines.wan.wan_i2v_pipeline import (
     WanImageToVideoValidationPipeline)
 from fastvideo.v1.training.training_pipeline import TrainingPipeline
-from fastvideo.v1.training.training_utils import shard_latents_across_sp
 from fastvideo.v1.utils import is_vsa_available
 
 vsa_available = is_vsa_available()
@@ -64,7 +63,7 @@ class WanI2VTrainingPipeline(TrainingPipeline):
         self.validation_pipeline = validation_pipeline
 
     def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
-        assert self.train_loader_iter is not None
+        assert self.training_args is not None
         assert self.train_dataloader is not None
 
         batch = next(self.train_loader_iter, None)  # type: ignore
@@ -76,39 +75,29 @@ class WanI2VTrainingPipeline(TrainingPipeline):
             # Get first batch of new epoch
             batch = next(self.train_loader_iter)
 
-        # latents, encoder_hidden_states, encoder_attention_mask, caption_text, extra_latents, infos = batch
-        # for key, value in batch.items():
-        #     if isinstance(value, torch.Tensor):
-        #         logger.info("key: %s, shape: %s", key, value.shape)
-        #     else:
-        #         logger.info("key: %s, value: %s", key, value)
-        # print("--------------------------------")
-        # logger.info("batch: %s", batch)
         latents = batch['vae_latent']
+        latents = latents[:, :, :self.training_args.num_latent_t]
         encoder_hidden_states = batch['text_embedding']
         encoder_attention_mask = batch['text_attention_mask']
         clip_features = batch['clip_feature']
         image_latents = batch['first_frame_latent']
+        image_latents = image_latents[:, :, :self.training_args.num_latent_t]
         pil_image = batch['pil_image']
-        # extra_latents = batch['extra_latents']
         infos = batch['info_list']
 
-        training_batch.latents = latents.to(get_torch_device(),
+        training_batch.latents = latents.to(get_local_torch_device(),
                                             dtype=torch.bfloat16)
         training_batch.encoder_hidden_states = encoder_hidden_states.to(
-            get_torch_device(), dtype=torch.bfloat16)
+            get_local_torch_device(), dtype=torch.bfloat16)
         training_batch.encoder_attention_mask = encoder_attention_mask.to(
-            get_torch_device(), dtype=torch.bfloat16)
+            get_local_torch_device(), dtype=torch.bfloat16)
-        training_batch.preprocessed_image = pil_image.to(get_torch_device())
+        training_batch.preprocessed_image = pil_image.to(
-        training_batch.image_embeds = clip_features.to(get_torch_device())
+            get_local_torch_device())
-        training_batch.image_latents = image_latents.to(get_torch_device())
+        training_batch.image_embeds = clip_features.to(get_local_torch_device())
-        # training_batch.extra_latents = extra_latents
+        training_batch.image_latents = image_latents.to(
+            get_local_torch_device())
         training_batch.infos = infos
 
-        assert training_batch.image_latents is not None
-        # if self.training_args and hasattr(self.training_args, 'num_latent_t'):
-        training_batch.image_latents = training_batch.image_latents[:, :, :self.training_args.num_latent_t]
-
         return training_batch
 
     def _prepare_dit_inputs(self,
@@ -123,12 +112,14 @@ class WanI2VTrainingPipeline(TrainingPipeline):
 
         # First, call parent method to prepare noise, timesteps, etc. for video latents
         training_batch = super()._prepare_dit_inputs(training_batch)
-        
+
-        image_latents = training_batch.image_latents.to(get_torch_device(), dtype=torch.bfloat16)
+        assert isinstance(training_batch.image_latents, torch.Tensor)
-        
+        image_latents = training_batch.image_latents.to(
+            get_local_torch_device(), dtype=torch.bfloat16)
+
         training_batch.noisy_model_input = torch.cat(
             [training_batch.noisy_model_input, image_latents], dim=1)
-        
+
         return training_batch
 
     def _build_input_kwargs(self,
@@ -143,7 +134,8 @@ class WanI2VTrainingPipeline(TrainingPipeline):
         # Image Embeds for conditioning
         image_embeds = training_batch.image_embeds
         assert torch.isnan(image_embeds).sum() == 0
-        image_embeds = image_embeds.to(get_torch_device(), dtype=torch.bfloat16)
+        image_embeds = image_embeds.to(get_local_torch_device(),
+                                       dtype=torch.bfloat16)
         encoder_hidden_states_image = image_embeds
 
         # NOTE: noisy_model_input already contains concatenated image_latents from _prepare_dit_inputs
@@ -153,7 +145,7 @@ class WanI2VTrainingPipeline(TrainingPipeline):
             "encoder_hidden_states":
             training_batch.encoder_hidden_states,
             "timestep":
-            training_batch.timesteps.to(get_torch_device(),
+            training_batch.timesteps.to(get_local_torch_device(),
                                         dtype=torch.bfloat16),
             "encoder_attention_mask":
             training_batch.encoder_attention_mask,
@@ -170,34 +162,16 @@ class WanI2VTrainingPipeline(TrainingPipeline):
             negative_prompt_embeds: torch.Tensor | None,
             negative_prompt_attention_mask: torch.Tensor | None
     ) -> ForwardBatch:
-        # latents, embeddings, masks, caption_text, extra_latents, infos = validation_batch
-        # latents = validation_batch['vae_latent']
         embeddings = validation_batch['text_embedding']
         masks = validation_batch['text_attention_mask']
         clip_features = validation_batch['clip_feature']
-        # extra_latents = validation_batch['extra_latents']
         pil_image = validation_batch['pil_image']
         infos = validation_batch['info_list']
         prompt = infos[0]['prompt']
 
-        prompt_embeds = embeddings.to(get_torch_device())
+        prompt_embeds = embeddings.to(get_local_torch_device())
-        prompt_attention_mask = masks.to(get_torch_device())
+        prompt_attention_mask = masks.to(get_local_torch_device())
-        clip_features = clip_features.to(get_torch_device())
+        clip_features = clip_features.to(get_local_torch_device())
-        # clip_features = extra_latents.get("clip_feature")
-        # first_frame_latent = extra_latents.get("first_frame_latent")
-        # pil_image = extra_latents.get("pil_image")
-
-        # if clip_features is not None and clip_features.numel() > 0:
-        #     clip_features = clip_features.to(get_torch_device())
-        # if first_frame_latent is not None and first_frame_latent.numel() > 0:
-        #     first_frame_latent = first_frame_latent.to(get_torch_device())
-        # if pil_image is not None and pil_image[0] is not None and pil_image[
-        #         0].numel() > 0:
-        #     pil_image = pil_image[0].to(get_torch_device())
-        # else:
-        #     clip_features = None
-        #     first_frame_latent = None
-        #     pil_image = None
 
         # Calculate sizes
         latents_size = [(sampling_param.num_frames - 1) // 4 + 1,

pyproject.toml

@@ -1,3 +1,4 @@
+# trigger test
 [build-system]
 requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"
@@ -19,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",
+    "timm==1.0.11", "peft>=0.15.0", "diffusers>=0.33.1", "bitsandbytes",
     "torch==2.7.1", "torchvision",
 
     # Acceleration & Optimization

scripts/finetune/finetune_v1.sh

@@ -36,7 +36,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --checkpoints_total_limit 3\
     --allow_tf32\
     --ema_start_step 0\
-    --cfg 0.0\
+    --training_cfg_rate 0.0\
     --output_dir="$DATA_DIR/outputs/wan_finetune"\
     --tracker_project_name wan_finetune \
     --num_height 480 \

scripts/finetune/finetune_v1_VSA.sh

@@ -42,7 +42,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
     --checkpoints_total_limit 3 \
     --allow_tf32 \
     --ema_start_step 0 \
-    --cfg 0.0 \
+    --training_cfg_rate 0.0 \
     --output_dir "$DATA_DIR/outputs/wan_finetune" \
     --tracker_project_name VSA_finetune \
     --num_height 448 \

scripts/inference/v1_inference_wan_VSA.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+num_gpus=1
+export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
+# change model path to local dir if you want to inference using your checkpoint
+export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
+# Note that the tp_size and sp_size should be the same and equal to the number
+# of GPUs. They are used for different parallel groups. sp_size is used for
+# dit model and tp_size is used for encoder models.
+fastvideo generate \
+    --model-path $MODEL_BASE \
+    --sp-size $num_gpus \
+    --tp-size $num_gpus \
+    --num-gpus $num_gpus \
+    --height 448 \
+    --width 832 \
+    --num-frames 77 \
+    --num-inference-steps 50 \
+    --fps 16 \
+    --guidance-scale 6.0 \
+    --flow-shift 8.0 \
+    --VSA-sparsity 0.9 \
+    --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 1024 \
+    --output-path outputs_video_1.3B_VSA/sparsity_0.9/