[Misc] clean up VSA finetuning examples. (#1111)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-02-18 11:37:20 -08:00
parent 8e39ce38c9
commit 76d3ba69e0
9 changed files with 224 additions and 87 deletions
@@ -44,8 +44,8 @@ VALIDATION_DATASET_FILE=your_validation_dataset_file
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

-# Training arguments
-training_args=(
+# Core training arguments
+core_training_args=(
  --tracker_project_name wan_i2v_VSA
  --output_dir "checkpoints/wan_i2v_finetune_VSA"
  --max_train_steps 4000
@@ -53,10 +53,14 @@ training_args=(
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
  --num_latent_t 21
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Validation generation shape arguments (used during training validation)
+validation_generation_shape_args=(
  --num_height 720
  --num_width 1280
  --num_frames 81
-  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
@@ -127,8 +131,9 @@ srun torchrun \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
-    "${training_args[@]}" \
+    "${core_training_args[@]}" \
+    "${validation_generation_shape_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${vsa_args[@]}"  
+    "${vsa_args[@]}"  
@@ -6,9 +6,7 @@ These are e2e example scripts for finetuning Wan2.1 T2V with VSA to accelerate i

 ## Make sure you have installed VSA

-```bash
-pip install vsa
-```
+Go to fastvideo-kernel/README.md for instructions.

 ### Download the synthetic dataset:

@@ -44,8 +44,8 @@ VALIDATION_DATASET_FILE=your_validation_dataset_file
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

-# Training arguments
-training_args=(
+# Core training arguments
+core_training_args=(
  --tracker_project_name wan_t2v_VSA
  --output_dir "checkpoints/wan_t2v_finetune_VSA"
  --max_train_steps 4000
@@ -53,10 +53,14 @@ training_args=(
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
  --num_latent_t 21
+  # --enable_gradient_checkpointing_type "full" # if OOM enable this
+)
+
+# Validation generation shape arguments (used during training validation)
+validation_generation_shape_args=(
  --num_height 480
  --num_width 832
  --num_frames 81
-  # --enable_gradient_checkpointing_type "full" # if OOM enable this
 )

 # Parallel arguments
@@ -127,8 +131,9 @@ srun torchrun \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
-    "${training_args[@]}" \
+    "${core_training_args[@]}" \
+    "${validation_generation_shape_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${vsa_args[@]}"  
+    "${vsa_args[@]}"  
@@ -0,0 +1,114 @@
+
+
+# Basic info
+export WANDB_MODE="online"
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export TRITON_CACHE_DIR="/tmp/triton_cache_${USER}_$$"
+export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+export MASTER_PORT="${MASTER_PORT:-29500}"
+export NODE_RANK=0
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
+export WANDB_API_KEY="your_wandb_api_key_here" # TODO: Replace with your actual key or load from a secure location
+
+# Configs
+NUM_GPUS=4
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR=data/Wan-Syn_77x448x832_600k
+VALIDATION_DATASET_FILE=examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
+
+# Core training arguments
+core_training_args=(
+  --tracker_project_name wan_t2v_VSA
+  --output_dir "checkpoints/wan_t2v_finetune_VSA"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 20
+  --enable_gradient_checkpointing_type "full" # if OOM enable this
+)
+
+# Validation generation shape arguments (used during training validation)
+validation_generation_shape_args=(
+  --num_height 448
+  --num_width 832
+  --num_frames 77
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim $NUM_GPUS
+  --hsdp_shard_dim 1
+)
+
+# 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 4
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 200
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "5.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 1e-6
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
+  --weight_decay 0.01
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --dit_precision "fp32"
+  --ema_start_step 0
+  --flow_shift 1
+  --seed 1000
+)
+
+# VSA arguments
+vsa_args=(
+  --VSA_decay_rate 0.03
+  --VSA_decay_interval_steps 50
+  --VSA_sparsity 0.9
+)
+
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node "$NUM_GPUS" \
+  --node_rank "$NODE_RANK" \
+  --rdzv_backend c10d \
+  --rdzv_endpoint "$MASTER_ADDR:$MASTER_PORT" \
+  fastvideo/training/wan_training_pipeline.py \
+  "${parallel_args[@]}" \
+  "${model_args[@]}" \
+  "${dataset_args[@]}" \
+  "${core_training_args[@]}" \
+  "${validation_generation_shape_args[@]}" \
+  "${optimizer_args[@]}" \
+  "${validation_args[@]}" \
+  "${miscellaneous_args[@]}" \
+  "${vsa_args[@]}"
@@ -44,8 +44,8 @@ VALIDATION_DATASET_FILE=your_validation_dataset_file
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

-# Training arguments
-training_args=(
+# Core training arguments
+core_training_args=(
  --tracker_project_name wan_t2v_VSA
  --output_dir "checkpoints/wan_t2v_finetune_VSA"
  --max_train_steps 4000
@@ -53,10 +53,14 @@ training_args=(
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
  --num_latent_t 21
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Validation generation shape arguments (used during training validation)
+validation_generation_shape_args=(
  --num_height 720
  --num_width 1280
  --num_frames 81
-  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
@@ -127,8 +131,9 @@ srun torchrun \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
-    "${training_args[@]}" \
+    "${core_training_args[@]}" \
+    "${validation_generation_shape_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${vsa_args[@]}"  
+    "${vsa_args[@]}"  
@@ -1,7 +1,8 @@
 #!/bin/bash
+# Note: for debugging, you do not need to download them all and just control c.
+

 # 480P dataset
-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"
-
+python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "data/Wan-Syn_77x448x832_600k" --repo_type "dataset"
 # 720P dataset
-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x768x1280_250k" --local_dir "FastVideo/Wan-Syn_77x768x1280_250k" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x768x1280_250k" --local_dir "data/Wan-Syn_77x768x1280_250k" --repo_type "dataset"
@@ -0,0 +1,36 @@
+{
+  "data": [
+    {
+      "caption": "In the video, a woman is elegantly showcasing her earrings, bringing attention to their intricate design with a gentle touch of her fingers. She is bathed in ambient purple and pink lighting, which casts a soft glow on her delicate features and enhances the vivid tones of her lipstick and eye makeup. Her hair is styled to frame her face smoothly, emphasizing the contours of her jawline and cheekbones. The background features a blurred neon light, adding an artistic and modern touch to the overall aesthetic.",
+      "video_path": "Fashion/mixkit-face-of-an-elegant-and-captivating-woman-41914_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.",
+      "video_path": "Man/mixkit-a-rancher-riding-a-horse-at-sunset-1143_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a dimly lit, eerie setting, a mysterious pink bottle labeled \"Authentic 100% organic POISON\" sits prominently in the foreground, casting a menacing aura. The bottle is accentuated by green fog, which swirls lightly around it, enhancing its sinister allure. Behind it, a shadowy golden bottle adorned with a spider emblem subtly emerges, adding an extra layer of mystery to the scene. Dim candles provide faint, flickering light, which complements the dark atmosphere, making the setting ideal for an illusion of hidden dangers.",
+      "video_path": "smoke/mixkit-poison-in-halloween-ritual-33879_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The video opens with a tranquil scene in the heart of a dense forest, emphasizing two large, textured tree trunks in the foreground framing the view. Sunlight filters through the canopy above, casting intricate patterns of light and shadow on the trees and the ground. Between the tree trunks, a clear view of a calm, muddy river unfolds, its surface shimmering under the gentle sunlight. The riverbank is decorated with a variety of small bushes and vibrant foliage, subtly transitioning into the deep greens of tall, leafy plants. In the background, the dense forest looms, filled with dark, towering trees, their branches intertwining to form an intricate canopy. The scene is bathed in the soft glow of the sun, creating a serene and picturesque setting. Occasional sunbeams pierce through the foliage, adding a magical aura to the landscape. The vibrant reds and oranges of the smaller plants add contrast, bringing warmth to the earthy tones of the scenery. Overall, this harmonious blend of natural elements creates a peaceful and idyllic forest setting.",
+      "video_path": "forest/mixkit-view-of-a-river-between-two-old-trees-560_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    }
+  ]
+}
@@ -95,6 +95,26 @@ class DP_SP_BatchSampler(Sampler[list[int]]):


 def get_parquet_files_and_length(path: str):
+    path = os.path.abspath(path)
+    if not os.path.exists(path):
+        raise FileNotFoundError(
+            f"Dataset path does not exist: {path}. "
+            "Expected a directory containing *.parquet files.")
+    if not os.path.isdir(path):
+        raise NotADirectoryError(
+            f"Dataset path is not a directory: {path}. "
+            "Expected a directory containing *.parquet files.")
+    has_parquet_file = False
+    for _, _, files in os.walk(path):
+        if any(file.endswith(".parquet") for file in files):
+            has_parquet_file = True
+            break
+    if not has_parquet_file:
+        raise FileNotFoundError(
+            "No parquet files found under dataset path: "
+            f"{path}. "
+            "Please verify this path points to preprocessed parquet data.")
+
    # Check if cached info exists
    cache_dir = os.path.join(path, "map_style_cache")
    cache_file = os.path.join(cache_dir, "file_info.pkl")
@@ -111,6 +131,10 @@ def get_parquet_files_and_length(path: str):
            try:
                with open(cache_file, "rb") as f:
                    file_names_sorted, lengths_sorted = pickle.load(f)
+                if len(file_names_sorted) == 0:
+                    raise ValueError(
+                        "Cached parquet metadata is empty in "
+                        f"{cache_file}.")
                cache_loaded = True
                logger.info("Successfully loaded cached file info")
            except Exception as e:
@@ -128,6 +152,12 @@ def get_parquet_files_and_length(path: str):
                    if file.endswith('.parquet'):
                        file_path = os.path.join(root, file)
                        file_names.append(file_path)
+            if len(file_names) == 0:
+                raise FileNotFoundError(
+                    "No parquet files found under dataset path: "
+                    f"{path}. "
+                    "Please verify this path points to preprocessed parquet "
+                    "data.")
            for file_path in tqdm.tqdm(
                    file_names, desc="Reading parquet files to get lengths"):
                num_rows = pq.ParquetFile(file_path).metadata.num_rows
@@ -138,9 +168,6 @@ def get_parquet_files_and_length(path: str):
                                                                strict=True),
                                                            key=lambda x: x[0]),
                                                    strict=True)
-            assert len(
-                file_names_sorted) != 0, "No parquet files found in the dataset"
-
            # Save the cache
            os.makedirs(cache_dir, exist_ok=True)
            with open(cache_file, "wb") as f:
@@ -155,6 +182,15 @@ def get_parquet_files_and_length(path: str):
    logger.info("Loading cached file info from %s after barrier", cache_file)
    with open(cache_file, "rb") as f:
        file_names_sorted, lengths_sorted = pickle.load(f)
+    if len(file_names_sorted) == 0:
+        raise RuntimeError(
+            "Cached parquet metadata is empty after synchronization at "
+            f"{cache_file}. "
+            "Please verify the dataset path and regenerate cache.")
+    if len(file_names_sorted) != len(lengths_sorted):
+        raise RuntimeError(
+            "Cached parquet metadata is corrupted at "
+            f"{cache_file}: file count and length count do not match.")

    return file_names_sorted, lengths_sorted

@@ -1,63 +0,0 @@
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=offline
-export WANDB_API_KEY='your_wandb_api_key'
-export TOKENIZERS_PARALLELISM=false
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-export TRITON_CACHE_DIR=/tmp/triton_cache
-DATA_DIR=~/train/
-VALIDATION_DATASET_FILE=[your validation dataset file]
-NUM_GPUS=8
-export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
-# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-CHECKPOINT_PATH="$DATA_DIR/outputs/wan_finetune/checkpoint-5"
-
-# Make sure that num_latent_t is a multiple of sp_size
-torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
-    fastvideo/training/wan_training_pipeline.py \
-    --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/ray/.cache" \
-    --data_path "$DATA_DIR" \
-    --validation_dataset_file "$VALIDATION_DATASET_FILE" \
-    --train_batch_size 1 \
-    --num_latent_t 16 \
-    --sp_size 1 \
-    --tp_size 1 \
-    --num_gpus $NUM_GPUS \
-    --hsdp_replicate_dim $NUM_GPUS \
-    --hsdp-shard-dim 1 \
-    --train_sp_batch_size 1 \
-    --dataloader_num_workers 4 \
-    --gradient_accumulation_steps 8 \
-    --max_train_steps 30000 \
-    --learning_rate 1e-6 \
-    --mixed_precision "bf16" \
-    --training_state_checkpointing_steps 6000 \
-    --validation_steps 100 \
-    --validation_sampling_steps "50" \
-    --log_validation \
-    --checkpoints_total_limit 3 \
-    --ema_start_step 0 \
-    --training_cfg_rate 0.0 \
-    --output_dir "$DATA_DIR/outputs/wan_finetune" \
-    --tracker_project_name VSA_finetune \
-    --num_height 448 \
-    --num_width 832 \
-    --num_frames  61 \
-    --flow_shift 3 \
-    --validation_guidance_scale "5.0" \
-    --num_euler_timesteps 50 \
-    --master_weight_type "fp32" \
-    --dit_precision "fp32" \
-    --weight_decay 0.01 \
-    --max_grad_norm 1.0 \
-    --VSA_decay_sparsity 0.9 \
-    --VSA_decay_rate 0.03 \
-    --VSA_decay_interval_steps 30 \
-    --VSA_val_sparsity 0.9 \
-    --enable_gradient_checkpointing_type "full"
-# --resume_from_checkpoint "$CHECKPOINT_PATH"