[Bug] Loss plateau (~0.4) and blurry validation when using Video Sparse Attention (VSA) on Wan2.1-1.3B

[xiaoxiao0406]

Describe the bug

I am experiencing training issues when enabling Video Sparse Attention (VSA) for finetuning Wan2.1-T2V-1.3B.

• The Problem: With VSA enabled, the loss plateaus around 0.4 and is difficult to decrease further. The validation videos are completely blurry.
• Comparison: Under the exact same hyperparameter settings, if I use the default attention backend (Flash Attention), both the loss curve and validation results are normal.
• Dataset: I am using the official FastVideo Synthetic Wan2.1 480P dataset.

Loss Curve:

Validation Video(3k steps):

68202673-3da3-4020-8405-bab519ac9590.mp4

Reproduction

#!/bin/bash
set -e -x

# Environment Setup
# source ~/conda/miniconda/bin/activate
# conda activate your_env

# Basic Info
# export SLURM_PROCID=$SLURM_NODEID
# export WANDB_MODE="online"
export NCCL_P2P_DISABLE=1
export NCCL_IB_DISABLE=1
export TORCH_NCCL_ENABLE_MONITORING=0
# different cache dir for different processes
export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
# mkdir -p /mnt/petrelfs/wangyating/tmp/triton_cache_${SLURM_PROCID}

# export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA

MASTER_PORT=${MASTER_PORT:-29506}
NODE_RANK=${SLURM_PROCID:-0}
NUM_GPUS_PER_NODE=8
NUM_NODES=${SLURM_JOB_NUM_NODES:-1}
WORLD_SIZE=$((NUM_NODES * NUM_GPUS_PER_NODE))
export MASTER_PORT=29506
export NODE_RANK=$SLURM_PROCID
if [[ -n "${SLURM_JOB_NODELIST}" ]]; then
  nodes=( $(scontrol show hostnames "$SLURM_JOB_NODELIST") )
  MASTER_ADDR=${nodes[0]}
else
  MASTER_ADDR=${MASTER_ADDR:-127.0.0.1}
fi

echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"

# Configs
NUM_GPUS=8
MODEL_PATH="/mnt/inspurfs/eb3d_t/share/hf_models/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR="/mnt/inspurfs/eb3d_t/wangyating/datasets"
VALIDATION_DATASET_FILE="/mnt/petrelfs/wangyating/projects/multimodal/FastVideo/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_64.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]

# Training arguments
training_args=(
  --tracker_project_name wan_t2v_finetune
  --output_dir "checkpoints/wan_t2v_finetune_gpu_32-export_attention"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
  --gradient_accumulation_steps 2
  --num_latent_t 16
  --num_height 448
  --num_width 832
  --num_frames 61
  --enable_gradient_checkpointing_type "full" # if OOM enable this
)

# Parallel arguments
parallel_args=(
  --num_gpus $WORLD_SIZE
  --sp_size 1
  --tp_size 1
  --hsdp_replicate_dim $WORLD_SIZE
  --hsdp_shard_dim 1
)

# Model arguments
model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path "/mnt/inspurfs/eb3d_t/share/hf_models/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
)

# 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 1000
  --validation_sampling_steps "50"
  --validation_guidance_scale "5.0"
  --wandb_run_name "wan_t2v_finetune_gpu_32-export_attention"
)

# Optimizer arguments
optimizer_args=(
  --learning_rate 1e-5
  --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
  --dit_layerwise_offload False
  --vae_cpu_offload False
)

# VSA arguments
vsa_args=(
  --VSA_decay_rate 0.03 \
  --VSA_decay_interval_steps 50 \
  --VSA_sparsity 0.9 \
)

torchrun \
  --nnodes $NUM_NODES \
  --nproc_per_node $NUM_GPUS_PER_NODE \
  --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[@]}" \
    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
    "${vsa_args[@]}"  

Environment

• Hardware: 32x NVIDIA A800 GPUs
• Python: 3.12.12
• PyTorch Version: 2.6.0
• CUDA Version: 12.4
• Model: Wan2.1-T2V-1.3B (Diffusers format)

PyTorch version: 2.6.0+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 11.2.0
Clang version: Could not collect
CMake version: version 3.30.9
Libc version: glibc-2.17

Python version: 3.12.12 | packaged by conda-forge | (main, Jan 26 2026, 23:51:32) [GCC 14.3.0] (64-bit runtime)
Python platform: Linux-3.10.0-957.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A800-SXM4-80GB
GPU 1: NVIDIA A800-SXM4-80GB
GPU 2: NVIDIA A800-SXM4-80GB
GPU 3: NVIDIA A800-SXM4-80GB
GPU 4: NVIDIA A800-SXM4-80GB
GPU 5: NVIDIA A800-SXM4-80GB
GPU 6: NVIDIA A800-SXM4-80GB
GPU 7: NVIDIA A800-SXM4-80GB

Nvidia driver version: 550.90.07
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:          x86_64
CPU op-mode(s):        32-bit, 64-bit
Byte Order:            Little Endian
CPU(s):                128
On-line CPU(s) list:   0-127
Thread(s) per core:    2
Core(s) per socket:    32
Socket(s):             2
NUMA node(s):          2
Vendor ID:             GenuineIntel
CPU family:            6
Model:                 106
Model name:            Intel(R) Xeon(R) Platinum 8358P CPU @ 2.60GHz
Stepping:              6
CPU MHz:               3199.853
CPU max MHz:           3400.0000
CPU min MHz:           800.0000
BogoMIPS:              5200.00
Virtualization:        VT-x
L1d cache:             48K
L1i cache:             32K
L2 cache:              1280K
L3 cache:              49152K
NUMA node0 CPU(s):     0-31,64-95
NUMA node1 CPU(s):     32-63,96-127
Flags:                 fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc aperfmperf eagerfpu pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch epb cat_l3 intel_pt ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq spec_ctrl intel_stibp flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] accelerate==1.0.1
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-cusparselt-cu12==0.6.2
[pip3] nvidia-ml-py==13.590.48
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] peft==0.18.1
[pip3] torch==2.6.0+cu124
[pip3] torchaudio==2.6.0+cu124
[pip3] torchdata==0.11.0
[pip3] torchvision==0.21.0+cu124
[pip3] transformers==4.57.3
[pip3] triton==3.2.0
[conda] accelerate                1.0.1                    pypi_0    pypi
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] nvidia-cublas-cu12        12.4.5.8                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.4.127                 pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.1.0.70                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.2.1.3                 pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.5.147               pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.6.1.9                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.3.1.170               pypi_0    pypi
[conda] nvidia-cusparselt-cu12    0.6.2                    pypi_0    pypi
[conda] nvidia-ml-py              13.590.48                pypi_0    pypi
[conda] nvidia-nccl-cu12          2.21.5                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.4.127                 pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.4.127                 pypi_0    pypi
[conda] peft                      0.18.1                   pypi_0    pypi
[conda] torch                     2.6.0+cu124              pypi_0    pypi
[conda] torchaudio                2.6.0+cu124              pypi_0    pypi
[conda] torchcodec                0.2.1                    pypi_0    pypi
[conda] torchdata                 0.11.0                   pypi_0    pypi
[conda] torchvision               0.21.0+cu124             pypi_0    pypi
[conda] transformers              4.57.3                   pypi_0    pypi
[conda] triton                    3.2.0                    pypi_0    pypi
FastVideo Version:
FastVideo Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
        GPU0    GPU1    GPU2    GPU3    GPU4    GPU5    GPU6    GPU7    NIC0    NIC1    NIC2    NIC3    NIC4    NIC5    CPU Affinity    NUMA Affinity   GPU NUMA ID
GPU0     X      NV8     NV8     NV8     NV8     NV8     NV8     NV8     PXB     NODE    SYS     SYS     NODE    PXB     0-31,64-95      0               N/A
GPU1    NV8      X      NV8     NV8     NV8     NV8     NV8     NV8     PXB     NODE    SYS     SYS     NODE    PXB     0-31,64-95      0               N/A
GPU2    NV8     NV8      X      NV8     NV8     NV8     NV8     NV8     NODE    PXB     SYS     SYS     NODE    NODE    0-31,64-95      0               N/A
GPU3    NV8     NV8     NV8      X      NV8     NV8     NV8     NV8     NODE    PXB     SYS     SYS     NODE    NODE    0-31,64-95      0               N/A
GPU4    NV8     NV8     NV8     NV8      X      NV8     NV8     NV8     SYS     SYS     PXB     NODE    SYS     SYS     32-63,96-127    1               N/A
GPU5    NV8     NV8     NV8     NV8     NV8      X      NV8     NV8     SYS     SYS     PXB     NODE    SYS     SYS     32-63,96-127    1               N/A
GPU6    NV8     NV8     NV8     NV8     NV8     NV8      X      NV8     SYS     SYS     NODE    PXB     SYS     SYS     32-63,96-127    1               N/A
GPU7    NV8     NV8     NV8     NV8     NV8     NV8     NV8      X      SYS     SYS     NODE    PXB     SYS     SYS     32-63,96-127    1               N/A
NIC0    PXB     PXB     NODE    NODE    SYS     SYS     SYS     SYS      X      NODE    SYS     SYS     NODE    PIX
NIC1    NODE    NODE    PXB     PXB     SYS     SYS     SYS     SYS     NODE     X      SYS     SYS     NODE    NODE
NIC2    SYS     SYS     SYS     SYS     PXB     PXB     NODE    NODE    SYS     SYS      X      NODE    SYS     SYS
NIC3    SYS     SYS     SYS     SYS     NODE    NODE    PXB     PXB     SYS     SYS     NODE     X      SYS     SYS
NIC4    NODE    NODE    NODE    NODE    SYS     SYS     SYS     SYS     NODE    NODE    SYS     SYS      X      NODE
NIC5    PXB     PXB     NODE    NODE    SYS     SYS     SYS     SYS     PIX     NODE    SYS     SYS     NODE     X

Legend:

  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinks

NIC Legend:

  NIC0: mlx5_2
  NIC1: mlx5_3
  NIC2: mlx5_4
  NIC3: mlx5_5
  NIC4: mlx5_bond_0
  NIC5: mlx5_bond_1


CUDA_HOME=/mnt/petrelfs/share/cuda-12.4
LD_LIBRARY_PATH=/mnt/petrelfs/share/ffmpeg-6.1.2-shared/lib:/mnt/petrelfs/wangyating/anaconda3/lib:/mnt/petrelfs/share/gcc/gmp-6.2.0/lib:/mnt/petrelfs/share/gcc/mpfr-4.1.0/lib:/mnt/petrelfs/share/gcc/mpc-0.8.1/lib:/mnt/petrelfs/share/cuda-12.4/lib64:/mnt/petrelfs/wangyating/anaconda3/lib:/mnt/petrelfs/share/gcc/gmp-6.2.0/lib:/mnt/petrelfs/share/gcc/mpfr-4.1.0/lib:/mnt/petrelfs/share/gcc/mpc-0.8.1/lib:/mnt/petrelfs/share/cuda-12.4/lib64:/mnt/petrelfs/share/ffmpeg-6.1.2-shared/lib:/mnt/petrelfs/wangyating/anaconda3/lib:/mnt/petrelfs/share/gcc/gmp-6.2.0/lib:/mnt/petrelfs/share/gcc/mpfr-4.1.0/lib:/mnt/petrelfs/share/gcc/mpc-0.8.1/lib:/mnt/petrelfs/share/cuda-12.4/lib64::/mnt/petrelfs/wangyating/.mujoco/mujoco210/bin:/usr/lib/nvidia:/mnt/petrelfs/wangyating/.mujoco/mujoco210/bin:/usr/lib/nvidia:/mnt/petrelfs/wangyating/.mujoco/mujoco210/bin:/usr/lib/nvidia

[SolitaryThinker]

Could you try aligning your parameters and settings with this example script and try on main: https://github.com/hao-ai-lab/FastVideo/blob/main/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm

I noticed that your LR is 1e-5 which is used in much older version of the code/script. In addition, you can also lower your sparsity to 0.8

[xiaoxiao0406]

Thank you for your response.

I initially set the learning rate following the Sparse Adaptation Setup section in the paper. Following your suggestion, I have aligned my parameters and settings with the example script you provided and tested it on the main branch.

However, the loss remains similar to my previous attempts. I also tried lowering the VSA_sparsity to 0.8 and even to 0.0, but the loss actually increased.