From c9005045dc9bcd37f7e145c6cf67e6e2eed7aecc Mon Sep 17 00:00:00 2001 From: Yongqi Chen <144848849+BrianChen1129@users.noreply.github.com> Date: Sat, 2 Aug 2025 02:35:46 -0400 Subject: [PATCH] [Feature[[Readme] Add VSA/DMD doc (#673) --- .gitignore | 1 + README.md | 2 +- docs/source/distillation/data_preprocess.md | 43 ++++++++++++ docs/source/distillation/dmd.md | 69 +++++++++++++++++++ docs/source/generate_examples.py | 34 +++++++-- docs/source/index.md | 17 +++-- docs/source/inference/optimizations.md | 12 ++++ .../sliding_tile_attention/installation.md | 4 +- docs/source/training/distillation.md | 25 ------- docs/source/training/finetune.md | 2 +- docs/source/video_sparse_attention/demo.md | 3 + .../video_sparse_attention/installation.md | 61 ++++++++++++++++ .../pipelines/basic/wan/wan_dmd_pipeline.py | 2 +- fastvideo/pipelines/stages/denoising.py | 6 +- 14 files changed, 238 insertions(+), 43 deletions(-) create mode 100644 docs/source/distillation/data_preprocess.md create mode 100644 docs/source/distillation/dmd.md delete mode 100644 docs/source/training/distillation.md create mode 100644 docs/source/video_sparse_attention/demo.md create mode 100644 docs/source/video_sparse_attention/installation.md diff --git a/.gitignore b/.gitignore index 57711d685..0196c653b 100644 --- a/.gitignore +++ b/.gitignore @@ -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/ diff --git a/README.md b/README.md index f6b9cc0cc..7e7bb2d05 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/docs/source/distillation/data_preprocess.md b/docs/source/distillation/data_preprocess.md new file mode 100644 index 000000000..9388d804d --- /dev/null +++ b/docs/source/distillation/data_preprocess.md @@ -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 +``` diff --git a/docs/source/distillation/dmd.md b/docs/source/distillation/dmd.md new file mode 100644 index 000000000..30a86e5d1 --- /dev/null +++ b/docs/source/distillation/dmd.md @@ -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 +``` diff --git a/docs/source/generate_examples.py b/docs/source/generate_examples.py index ba781c282..a2e4f6696 100644 --- a/docs/source/generate_examples.py +++ b/docs/source/generate_examples.py @@ -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: diff --git a/docs/source/index.md b/docs/source/index.md index 0fdfac364..c8c0a3a6a 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -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 ::: - - +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 diff --git a/docs/source/inference/optimizations.md b/docs/source/inference/optimizations.md index 7cff763d1..609f3f46f 100644 --- a/docs/source/inference/optimizations.md +++ b/docs/source/inference/optimizations.md @@ -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`** diff --git a/docs/source/sliding_tile_attention/installation.md b/docs/source/sliding_tile_attention/installation.md index 8562a4b55..9069bcdd0 100644 --- a/docs/source/sliding_tile_attention/installation.md +++ b/docs/source/sliding_tile_attention/installation.md @@ -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 diff --git a/docs/source/training/distillation.md b/docs/source/training/distillation.md deleted file mode 100644 index 44d331512..000000000 --- a/docs/source/training/distillation.md +++ /dev/null @@ -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. diff --git a/docs/source/training/finetune.md b/docs/source/training/finetune.md index 391d6ce2c..92636779c 100644 --- a/docs/source/training/finetune.md +++ b/docs/source/training/finetune.md @@ -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: diff --git a/docs/source/video_sparse_attention/demo.md b/docs/source/video_sparse_attention/demo.md new file mode 100644 index 000000000..af048428b --- /dev/null +++ b/docs/source/video_sparse_attention/demo.md @@ -0,0 +1,3 @@ +(vsa-demo)= + +# ๐ŸŽฌ Demo diff --git a/docs/source/video_sparse_attention/installation.md b/docs/source/video_sparse_attention/installation.md new file mode 100644 index 000000000..78a9a907c --- /dev/null +++ b/docs/source/video_sparse_attention/installation.md @@ -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) + +``` diff --git a/fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py index 81ece4609..79d928f0d 100644 --- a/fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py @@ -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) diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 07d627c4a..176fbdb9e 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -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( )