[Feature[[Readme] Add VSA/DMD doc (#673)

2025-08-02 02:35:46 -04:00
parent 6c81befc87
commit c9005045dc
14 changed files with 238 additions and 43 deletions
@@ -42,6 +42,7 @@ docs/_build/
 docs/source/getting_started/examples/
 docs/source/inference/examples/
 docs/source/training/examples/
+docs/source/distillation/examples/

 # VSCode
 .vscode/
@@ -95,7 +95,7 @@ For a more detailed guide, please see our [inference quick start](https://hao-ai
 - [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)

 ## Distillation and Finetuning
- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/training/distillation.html)
+- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html)
 - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html)

 ## 📑 Development Plan
@@ -0,0 +1,43 @@
+(v0-data-preprocess)=
+
+# 🧱 Data Preprocess for Distillation
+
+For distillation, we use the same data preprocessing pipeline as training. Please refer to the [Training Data Preprocess](../training/data_preprocess.md) for general preprocessing steps.
+
+## Distillation-Specific Datasets
+
+### FastVideo 480P Synthetic Wan Dataset
+
+For Wan2.1 T2V distillation, we use the **FastVideo 480P Synthetic Wan dataset** ([FastVideo/Wan-Syn_77x448x832_600k](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)) which contains 600k synthetic latents.
+
+```bash
+# Download the preprocessed dataset
+python scripts/huggingface/download_hf.py \
+    --repo_id "FastVideo/Wan-Syn_77x448x832_600k" \
+    --local_dir "FastVideo/Wan-Syn_77x448x832_600k" \
+    --repo_type "dataset"
+```
+
+### Crush Smol Dataset
+
+For Wan2.2 TI2V distillation, we use the crush_smol dataset which includes both raw videos and preprocessed latents.
+
+```bash
+# Download dataset
+python scripts/huggingface/download_hf.py \
+    --repo_id=FastVideo/mini_i2v_dataset \
+    --local_dir=data/mini_i2v_dataset \
+    --repo_type=dataset
+```
+
+## Preprocessing for Distillation
+
+The preprocessing steps are identical to training. Run the appropriate preprocessing script based on your model:
+
+```bash
+# For Wan2.1 T2V
+bash scripts/preprocess/v1_preprocess_wan_data_t2v
+
+# For Wan2.2 TI2V  
+bash examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/preprocess_wan_data_ti2v_5b.sh
+```
@@ -0,0 +1,69 @@
+# 🎯 Distillation
+
+We introduce a new finetuning strategy - **Sparse-distill**, which jointly integrates **[DMD](https://arxiv.org/abs/2405.14867)** and **[VSA](https://arxiv.org/abs/2505.13389)** in a single training process. This approach combines the benefits of both distillation to shorten diffusion steps and sparse attention to reduce attention computations, enabling much faster video generation.
+
+## 📊 Model Overview
+
+We provide two distilled models:
+
+- **[FastWan2.1-T2V-1.3B-Diffusers](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers)**: 3-step inference, up to **20 FPS** on H100 GPU
+- **[FastWan2.1-T2V-14B-480P-Diffusers](https://huggingface.co/FastVideo/FastWan2.1-T2V-14B-480P-Diffusers)**: 3-step inference, up to **50x speed up** at 480P, **70x speed up** at 720P for denoising loop
+
+Both models are trained on **61×448×832** resolution but support generating videos with **any resolution** (1.3B  model mainly support 480P, 14B model support 480P and 720P, quality may degrade for different resolutions).
+
+## 🗂️ Dataset
+
+We use the **FastVideo 480P Synthetic Wan dataset** ([FastVideo/Wan-Syn_77x448x832_600k](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)) for distillation, which contains 600k synthetic latents.
+
+### Download Dataset
+
+```bash
+# Download the preprocessed dataset
+python scripts/huggingface/download_hf.py \
+    --repo_id "FastVideo/Wan-Syn_77x448x832_600k" \
+    --local_dir "FastVideo/Wan-Syn_77x448x832_600k" \
+    --repo_type "dataset"
+```
+
+## 🚀 Training Scripts
+
+### 1.3B Model Sparse-Distill
+
+For the 1.3B model, we use **4 nodes with 32 H200 GPUs** (8 GPUs per node):
+
+```bash
+# Multi-node training (8 nodes, 64 GPUs total)
+sbatch examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm
+```
+
+**Key Configuration:**
+- Global batch size: 64
+- Gradient accumulation steps: 2
+- Learning rate: 1e-5
+- VSA attention sparsity: 0.8
+- Training steps: 4000 (~12 hours)
+
+### 14B Model Sparse-Distill
+
+For the 14B model, we use **8 nodes with 64 H200 GPUs** (8 GPUs per node):
+
+```bash
+# Multi-node training (8 nodes, 64 GPUs total)
+sbatch examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm
+```
+
+**Key Configuration:**
+- Global batch size: 64
+- Sequence parallel size: 4
+- Gradient accumulation steps: 4
+- Learning rate: 1e-5
+- VSA attention sparsity: 0.9
+- Training steps: 3000 (~52 hours)
+- HSDP shard dim: 8
+
+## ⚙️ Inference
+Set `MODEL_BASE` to your own model path and run:
+
+```bash
+bash scripts/inference/v1_inference_wan_dmd.sh
+```
@@ -287,11 +287,35 @@ def create_nested_structures(
        relative_path = example.path.relative_to(category_dir)
        path_parts = relative_path.parts

-        # For nested examples like finetune/wan_i2v_14b_480p/crush_smol
-        if len(path_parts) >= 3:
-            method = path_parts[0]  # e.g., "finetune"
-            model = path_parts[1]  # e.g., "wan_i2v_14b_480p"
-            dataset = path_parts[2]  # e.g., "crush_smol"
+        if example.category == "training":
+            # For training examples like finetune/wan_i2v_14b_480p/crush_smol
+            if len(path_parts) >= 3:
+                method = path_parts[0]  # e.g., "finetune"
+                model = path_parts[1]  # e.g., "wan_i2v_14b_480p"
+                dataset = path_parts[2]  # e.g., "crush_smol"
+
+                # Initialize nested structure
+                if example.category not in nested_structures:
+                    nested_structures[example.category] = {}
+                if method not in nested_structures[example.category]:
+                    nested_structures[example.category][method] = {}
+                if model not in nested_structures[example.category][method]:
+                    nested_structures[example.category][method][model] = {}
+
+                # Store the nested structure
+                nested_structures[
+                    example.category][method][model][dataset] = NestedStructure(
+                        category=example.category,
+                        method=method,
+                        model=model,
+                        dataset=dataset,
+                        example=example)
+
+        elif example.category == "distillation" and len(path_parts) >= 2:
+            # For distillation examples like Wan2.1-T2V/Wan-Syn-Data-480P
+            model = path_parts[0]  # e.g., "Wan2.1-T2V"
+            dataset = path_parts[1]  # e.g., "Wan-Syn-Data-480P"
+            method = "DMD"  # Default method for distillation

            # Initialize nested structure
            if example.category not in nested_structures:
@@ -33,6 +33,7 @@ With FastVideo's optimizations, you can achieve more than 3x inference improveme
 FastVideo has the following features:
 - State-of-the-art performance optimizations for inference
  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [Video Sparse Attention](https://arxiv.org/pdf/2505.13389)
  - [TeaCache](https://arxiv.org/pdf/2411.19108)
  - [Sage Attention](https://arxiv.org/abs/2410.02367)
 - Cutting edge models
@@ -78,18 +79,16 @@ inference/add_pipeline

 training/examples/examples_training_index
 training/data_preprocess
-training/distillation
 <!-- training/finetune -->
 :::

-<!-- :::{toctree}
+:::{toctree}
 :caption: Distillation
 :maxdepth: 1

 distillation/examples/examples_distillation_index
 distillation/data_preprocess
-distillation/dmd -->
-<!-- training/finetune -->
+distillation/dmd
 :::

 % What is STA Kernel?
@@ -102,6 +101,16 @@ sliding_tile_attention/installation
 sliding_tile_attention/demo
 :::

+% What is VSA Kernel?
+
+:::{toctree}
+:caption: Video Sparse Attention
+:maxdepth: 1
+
+video_sparse_attention/installation
+video_sparse_attention/demo
+:::
+
 :::{toctree}
 :caption: Design
 :maxdepth: 1
@@ -19,6 +19,7 @@ This page describes the various options for speeding up generation times in Fast
 - Torch SDPA: `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`
 - Flash Attention 2 and 3: `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN`
 - Sliding Tile Attention: `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
+- Video Sparse Attention: `FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN`
 - Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`

 ### Configuring Backends
@@ -74,6 +75,17 @@ pip install st_attn==0.0.4

 Please see [this page](#sta-installation) for more installation instructions.

+(optimizations-vsa)=
+### Video Sparse Attention
+**`VIDEO_SPARSE_ATTN`**
+
+```bash
+git submodule update --init --recursive
+python setup_vsa.py install
+```
+
+Please see [this page](#vsa-installation) for more installation instructions.
+
 (optimizations-sage)=
 ### Sage Attention
 **`SAGE_ATTN`**
@@ -29,13 +29,13 @@ export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 git submodule update --init --recursive
-python setup.py install
+python setup_sta.py install
 ```

 # 🧪 Test

 ```bash
-python test/test_sta.py
+python csrc/attn/tests/test_sta.py
 ```

 # 📋 Usage
@@ -1,25 +0,0 @@
-(v0-distill)=
-# 🎯 Distill
-Our distillation recipe is based on [Phased Consistency Model](https://github.com/G-U-N/Phased-Consistency-Model). We did not find significant improvement using multi-phase distillation, so we keep the one phase setup similar to the original latent consistency model's recipe.
-We use the [MixKit](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/all_mixkit) dataset for distillation. To avoid running the text encoder and VAE during training, we prprocess all data to generate text embeddings and VAE latents.
-Preprocessing instructions can be found [data_preprocess.md](#v0-data-preprocess). For convenience, we also provide preprocessed data that can be downloaded directly using the following command:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/HD-Mixkit-Finetune-Hunyuan --local_dir=data/HD-Mixkit-Finetune-Hunyuan --repo_type=dataset
-```
-
-Next, download the original model weights with:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model # original hunyuan
-python scripts/huggingface/download_hf.py --repo_id=genmo/mochi-1-preview --local_dir=data/mochi --repo_type=model # original mochi
-```
-
-To launch the distillation process, use the following commands:
-
-```
-bash scripts/distill/distill_hunyuan.sh # for hunyuan
-bash scripts/distill/distill_mochi.sh # for mochi
-```
-
-We also provide an optional script for distillation with adversarial loss, located at `fastvideo/distill_adv.py`. Although we tried adversarial loss, we did not observe significant improvements.
@@ -7,7 +7,7 @@ Ensure your data is prepared and preprocessed in the format specified in [data_p
 python scripts/huggingface/download_hf.py --repo_id=FastVideo/Mochi-Black-Myth --local_dir=data/Mochi-Black-Myth --repo_type=dataset
 ```

-Download the original model weights as specified in [Distill Section](#v0-distill):
+Download the original model weights as specified in the [Distillation Section](../distillation/dmd.md):

 Then you can run the finetune with:

@@ -0,0 +1,3 @@
+(vsa-demo)=
+
+# 🎬 Demo
@@ -0,0 +1,61 @@
+(vsa-installation)=
+
+# 🔧 Installation
+You can install the Video Sparse Attention package using
+
+```bash
+git submodule update --init --recursive
+python setup_vsa.py install
+```
+
+# Building from Source
+We support H100 (via ThunderKittens) and any other GPU (via Triton) for VSA.
+
+First, install C++20 for ThunderKittens (if using H100):
+
+```bash
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+
+Set up CUDA environment (if using CUDA 12.4):
+
+```bash
+export CUDA_HOME=/usr/local/cuda-12.4
+export PATH=${CUDA_HOME}/bin:${PATH} 
+export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+Install VSA:
+
+```bash
+cd csrc/attn/
+git submodule update --init --recursive
+python setup_vsa.py install
+```
+
+# 🧪 Test
+
+```bash
+python csrc/attn/tests/test_vsa.py
+```
+
+# 📋 Usage
+
+```python
+from vsa import video_sparse_attn
+
+# q, k, v: [batch_size, num_heads, seq_len, head_dim]
+# variable_block_sizes: [num_blocks] - number of valid tokens in each block
+# topk: int - number of top-k blocks to attend to
+# block_size: int or tuple of 3 ints - size of each block (default: 64 tokens)
+# compress_attn_weight: optional weight for compressed attention branch
+
+output = video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight)
+
+``` 
@@ -33,7 +33,7 @@ class WanDMDPipeline(LoRAPipeline, ComposedPipelineBase):
    ]

    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        # We use UniPCMScheduler from Wan2.1 official repo, not the one in diffusers.
+
        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
            shift=fastvideo_args.pipeline_config.flow_shift)

@@ -715,10 +715,8 @@ class DmdDenoisingStage(DenoisingStage):
                with torch.autocast(device_type="cuda",
                                    dtype=target_dtype,
                                    enabled=autocast_enabled):
-                    if (st_attn_available
-                            and self.attn_backend == SlidingTileAttentionBackend
-                        ) or (vsa_available and self.attn_backend
-                              == VideoSparseAttentionBackend):
+                    if (vsa_available and self.attn_backend
+                            == VideoSparseAttentionBackend):
                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
                        )