nodes/aigc-apps-VideoX-Fun
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update Sparkwan, Reward Model Training, FastVideo (#525)

Browse Source
This commit is contained in:
Bubbliiiing
2026-10-28 11:34:50 +08:00
committed by GitHub
parent 9dc6f1d49a
commit e8c314fbfa
81 changed files with 29460 additions and 1105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/fastvideo/README_TRAIN_DISTILL.md
+648
View File
@@ -0,0 +1,648 @@
# FastVideo VSA Distill Training Guide
This document provides a complete workflow for distilling Wan2.1 into FastWan2.1 with VSA (Video Sparse Attention), including environment configuration, data preparation, distributed training, and inference testing.
> **Note**: FastVideo VSA combines a block-sparse video attention (a learned gate keeps ~20% of key blocks per query) with DMD few-step distillation, so the student runs CFG-free in ~3 steps while attending sparsely. The training recipe, data pipeline, and distributed backends are identical to the TurboDiffusion distill guide; only the sparse student model class differs.
---
## Table of Contents
- [1. Environment Configuration](#1-environment-configuration)
- [2. Data Preparation](#2-data-preparation)
- [2.1 Quick Test Dataset](#21-quick-test-dataset)
- [2.2 Dataset Structure](#22-dataset-structure)
- [2.3 metadata.json Format](#23-metadatajson-format)
- [2.4 Relative vs Absolute Path Usage](#24-relative-vs-absolute-path-usage)
- [3. Distill Training](#3-distill-training)
- [3.1 Download Pretrained Model](#31-download-pretrained-model)
- [3.2 Quick Start (DeepSpeed-Zero-2)](#32-quick-start-deepspeed-zero-2)
- [3.3 Common Training Parameters](#33-common-training-parameters)
- [3.4 Training Validation](#34-training-validation)
- [3.5 Training with FSDP](#35-training-with-fsdp)
- [3.6 Other Backends](#36-other-backends)
- [3.7 Multi-Machine Distributed Training](#37-multi-machine-distributed-training)
- [4. Inference Testing](#4-inference-testing)
- [4.1 Inference Parameters](#41-inference-parameters)
- [4.2 Text-to-Video (T2V) Inference](#42-text-to-video-t2v-inference)
- [4.3 Single-GPU Only](#43-single-gpu-only)
- [5. Additional Resources](#5-additional-resources)
---
## 1. Environment Configuration
**Method 1: Using requirements.txt**
```bash
pip install -r requirements.txt
```
**Method 2: Manual Dependency Installation**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 💡 The VSA block-sparse attention runs on a Triton kernel; Triton ships with PyTorch (>=2.1) on Linux, so no extra install is needed.
**Method 3: Using Docker**
When using Docker, please ensure that the GPU driver and CUDA environment are correctly installed on your machine, then execute the following commands:
```
# pull image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# enter image
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 2. Data Preparation
### 2.1 Quick Test Dataset
We provide a test dataset containing several training samples.
```bash
# Download official demo dataset
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 Dataset Structure
```
📦 datasets/
├── 📂 my_dataset/
│ ├── 📂 train/
│ │ ├── 📄 video001.mp4
│ │ ├── 📄 video002.mp4
│ │ └── 📄 ...
│ └── 📄 metadata.json
```
### 2.3 metadata.json Format
**Relative Path Format** (example):
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
},
{
"file_path": "train/video002.mp4",
"text": "A person walking through a forest, cinematic view",
"type": "video",
"width": 1328,
"height": 1328
}
]
```
**Absolute Path Format**:
```json
[
{
"file_path": "/mnt/data/videos/sunset.mp4",
"text": "A beautiful sunset over the ocean",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
**Key Field Descriptions**:
- `file_path`: Video path (relative or absolute path)
- `text`: Video description (English prompt)
- `type`: Data type, fixed as `"video"`
- `width` / `height`: Video dimensions (**recommended** to provide for bucket training. If not provided, they will be automatically read during training, which may slow down training when data is stored on slow systems like OSS).
- You can use `scripts/process_json_add_width_and_height.py` to extract width and height fields for JSON files without them, supporting both images and videos.
- Usage: `python scripts/process_json_add_width_and_height.py --input_file datasets/X-Fun-Videos-Demo/metadata.json --output_file datasets/X-Fun-Videos-Demo/metadata_add_width_height.json`.
### 2.4 Relative vs Absolute Path Usage
**Relative Path**:
If your data uses relative paths, configure in the training script:
```bash
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
```
**Absolute Path**:
If your data uses absolute paths, configure in the training script:
```bash
export DATASET_NAME=""
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
```
> 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths.
---
## 3. Distill Training
The distillation is a DMD (distribution matching distillation) run with three networks: a sparse **generator** (the FastWan student, `FastVideoWanTransformer3DModel`), a sparse **fake_score** critic (same class), and a frozen **dense** `real_score` teacher (`WanTransformer3DModel`). The student and critic carry the VSA block-sparse self-attention; the teacher keeps full attention.
### 3.1 Download Pretrained Model
```bash
# Create model directory
mkdir -p models/Diffusion_Transformer
# Download Wan2.1 official weights
modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
```
### 3.2 Quick Start (DeepSpeed-Zero-2)
After downloading the data as per **2.1 Quick Test Dataset** and the weights as per **3.1 Download Pretrained Model**, you can directly copy and run the quick start command.
DeepSpeed-Zero-2 and FSDP are recommended for training. Here we use DeepSpeed-Zero-2 as an example.
The difference between DeepSpeed-Zero-2 and FSDP is whether the model weights are sharded. **If you experience insufficient GPU memory when using multiple GPUs with DeepSpeed-Zero-2**, you can switch to FSDP for training.
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.3 Common Training Parameters
**Key Parameter Descriptions**:
| Parameter | Description | Example Value |
|-----|------|-------|
| `--config_path` | Config file path | `config/wan2.1/wan_civitai.yaml` |
| `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` |
| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` |
| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` |
| `--train_batch_size` | Batch size per GPU | 1 |
| `--image_sample_size` | Maximum image training resolution | 640 |
| `--video_sample_size` | Maximum video training resolution | 640 |
| `--token_sample_size` | Token sample size | 640 |
| `--video_sample_stride` | Video sampling stride | 2 |
| `--video_sample_n_frames` | Video sampling frames | 81 |
| `--gradient_accumulation_steps` | Gradient accumulation steps (effectively increases batch) | 1 |
| `--dataloader_num_workers` | DataLoader worker processes | 8 |
| `--num_train_epochs` | Number of training epochs | 100 |
| `--checkpointing_steps` | Save checkpoint every N steps | 50 |
| `--learning_rate` | Initial learning rate (generator) | 2e-06 |
| `--learning_rate_critic` | Initial learning rate (critic) | 2e-06 |
| `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` |
| `--lr_warmup_steps` | Learning rate warmup steps | 100 |
| `--seed` | Random seed | 42 |
| `--output_dir` | Output directory | `output_dir_fastvideo_wan2.1_1.3b` |
| `--gradient_checkpointing` | Enable gradient checkpointing | - |
| `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
| `--adam_weight_decay` | AdamW weight decay | 3e-2 |
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--vae_mini_batch` | VAE encoding mini-batch size | 1 |
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
| `--enable_bucket` | Enable bucket training, no cropping, group by resolution | - |
| `--random_hw_adapt` | Auto-scale images/videos to random size in `[min_size, max_size]` range | - |
| `--training_with_video_token_length` | Train based on token length, supports arbitrary resolutions | - |
| `--uniform_sampling` | Uniform timestep sampling | - |
| `--low_vram` | Low VRAM mode | - |
| `--train_mode` | Training mode: `normal` (standard) or `i2v` (image-to-video) | `normal` |
| `--resume_from_checkpoint` | Resume training path, use `"latest"` to auto-select latest checkpoint | None |
| `--validation_steps` | Run validation every N steps | 2000 |
| `--validation_epochs` | Run validation every N epochs | 5 |
| `--validation_prompts` | Prompts for video generation validation | `"A dog shaking head..."` |
| `--trainable_modules` | Trainable modules (`"."` means all modules) | `"."` |
**Distill-Specific Parameters**:
| Parameter | Description | Example Value |
|-----|------|-------|
| `--denoising_step_indices_list` | Denoising step list (core distill parameter) | `1000 750 500 250` |
| `--real_guidance_scale` | Real guidance scale for scoring | 6.0 |
| `--fake_guidance_scale` | Fake guidance scale for scoring | 0.0 |
| `--gen_update_interval` | Generator update interval | 5 |
| `--negative_prompt` | Negative prompt for distillation | Chinese negative prompt |
| `--validation_paths` | Validation image paths for I2V mode | Image path list |
| `--train_sampling_steps` | Training sampling steps | 1000 |
**VSA Sparsity**:
- The student/critic self-attention is VSA block-sparse. The keep-ratio `student_topk` defaults to `0.20` (i.e. 80% sparsity) inside `FastVideoWanTransformer3DModel`; it is a model-class default, not a CLI flag.
- The learned compression gate `to_gate_compress` is created zero-initialized and becomes trainable automatically under `--trainable_modules "."` (its parameter names contain a dot, like every other module).
- The `real_score` teacher stays dense (`WanTransformer3DModel`) and frozen, so the sparsity is distilled into the student rather than copied from the teacher.
- Keep the inference-time `student_topk` equal to the value the checkpoint was trained at; changing it rebuilds the attention op and resets the learned gate.
**Sample Size Configuration Guide**:
- `video_sample_size` represents the resolution size of videos; when `random_hw_adapt` is True, it represents the minimum value between video and image resolutions.
- `image_sample_size` represents the resolution size of images; when `random_hw_adapt` is True, it represents the maximum value between video and image resolutions.
- `token_sample_size` represents the resolution corresponding to the maximum token length when `training_with_video_token_length` is True.
- Due to potential confusion in configuration, **if you don't require arbitrary resolution for finetuning**, it is recommended to set `video_sample_size`, `image_sample_size`, and `token_sample_size` to the same fixed value, such as **(320, 480, 512, 640, 960)**.
- **All set to 320** represents **240P**.
- **All set to 480** represents **320P**.
- **All set to 640** represents **480P**.
- **All set to 960** represents **720P**.
**Other Parameter Explanations**:
- `enable_bucket` is used to enable bucket training. When enabled, the model does not crop the images and videos at the center, but instead, it trains the entire images and videos after grouping them into buckets based on resolution.
- `random_hw_adapt` is used to enable automatic height and width scaling for images and videos. When `random_hw_adapt` is enabled, the training images will have their height and width set to `image_sample_size` as the maximum and `min(video_sample_size, 512)` as the minimum.
- `training_with_video_token_length` specifies training the model according to token length. For training images and videos, the height and width will be set to `image_sample_size` as the maximum and `video_sample_size` as the minimum.
- `train_mode` is used to specify the training mode, which can be either normal or i2v. If you only wish to achieve text-to-video generation, keep `normal`.
- `resume_from_checkpoint` is used to set the training should be resumed from a previous checkpoint. Use a path or `"latest"` to automatically select the last available checkpoint.
### 3.4 Training Validation
You can configure validation parameters to periodically generate test videos during training, allowing you to monitor training progress and model quality.
**Validation Parameter Descriptions**:
| Parameter | Description | Recommended Value |
|------|------|--------|
| `--validation_steps` | Run validation every N steps | 2000 |
| `--validation_epochs` | Run validation every N epochs | 5 |
| `--validation_prompts` | Prompts for video generation validation | English prompts |
| `--validation_paths` | Validation image paths for I2V mode (i2v mode only) | `"asset/1.png"` |
**Normal Mode Example** (T2V Validation):
```bash
--validation_steps=2000 \
--validation_epochs=5 \
--validation_prompts="A dog shaking head. The video is of high quality, and the view is very clear."
```
**Notes**:
- Validation videos will be saved to the `output_dir` directory.
- Multiple prompts validation format: `--validation_prompts "prompt1" "prompt2" "prompt3"`.
- Distill model validation uses the few steps defined in `denoising_step_indices_list` and runs CFG-free (`guidance_scale=1.0`), matching inference.
### 3.5 Training with FSDP
**If you experience insufficient GPU memory when using multiple GPUs with DeepSpeed-Zero-2**, you can switch to FSDP for training. `scripts/fastvideo/train_distill_wan2.1.sh` is a ready-to-run FSDP launcher (it wraps `WanAttentionBlock` with `FULL_SHARD`).
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.6 Other Backends
#### 3.6.1 Training with DeepSpeed-Zero-3
DeepSpeed Zero-3 is not highly recommended at the moment. In this repository, using FSDP has fewer errors and is more stable.
DeepSpeed Zero-3 is suitable for 14B Wan at high resolutions. After training, you can use the following command to get the final model:
```bash
python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/checkpoint-{your-num-steps}-outputs --max_shard_size 80GB --safe_serialization
```
Training shell command is as follows:
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
#### 3.6.2 Training Without DeepSpeed and FSDP
**This approach is not recommended because without memory-saving backends, it easily causes insufficient GPU memory**. This is only provided as a reference shell for training.
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.7 Multi-Machine Distributed Training
**Suitable for**: Ultra-large datasets, faster training speed
#### 3.7.1 Environment Configuration
Assuming 2 machines, each with 8 GPUs:
**Machine 0 (Master)**:
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export MASTER_ADDR="192.168.1.100" # Master machine IP
export MASTER_PORT=10086
export WORLD_SIZE=2 # Total number of machines
export NUM_PROCESS=16 # Total processes = machines × 8
export RANK=0 # Current machine rank (0 or 1)
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
**Machine 1 (Worker)**:
```bash
export MASTER_ADDR="192.168.1.100" # Same as Master
export MASTER_PORT=10086
export WORLD_SIZE=2
export NUM_PROCESS=16
export RANK=1 # Note this is 1
NCCL_DEBUG=INFO
# Use the same accelerate launch command as Machine 0
```
#### 3.7.2 Multi-Machine Training Notes
- **Network Requirements**:
- RDMA/InfiniBand recommended (high performance)
- Without RDMA, add environment variables:
```bash
export NCCL_IB_DISABLE=1
export NCCL_P2P_DISABLE=1
```
- **Data Synchronization**: All machines must be able to access the same data paths (NFS/shared storage)
---
## 4. Inference Testing
### 4.1 Inference Parameters
**Key Parameter Descriptions** (from `examples/fastvideo/predict_t2v_wan2.1.py`):
| Parameter | Description | Example Value |
|------|------|-------|
| `GPU_memory_mode` | GPU memory mode, see table below for options | `model_cpu_offload` |
| `ulysses_degree` | Ulysses parallelism degree (VSA is single-GPU only, keep 1) | 1 |
| `ring_degree` | Ring parallelism degree (VSA is single-GPU only, keep 1) | 1 |
| `config_path` | Config file path | `config/wan2.1/wan_civitai.yaml` |
| `model_name` | Base model path (config + dense backbone + VAE + T5) | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
| `student_topk` | VSA keep-ratio; must equal the trained value | `0.20` |
| `sampler_name` | Sampler type: `Flow`, `Flow_Unipc`, `Flow_DPM++` | `Flow_Unipc` |
| `shift` | Noise schedule shift (official FastWan flow_shift) | `3.0` |
| `transformer_path` | Path to the FastWan student Transformer weights | trained checkpoint / official FastWan |
| `vae_path` | Path to trained VAE weights | `None` |
| `lora_path` | LoRA weights path | `None` |
| `sample_size` | Generated video resolution `[height, width]` | `[480, 848]` |
| `video_length` | Number of generated frames | `81` |
| `fps` | Frames per second | `16` |
| `weight_dtype` | Model weight precision, use `torch.float16` for GPUs without bf16 support | `torch.bfloat16` |
| `prompt` | Positive prompt describing generated content | `"A cinematic aerial shot..."` |
| `negative_prompt` | Negative prompt to avoid certain content | Chinese negative prompt |
| `guidance_scale` | Guidance strength (the DMD student is CFG-free) | 1.0 |
| `seed` | Random seed for reproducibility | 1024 |
| `num_inference_steps` | Number of inference steps (FastWan uses 3) | 3 |
| `save_path` | Path to save generated videos | `samples/fastvideo-wan2.1-official` |
**GPU Memory Mode Descriptions**:
| Mode | Description | Memory Usage |
|------|------|---------|
| `model_full_load` | Full model loaded to GPU | Highest |
| `model_full_load_and_qfloat8` | Full load + FP8 quantization | High |
| `model_cpu_offload` | Offload model to CPU after use | Medium |
| `model_cpu_offload_and_qfloat8` | CPU offload + FP8 quantization | Medium-Low |
| `model_group_offload` | Layer groups switch between CPU/CUDA | Low |
| `sequential_cpu_offload` | Layer-by-layer offload (slowest) | Lowest |
### 4.2 Text-to-Video (T2V) Inference
Run single-GPU inference:
```bash
python examples/fastvideo/predict_t2v_wan2.1.py
```
Edit `examples/fastvideo/predict_t2v_wan2.1.py` according to your needs. For first-time inference, focus on the following key parameters. For other parameters, refer to the inference parameter descriptions above.
```python
# Choose based on GPU memory
GPU_memory_mode = "model_cpu_offload"
# Base model path (config + dense backbone + VAE + T5)
model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
# VSA keep-ratio; must equal the value the checkpoint was trained at
student_topk = 0.20
# Path to the FastWan student weights (your trained checkpoint, or the official FastWan release)
transformer_path = "models/Personalized_Model/FastWan2.1-T2V-1.3B.pth"
# FastWan sampler + shift (CFG-free few-step student)
sampler_name = "Flow_Unipc"
shift = 3.0
guidance_scale = 1.0
num_inference_steps = 3
# Write based on your generation content
prompt = "A cinematic aerial shot of a snow-covered mountain village at dawn..."
# ...
```
> 💡 The official `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` checkpoint ships in diffusers key naming; `examples/fastvideo/predict_t2v_wan2.1.py` remaps it to the VideoX-Fun civitai layout (including the VSA `to_gate_compress` gate) before loading. A checkpoint produced by this training guide is already in the native layout.
### 4.3 Single-GPU Only
The VSA block-sparse kernel has no sequence-parallel path, so a FastWan student runs on a **single GPU**. Keep `ulysses_degree = ring_degree = 1`; calling `enable_multi_gpus_inference()` raises otherwise. To fit larger resolutions on one card, lower the memory footprint via `GPU_memory_mode` (e.g. `model_cpu_offload` or `sequential_cpu_offload`) rather than splitting across GPUs.
---
## 5. Additional Resources
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
- **FastVideo GitHub**: https://github.com/hao-ai-lab/FastVideo
- **VSA Paper (Video Sparse Attention)**: https://arxiv.org/abs/2505.13389
scripts/fastvideo/README_TRAIN_DISTILL_zh-CN.md
+648
View File
@@ -0,0 +1,648 @@
# FastVideo VSA 蒸馏训练指南
本文档提供了将 Wan2.1 蒸馏为 FastWan2.1（采用 VSA，Video Sparse Attention）的完整工作流，包括环境配置、数据准备、分布式训练和推理测试。
> **注意**: FastVideo VSA 将块稀疏视频注意力（由可学习门控为每个 query 保留约 20% 的 key 块）与 DMD 少步蒸馏结合，使学生模型在约 3 步、无 CFG 的条件下以稀疏注意力推理。训练配方、数据管线与分布式后端与 TurboDiffusion 蒸馏指南完全一致，仅稀疏学生模型类不同。
---
## 目录
- [一、环境配置](#一环境配置)
- [二、数据准备](#二数据准备)
- [2.1 快速测试数据集](#21-快速测试数据集)
- [2.2 数据集结构](#22-数据集结构)
- [2.3 metadata.json 格式](#23-metadatajson-格式)
- [2.4 相对路径与绝对路径使用方案](#24-相对路径与绝对路径使用方案)
- [三、蒸馏训练](#三蒸馏训练)
- [3.1 下载预训练模型](#31-下载预训练模型)
- [3.2 快速开始（DeepSpeed-Zero-2）](#32-快速开始deepspeed-zero-2)
- [3.3 常见训练参数](#33-常见训练参数)
- [3.4 训练验证](#34-训练验证)
- [3.5 使用 FSDP 训练](#35-使用-fsdp-训练)
- [3.6 其他后端](#36-其他后端)
- [3.7 多机分布式训练](#37-多机分布式训练)
- [四、推理测试](#四推理测试)
- [4.1 推理参数解析](#41-推理参数解析)
- [4.2 文本生成视频（T2V）推理](#42-文本生成视频t2v推理)
- [4.3 仅支持单卡](#43-仅支持单卡)
- [五、更多资源](#五更多资源)
---
## 一、环境配置
**方式一：使用 requirements.txt**
```bash
pip install -r requirements.txt
```
**方式二：手动安装依赖**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 💡 VSA 块稀疏注意力运行在 Triton 内核上；Triton 随 PyTorch（>=2.1）在 Linux 上一并提供，无需额外安装。
**方式三：使用 Docker**
使用 Docker 时，请确保机器上已正确安装 GPU 驱动与 CUDA 环境，然后执行以下命令：
```
# 拉取镜像
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# 进入容器
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 二、数据准备
### 2.1 快速测试数据集
我们提供了一个包含若干训练样本的测试数据集。
```bash
# 下载官方演示数据集
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 数据集结构
```
📦 datasets/
├── 📂 my_dataset/
│ ├── 📂 train/
│ │ ├── 📄 video001.mp4
│ │ ├── 📄 video002.mp4
│ │ └── 📄 ...
│ └── 📄 metadata.json
```
### 2.3 metadata.json 格式
**相对路径格式**（示例）：
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
},
{
"file_path": "train/video002.mp4",
"text": "A person walking through a forest, cinematic view",
"type": "video",
"width": 1328,
"height": 1328
}
]
```
**绝对路径格式**：
```json
[
{
"file_path": "/mnt/data/videos/sunset.mp4",
"text": "A beautiful sunset over the ocean",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
**关键字段说明**：
- `file_path`：视频路径（相对或绝对路径）
- `text`：视频描述（英文 prompt）
- `type`：数据类型，固定为 `"video"`
- `width` / `height`：视频尺寸（**建议**提供以便分桶训练。若不提供，训练时会自动读取，当数据存放在 OSS 等较慢系统上时会拖慢训练）。
- 可使用 `scripts/process_json_add_width_and_height.py` 为缺少宽高字段的 JSON 文件提取 width 和 height，同时支持图片与视频。
- 用法：`python scripts/process_json_add_width_and_height.py --input_file datasets/X-Fun-Videos-Demo/metadata.json --output_file datasets/X-Fun-Videos-Demo/metadata_add_width_height.json`。
### 2.4 相对路径与绝对路径使用方案
**相对路径**：
若数据使用相对路径，在训练脚本中配置：
```bash
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
```
**绝对路径**：
若数据使用绝对路径，在训练脚本中配置：
```bash
export DATASET_NAME=""
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
```
> 💡 **建议**：数据集较小且存放在本地时使用相对路径；数据集存放在外部存储（如 NAS、OSS）或跨多台机器共享时使用绝对路径。
---
## 三、蒸馏训练
蒸馏是一次 DMD（distribution matching distillation）训练，包含三个网络：稀疏的 **generator**（FastWan 学生，`FastVideoWanTransformer3DModel`）、稀疏的 **fake_score** critic（同一个类），以及冻结的 **dense** `real_score` 教师（`WanTransformer3DModel`）。学生与 critic 携带 VSA 块稀疏自注意力，教师保持全注意力。
### 3.1 下载预训练模型
```bash
# 创建模型目录
mkdir -p models/Diffusion_Transformer
# 下载 Wan2.1 官方权重
modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
```
### 3.2 快速开始（DeepSpeed-Zero-2）
按照 **2.1 快速测试数据集** 下载数据、按照 **3.1 下载预训练模型** 下载权重后，即可直接复制运行快速开始命令。
推荐使用 DeepSpeed-Zero-2 或 FSDP 训练，这里以 DeepSpeed-Zero-2 为例。
DeepSpeed-Zero-2 与 FSDP 的区别在于模型权重是否分片。**若使用 DeepSpeed-Zero-2 多卡训练时显存不足**，可切换到 FSDP 训练。
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.3 常见训练参数
**关键参数说明**：
| 参数 | 描述 | 示例值 |
|-----|------|-------|
| `--config_path` | 配置文件路径 | `config/wan2.1/wan_civitai.yaml` |
| `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` |
| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` |
| `--train_data_meta` | 训练数据元信息文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` |
| `--train_batch_size` | 每张卡的 batch size | 1 |
| `--image_sample_size` | 图片训练最大分辨率 | 640 |
| `--video_sample_size` | 视频训练最大分辨率 | 640 |
| `--token_sample_size` | Token 采样尺寸 | 640 |
| `--video_sample_stride` | 视频采样步长 | 2 |
| `--video_sample_n_frames` | 视频采样帧数 | 81 |
| `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 |
| `--dataloader_num_workers` | DataLoader 工作进程数 | 8 |
| `--num_train_epochs` | 训练轮数 | 100 |
| `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 |
| `--learning_rate` | 初始学习率（generator） | 2e-06 |
| `--learning_rate_critic` | 初始学习率（critic） | 2e-06 |
| `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
| `--lr_warmup_steps` | 学习率 warmup 步数 | 100 |
| `--seed` | 随机种子 | 42 |
| `--output_dir` | 输出目录 | `output_dir_fastvideo_wan2.1_1.3b` |
| `--gradient_checkpointing` | 启用梯度检查点 | - |
| `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--vae_mini_batch` | VAE 编码 mini-batch 大小 | 1 |
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
| `--enable_bucket` | 启用分桶训练，不裁剪，按分辨率分组 | - |
| `--random_hw_adapt` | 将图片/视频自动缩放到 `[min_size, max_size]` 范围内的随机尺寸 | - |
| `--training_with_video_token_length` | 基于 token 长度训练，支持任意分辨率 | - |
| `--uniform_sampling` | 均匀时间步采样 | - |
| `--low_vram` | 低显存模式 | - |
| `--train_mode` | 训练模式：`normal`（标准）或 `i2v`（图生视频） | `normal` |
| `--resume_from_checkpoint` | 续训路径，使用 `"latest"` 自动选择最新 checkpoint | None |
| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
| `--validation_epochs` | 每 N 轮运行一次验证 | 5 |
| `--validation_prompts` | 用于视频生成验证的 prompt | `"A dog shaking head..."` |
| `--trainable_modules` | 可训练模块（`"."` 表示全部模块） | `"."` |
**蒸馏专用参数**：
| 参数 | 描述 | 示例值 |
|-----|------|-------|
| `--denoising_step_indices_list` | 去噪步列表（蒸馏核心参数） | `1000 750 500 250` |
| `--real_guidance_scale` | 打分用的 real guidance scale | 6.0 |
| `--fake_guidance_scale` | 打分用的 fake guidance scale | 0.0 |
| `--gen_update_interval` | generator 更新间隔 | 5 |
| `--negative_prompt` | 蒸馏用的负向 prompt | 中文负向 prompt |
| `--validation_paths` | I2V 模式的验证图片路径 | 图片路径列表 |
| `--train_sampling_steps` | 训练采样步数 | 1000 |
**VSA 稀疏度**：
- 学生/critic 的自注意力为 VSA 块稀疏。保留比 `student_topk` 在 `FastVideoWanTransformer3DModel` 内默认为 `0.20`（即 80% 稀疏度）；它是模型类默认值，而非命令行旗标。
- 可学习的压缩门控 `to_gate_compress` 以零初始化创建，并在 `--trainable_modules "."` 下自动变为可训练（其参数名与其它模块一样包含点号）。
- `real_score` 教师保持 dense（`WanTransformer3DModel`）且冻结，因此稀疏性是被蒸馏进学生，而非从教师拷贝。
- 推理时的 `student_topk` 须与 checkpoint 训练时的取值一致；修改它会重建注意力算子并重置已训练好的门控。
**采样尺寸配置指南**：
- `video_sample_size` 表示视频的分辨率尺寸；当 `random_hw_adapt` 为 True 时，表示视频与图片分辨率之间的最小值。
- `image_sample_size` 表示图片的分辨率尺寸；当 `random_hw_adapt` 为 True 时，表示视频与图片分辨率之间的最大值。
- `token_sample_size` 表示 `training_with_video_token_length` 为 True 时最大 token 长度对应的分辨率。
- 由于配置容易混淆，**若微调不需要任意分辨率**，建议将 `video_sample_size`、`image_sample_size`、`token_sample_size` 设为同一固定值，例如 **(320, 480, 512, 640, 960)**。
- **全部设为 320** 代表 **240P**。
- **全部设为 480** 代表 **320P**。
- **全部设为 640** 代表 **480P**。
- **全部设为 960** 代表 **720P**。
**其它参数说明**：
- `enable_bucket` 用于启用分桶训练。启用后模型不在中心裁剪图片与视频，而是按分辨率分桶后对完整图片与视频进行训练。
- `random_hw_adapt` 用于对图片与视频自动高宽缩放。启用后，训练图片的高宽以 `image_sample_size` 为最大、`min(video_sample_size, 512)` 为最小。
- `training_with_video_token_length` 指定按 token 长度训练。训练图片与视频的高宽以 `image_sample_size` 为最大、`video_sample_size` 为最小。
- `train_mode` 用于指定训练模式，可为 normal 或 i2v。若只需文本生成视频，保持 `normal`。
- `resume_from_checkpoint` 用于设置从之前的 checkpoint 续训。使用路径或 `"latest"` 自动选择最近可用的 checkpoint。
### 3.4 训练验证
你可以配置验证参数，在训练期间周期性生成测试视频，以便监控训练进度与模型质量。
**验证参数说明**：
| 参数 | 描述 | 推荐值 |
|------|------|--------|
| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
| `--validation_epochs` | 每 N 轮运行一次验证 | 5 |
| `--validation_prompts` | 用于视频生成验证的 prompt | 英文 prompt |
| `--validation_paths` | I2V 模式的验证图片路径（仅 i2v 模式） | `"asset/1.png"` |
**Normal 模式示例**（T2V 验证）：
```bash
--validation_steps=2000 \
--validation_epochs=5 \
--validation_prompts="A dog shaking head. The video is of high quality, and the view is very clear."
```
**说明**：
- 验证视频会保存到 `output_dir` 目录。
- 多 prompt 验证格式：`--validation_prompts "prompt1" "prompt2" "prompt3"`。
- 蒸馏模型验证使用 `denoising_step_indices_list` 定义的少步、并以无 CFG（`guidance_scale=1.0`）运行，与推理保持一致。
### 3.5 使用 FSDP 训练
**若使用 DeepSpeed-Zero-2 多卡训练时显存不足**，可切换到 FSDP 训练。`scripts/fastvideo/train_distill_wan2.1.sh` 是一个开箱即用的 FSDP launcher（它以 `FULL_SHARD` 包装 `WanAttentionBlock`）。
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.6 其他后端
#### 3.6.1 使用 DeepSpeed-Zero-3 进行训练
目前不强烈推荐 DeepSpeed Zero-3。在本仓库中，使用 FSDP 报错更少、更稳定。
DeepSpeed Zero-3 适合高分辨率下的 14B Wan。训练后可用以下命令得到最终模型：
```bash
python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/checkpoint-{your-num-steps}-outputs --max_shard_size 80GB --safe_serialization
```
训练 shell 命令如下：
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
#### 3.6.2 不使用 DeepSpeed 与 FSDP 训练
**不推荐这种方式，因为没有省显存后端时极易导致显存不足**。这里仅作为训练参考 shell 提供。
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.7 多机分布式训练
**适用于**：超大数据集、更快的训练速度
#### 3.7.1 环境配置
假设有 2 台机器，每台 8 张 GPU：
**机器 0（Master）**：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export MASTER_ADDR="192.168.1.100" # Master 机器 IP
export MASTER_PORT=10086
export WORLD_SIZE=2 # 机器总数
export NUM_PROCESS=16 # 总进程数 = 机器数 × 8
export RANK=0 # 当前机器 rank（0 或 1）
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
**机器 1（Worker）**：
```bash
export MASTER_ADDR="192.168.1.100" # 与 Master 相同
export MASTER_PORT=10086
export WORLD_SIZE=2
export NUM_PROCESS=16
export RANK=1 # 注意这里是 1
NCCL_DEBUG=INFO
# 使用与机器 0 相同的 accelerate launch 命令
```
#### 3.7.2 多机训练注意事项
- **网络要求**：
- 推荐 RDMA/InfiniBand（高性能）
- 没有 RDMA 时，添加环境变量：
```bash
export NCCL_IB_DISABLE=1
export NCCL_P2P_DISABLE=1
```
- **数据同步**：所有机器必须能访问相同的数据路径（NFS/共享存储）
---
## 四、推理测试
### 4.1 推理参数解析
**关键参数说明**（取自 `examples/fastvideo/predict_t2v_wan2.1.py`）：
| 参数 | 描述 | 示例值 |
|------|------|-------|
| `GPU_memory_mode` | GPU 显存模式，选项见下表 | `model_cpu_offload` |
| `ulysses_degree` | Ulysses 并行度（VSA 仅单卡，保持 1） | 1 |
| `ring_degree` | Ring 并行度（VSA 仅单卡，保持 1） | 1 |
| `config_path` | 配置文件路径 | `config/wan2.1/wan_civitai.yaml` |
| `model_name` | 基座模型路径（config + dense 主干 + VAE + T5） | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
| `student_topk` | VSA 保留比；须与训练时取值一致 | `0.20` |
| `sampler_name` | 采样器类型：`Flow`、`Flow_Unipc`、`Flow_DPM++` | `Flow_Unipc` |
| `shift` | 噪声调度 shift（官方 FastWan flow_shift） | `3.0` |
| `transformer_path` | FastWan 学生 Transformer 权重路径 | 训练 checkpoint / 官方 FastWan |
| `vae_path` | 训练好的 VAE 权重路径 | `None` |
| `lora_path` | LoRA 权重路径 | `None` |
| `sample_size` | 生成视频分辨率 `[height, width]` | `[480, 848]` |
| `video_length` | 生成帧数 | `81` |
| `fps` | 每秒帧数 | `16` |
| `weight_dtype` | 模型权重精度，不支持 bf16 的 GPU 用 `torch.float16` | `torch.bfloat16` |
| `prompt` | 描述生成内容的正向 prompt | `"A cinematic aerial shot..."` |
| `negative_prompt` | 用于规避特定内容的负向 prompt | 中文负向 prompt |
| `guidance_scale` | 引导强度（DMD 学生无 CFG） | 1.0 |
| `seed` | 复现用随机种子 | 1024 |
| `num_inference_steps` | 推理步数（FastWan 用 3 步） | 3 |
| `save_path` | 生成视频保存路径 | `samples/fastvideo-wan2.1-official` |
**GPU 显存模式说明**：
| 模式 | 描述 | 显存占用 |
|------|------|---------|
| `model_full_load` | 整个模型加载到 GPU | 最高 |
| `model_full_load_and_qfloat8` | 全量加载 + FP8 量化 | 高 |
| `model_cpu_offload` | 用后将模型卸载到 CPU | 中 |
| `model_cpu_offload_and_qfloat8` | CPU 卸载 + FP8 量化 | 中低 |
| `model_group_offload` | 层组在 CPU/CUDA 间切换 | 低 |
| `sequential_cpu_offload` | 逐层卸载（最慢） | 最低 |
### 4.2 文本生成视频（T2V）推理
运行单卡推理：
```bash
python examples/fastvideo/predict_t2v_wan2.1.py
```
根据需求编辑 `examples/fastvideo/predict_t2v_wan2.1.py`。首次推理重点关注以下关键参数，其它参数参见上文推理参数说明。
```python
# 根据 GPU 显存选择
GPU_memory_mode = "model_cpu_offload"
# 基座模型路径（config + dense 主干 + VAE + T5）
model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
# VSA 保留比；须与 checkpoint 训练时取值一致
student_topk = 0.20
# FastWan 学生权重路径（你训练的 checkpoint，或官方 FastWan 发布权重）
transformer_path = "models/Personalized_Model/FastWan2.1-T2V-1.3B.pth"
# FastWan 采样器 + shift（无 CFG 少步学生）
sampler_name = "Flow_Unipc"
shift = 3.0
guidance_scale = 1.0
num_inference_steps = 3
# 根据你的生成内容编写
prompt = "A cinematic aerial shot of a snow-covered mountain village at dawn..."
# ...
```
> 💡 官方 `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` checkpoint 采用 diffusers 键名；`examples/fastvideo/predict_t2v_wan2.1.py` 会在加载前将其 remap 到 VideoX-Fun civitai 布局（包含 VSA `to_gate_compress` 门控）。本训练指南产出的 checkpoint 已是原生布局。
### 4.3 仅支持单卡
VSA 块稀疏内核没有序列并行路径，因此 FastWan 学生只能在**单卡**上运行。请保持 `ulysses_degree = ring_degree = 1`；否则调用 `enable_multi_gpus_inference()` 会 raise。要在单卡上容纳更大分辨率，请通过 `GPU_memory_mode`（如 `model_cpu_offload` 或 `sequential_cpu_offload`）降低显存占用，而不是跨卡拆分。
---
## 五、更多资源
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
- **FastVideo GitHub**：https://github.com/hao-ai-lab/FastVideo
- **VSA 论文（Video Sparse Attention）**：https://arxiv.org/abs/2505.13389
scripts/fastvideo/train_distill_wan2.1.py
+2411
View File
File diff suppressed because it is too large Load Diff
scripts/fastvideo/train_distill_wan2.1.sh
+48
View File
@@ -0,0 +1,48 @@
# This shell script is used to distill a regular Wan2.1-1.3B into FastWan2.1-1.3B (VSA sparse student)
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1.md
+452
View File
@@ -0,0 +1,452 @@
# SparkDiffusion Sparse Distillation Training Guide (Wan2.1-T2V-1.3B)
This document provides the complete workflow for **SparkDiffusion** on Wan2.1-T2V-1.3B: a two-stage pipeline that first warms up a **sparse (RoLa) student** and then **distills** it into a few-step, CFG-free generator. It covers environment setup, data preparation, the two launchers (Stage-1 warm-up + Stage-2 DMD distill), and sparse few-step inference. Wan2.1-T2V-1.3B is a **single full-range expert**, so this is the leaner counterpart of the Wan2.2 MoE guide ([README_TRAIN_DISTILL_wan2.2.md](./README_TRAIN_DISTILL_wan2.2.md)): no MoE boundary, no high/low expert split, and no init-rollout relay.
> **What is SparkDiffusion?**
>
> SparkDiffusion (arXiv:2609.23153, *Mitigating the High-Sparsity Trap*) combines a **compensating sparse attention** (RoLa) with a **DMD** few-step distillation recipe. Wan2.1-T2V-1.3B is a single-expert diffusion transformer that spans the whole noise range, so the student is warmed up and distilled **once over the full rf-time range `[0, 1]`** — there is no MoE boundary to split on.
>
> - **Stage 1 — sparse warm-up** ([`train_finetune_sparse_wan2.1.py`](./train_finetune_sparse_wan2.1.py)): insert the RoLa sparse self-attention into the dense backbone and briefly fine-tune it with the plain step-local flow-matching objective, so distillation starts from a usable coarse prior instead of a zero-initialized sparse branch. This stage is intentionally **short**.
> - **Stage 2 — DMD distill** ([`train_distill_wan2.1.py`](./train_distill_wan2.1.py)): distill the sparse student into a few-step generator against the **dense** teacher. It is plain Distribution Matching Distillation (DMD) with three co-resident transformers — a trainable sparse `generator`, a frozen dense `real_score` teacher, and a trainable dense `fake_score` critic — rolled out over a discrete few-step grid. An optional **CrossDistill** PCM diversity term (`--loss_scale_div > 0`, see §3.3.1) can be layered on top to preserve sample diversity.
>
> | Stage | Launcher | Output |
> |-------|----------|--------|
> | **1 warm-up** | [`train_finetune_sparse_wan2.1.sh`](./train_finetune_sparse_wan2.1.sh) | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola` |
> | **2 distill** | [`train_distill_wan2.1.sh`](./train_distill_wan2.1.sh) | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill` |
>
> Run the two launchers in order: the Stage-2 generator is seeded from the Stage-1 warm-up product, and the single distilled checkpoint is the whole few-step student at inference.
---
## Table of Contents
- [1. Environment Setup](#1-environment-setup)
- [2. Data Preparation](#2-data-preparation)
- [3. Training](#3-training)
- [3.1 Download Pretrained Models](#31-download-pretrained-models)
- [3.2 Stage 1 — Sparse Warm-up](#32-stage-1--sparse-warm-up)
- [3.3 Stage 2 — DMD Distill](#33-stage-2--dmd-distill)
- [3.4 Common Training Parameters](#34-common-training-parameters)
- [3.5 FSDP Backend and Memory Notes](#35-fsdp-backend-and-memory-notes)
- [3.6 Other Backends](#36-other-backends)
- [3.7 Multi-Node Distributed Training](#37-multi-node-distributed-training)
- [4. Inference — Sparse Few-Step Student](#4-inference--sparse-few-step-student)
- [5. Additional Resources](#5-additional-resources)
---
## 1. Environment Setup
**Method 1: Using requirements.txt**
```bash
pip install -r requirements.txt
```
**Method 2: Manual Installation**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 💡 The RoLa sparse attention relies on the in-repo Triton block-sparse kernel (`videox_fun/models/attention_utils.py`), so a working Triton install is required for the sparse student.
**Method 3: Using Docker**
```bash
# pull image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# enter image
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 2. Data Preparation
SparkDiffusion consumes ordinary video-caption data (the same format as the standard Wan2.1 training).
### 2.1 Quick Test Dataset
```bash
# Download official example dataset
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 metadata.json Format
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
- `file_path`: video path (relative to `DATASET_NAME`, or absolute).
- `text`: caption / prompt.
- `type`: fixed as `"video"`.
- `width` / `height`: recommended for bucket training; add them with `scripts/process_json_add_width_and_height.py` if absent.
### 2.3 Relative vs Absolute Paths
```bash
# relative paths
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# absolute paths (data on NAS/OSS shared across machines)
export DATASET_NAME=""
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
```
---
## 3. Training
The full order is **Stage-1 warm-up → Stage-2 DMD distill**. Run the two launchers by hand, in that order.
### 3.1 Download Pretrained Models
```bash
# Create model directory
mkdir -p models/Diffusion_Transformer
# Download Wan2.1 T2V official weights
modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
```
The base checkpoint seeds all three Stage-2 transformers (the sparse `generator` is additionally initialized from the Stage-1 warm-up product). Its layout is the civitai single-file form read by `config/wan2.1/wan_civitai.yaml`: the transformer sits at the model root (`transformer_subpath: ./`), alongside `Wan2.1_VAE.pth`, the UMT5 encoder `models_t5_umt5-xxl-enc-bf16.pth` and the `google/umt5-xxl` tokenizer. The scheduler `shift` is `5.0`, supplied by the config `scheduler_kwargs`; the distill launcher does not override it.
### 3.2 Stage 1 — Sparse Warm-up
Stage 1 inserts the RoLa sparse self-attention and briefly warms it up with the step-local flow-matching objective. It is a **short** stage — a brief warm-up is enough to establish a usable prior for distillation, so the launcher caps it at `--max_train_steps=2000`. Do **not** judge this stage by its own sample quality; the sharpness is established in Stage 2.
Run [scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh](./train_finetune_sparse_wan2.1.sh):
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--max_train_steps=2000 \
--checkpointing_steps=200 \
--learning_rate=1e-4 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--random_hw_adapt \
--enable_bucket \
--uniform_sampling \
--trainable_modules "attn_op" \
--student_topk=0.10 \
--train_mode="normal" \
--resume_from_checkpoint="latest"
```
Or simply:
```bash
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh
```
Stage-1 keys:
| Flag | Meaning |
|------|---------|
| `--student_topk` | Ratio of key blocks kept per query in the sparse attention (sparsity = `1 - topk`). `0.10` ≈ 90% sparse. **Must match** the value used in Stage 2. |
| `--trainable_modules` | `"attn_op"` = Stage A (train only the sparse branch); `"."` = Stage B (full sparse finetune, load Stage A via `--transformer_path`). |
| `--max_train_steps` | Warm-up budget; keep it short. |
On completion the warm-up product is `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`.
### 3.3 Stage 2 — DMD Distill
Stage 2 distills the sparse student into a few-step, CFG-free generator against the dense teacher. `--sparse_student` turns the `generator` into `SparkWanTransformer3DModel` while the teacher and critic stay dense. `--generator_transformer_path` loads the Stage-1 warm-up checkpoint. Being a single full-range expert, the student distills over the whole rf-time range `[0, 1]` — there is no boundary segment and no init-rollout relay. This is the plain-DMD baseline: the generator rolls out a discrete few-step grid (`--denoising_step_indices_list 1000 750 500 250` with `--flow_euler_rollout`), and the generator/critic alternate via `--gen_update_interval`. On top of this discrete DMD host, an optional CrossDistill PCM diversity term (`--loss_scale_div > 0`, see §3.3.1 below) adds a phase-0 boundary-alignment loss to preserve sample diversity. Two further opt-in switches (§3.3.2) replace parts of this baseline — `--tdm` (the official SparkDiffusion TDM student loss) and `--joint_backward_timesteps` (define the few-step grid directly with rf knots) — but neither is enabled in the launcher, so the default run stays plain DMD.
Run [scripts/sparkdiffusion/train_distill_wan2.1.sh](./train_distill_wan2.1.sh):
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Stage-1 warm-up product. Env-overridable; --student_topk below MUST match Stage 1.
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--random_hw_adapt \
--enable_bucket \
--trainable_modules "." \
--train_mode="normal" \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--flow_euler_rollout \
--randomize_step_indices \
--gen_update_interval 5 \
--div_mode pcm_random \
--loss_scale_div 5.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 5.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--sparse_student \
--student_topk=0.10 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
```
Or simply:
```bash
bash scripts/sparkdiffusion/train_distill_wan2.1.sh
```
> ⚠️ **The distillation phase has no step cap.** It runs until you stop it manually (kill the `train_distill_wan2.1.py` processes); checkpoints are written every `--checkpointing_steps=200`.
#### 3.3.1 CrossDistill Diversity (PCM)
Setting `--loss_scale_div > 0` layers the official SparkDiffusion **CrossDistill** PCM diversity term on top of the discrete DMD host, to counter the diversity loss that pure distribution matching can cause. It is a self-contained additive loss that leaves the DMD rollout untouched:
- **What it does** — a phase-0 boundary-alignment MSE: the student's one-step jump from a noisy anchor to the phase-0 boundary (`--div_rf_boundary`, the official `backward_timesteps[0]`) must match a teacher K-step Euler rollout (`--dcm_skipping_interval_steps` sub-steps) followed by the student's own target jump. It runs on its **own continuous rf grid** (`--dcm_total_steps` subdivisions shifted by `--teacher_timestep_shift`), decoupled from the discrete DMD grid.
- **Anchor** — the paired **real video latent** (faithful to the official `_compute_pcm_div_loss`, which distills on real `x0`), encoded on demand only when diversity is on.
- **Merge** — `generator_loss = dmd_loss + loss_scale_div * div_loss` in a **single** `accelerator.backward` (FSDP allows one backward per wrapped module per iteration); the teacher rollout and the student target run under `no_grad` before the graded student forward.
- **Scale & tuning** — `--loss_scale_div` is in **per-element-mean** space (aligned with the mean-reduced `dmd_loss`); the official uses `10.0` in sum space, so `5.0` is the recommended start. On the first run watch `train_div_loss` vs `train_dmd_loss` in the logs and adjust so diversity does not overwhelm DMD; `--div_teacher_guidance` is the diversity teacher's CFG scale, independent of `--real_guidance_scale`. Set `--loss_scale_div 0` to recover the exact plain-DMD baseline.
#### 3.3.2 TDM Loss & Direct rf Grid (opt-in)
Two opt-in switches replace parts of the plain-DMD baseline above; **neither is enabled in `train_distill_wan2.1.sh`**, so the default run stays plain DMD.
- **`--tdm`** — replaces the plain-DMD student loss with **TDM** (Trajectory Distribution Matching, the official SparkDiffusion student loss). Three deltas from DMD: (1) the last rollout segment does not collapse to `x0` but takes one more rf-Euler step to `s_{i-1}` (the next grid knot, `0.0` when exhausted), so the anchor is `G_x_{t_{i-1}}`; (2) `D_time` is clamped strictly above `s_{i-1}`; (3) the fake/teacher `x0` preds are mapped back to `s_{i-1}` via the rf PF-ODE analytic step and the surrogate is compared in the `x_{t_{i-1}}` domain. When `s_{i-1}==0` (grid exhausted) it reduces **exactly** to DMD. Requires `--flow_euler_rollout`; mutually exclusive with `--dfd` / `--decoupled_dmd`; compatible with the diversity term.
- **`--joint_backward_timesteps`** — descending rf knots (in `[0,1]`) that define the few-step grid **directly**, replacing `--denoising_step_indices_list`. Wan2.1 is a single expert (no MoE boundary split), so the whole list is used: the grid becomes `[1.0] + [k for k in knots if 0<k<1]` and the rf values serve directly as Euler sigmas (`sigma == rf == timestep/1000`). Official Wan2.1 rCM knots: `0.933781 0.852895 0.608979`. When omitted, the grid falls back to `--denoising_step_indices_list`.
### 3.4 Common Training Parameters
| Parameter | Description | Example Value |
|-----------|-------------|---------------|
| `--config_path` | Wan2.1 config (single expert, no `boundary`) | `config/wan2.1/wan_civitai.yaml` |
| `--pretrained_model_name_or_path` | Wan2.1-T2V-1.3B base (transformer + VAE + UMT5) | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
| `--train_data_dir` / `--train_data_meta` | Dataset root / metadata JSON | see §2 |
| `--train_mode` | SparkDiffusion supports `normal` (T2V) only | `normal` |
| `--student_topk` | Sparse attention keep-ratio (sparsity = 1 − topk); Stage 1 and Stage 2 must match | `0.10` |
| `--sparse_student` | (Stage 2) build the generator as the sparse student | - |
| `--generator_transformer_path` | (Stage 2) Stage-1 warm-up checkpoint loaded into the generator | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` |
| `--denoising_step_indices_list` | Discrete few-step grid the generator rolls out; validation reuses it (`num_inference_steps` = its length) | `1000 750 500 250` |
| `--flow_euler_rollout` | Roll the generator in flow/velocity space with deterministic Euler steps (LightX2V-style), no x0 re-noising | - |
| `--tdm` | (opt-in) Use TDM — the official SparkDiffusion student loss — instead of plain DMD; reduces exactly to DMD when `s_{i-1}==0`. Requires `--flow_euler_rollout`, mutually exclusive with `--dfd` / `--decoupled_dmd` | - |
| `--joint_backward_timesteps` | (opt-in) Descending rf knots that define the few-step grid directly, replacing `--denoising_step_indices_list` (single expert: the whole list, no MoE split) | `0.933781 0.852895 0.608979` |
| `--uniform_sampling` | Sample the DMD noise level uniformly over the training schedule | - |
| `--randomize_step_indices` | Jitter the grid indices each step for robustness | - |
| `--gen_update_interval` | Update the generator once every N steps (the critic updates every step) | `5` |
| `--real_guidance_scale` | CFG scale of the frozen dense teacher (not overridden; uses the script default) | `6.0` |
| `--div_mode` | CrossDistill diversity mode: `pcm_random` (random phase-0 offset) or `pcm_fixed` | `pcm_random` |
| `--loss_scale_div` | Weight of the CrossDistill PCM diversity term (per-element-mean space); `0` disables it and recovers plain DMD | `5.0` |
| `--dcm_total_steps` | Subdivision of the diversity term's continuous rf grid (`dw = 1 / N`) | `48` |
| `--dcm_skipping_interval_steps` | Teacher Euler sub-steps `K` within phase 0 | `1` |
| `--teacher_timestep_shift` | rf timestep shift for the diversity grid (matches the Wan2.1 scheduler shift) | `5.0` |
| `--div_rf_boundary` | Phase-0 right boundary `rf_b` (official `backward_timesteps[0]`) | `0.933781` |
| `--div_teacher_guidance` | CFG scale of the diversity teacher rollout (independent of `--real_guidance_scale`) | `5.0` |
| `--learning_rate` / `--learning_rate_critic` | Generator / critic learning rate | `2e-06` / `2e-06` |
| `--checkpointing_steps` | Save checkpoint every N steps | `200` |
| `--max_grad_norm` | Gradient clipping threshold | `0.05` |
| `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` |
> 💡 The few-step grid is set by `--denoising_step_indices_list` (indices into the `--train_sampling_steps=1000` schedule); validation reuses the same grid, so `num_inference_steps` equals its length (`4` here).
### 3.5 FSDP Backend and Memory Notes
Stage 2 keeps **three** transformers co-resident (sparse generator + dense teacher + dense critic). At 1.3B each this is far lighter than the Wan2.2 14B experts, but the optimizer states still dominate, so the distill launcher keeps FSDP `FULL_SHARD` plus `--gradient_checkpointing`:
- `--fsdp_sharding_strategy "FULL_SHARD"` — shard params/grads/optimizer across GPUs.
- `--fsdp_backward_prefetch "BACKWARD_POST"` — lower median residency during backward.
- `--gradient_checkpointing` — recompute activations in backward.
- **Resolution, not frame count, is the OOM lever.** Memory scales with `seq_len`; the launchers bucket at `--image_sample_size=640 --video_sample_size=640 --token_sample_size=640` with `--random_hw_adapt --enable_bucket` at `--video_sample_n_frames=81`. If it OOMs, lower the sample sizes (e.g. `640 → 480`) rather than dropping frames.
- Frame count must satisfy `(N − 1) % 4 == 0`.
> 💡 At 1.3B the three transformers usually fit without CPU offload; if it OOMs, add `--fsdp_offload_params true` to the `accelerate launch` line to offload params/grads/optimizer to CPU (slower). `--low_vram` only parks the frozen side on CPU; it complements, but does not replace, sharding.
### 3.6 Other Backends
The Stage-2 launcher defaults to FSDP `FULL_SHARD`. The same [`train_distill_wan2.1.py`](./train_distill_wan2.1.py) also runs under DeepSpeed Zero-2 or without a sharding backend; at 1.3B these are more viable than for the Wan2.2 14B experts. Keep every script argument identical to §3.3 — only the `accelerate launch` header changes.
#### 3.6.1 Training with DeepSpeed Zero-2
DeepSpeed Zero-2 shards optimizer states and gradients but **not** parameters. For three co-resident 1.3B transformers it fits comfortably at the default resolution:
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/sparkdiffusion/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--random_hw_adapt \
--enable_bucket \
--trainable_modules "." \
--train_mode="normal" \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--flow_euler_rollout \
--randomize_step_indices \
--gen_update_interval 5 \
--div_mode pcm_random \
--loss_scale_div 5.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 5.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--sparse_student \
--student_topk=0.10 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
```
#### 3.6.2 Training Without DeepSpeed and FSDP
Drop the sharding backend entirely — three co-resident 1.3B transformers fit on a single high-memory GPU, though FSDP is still recommended for multi-GPU. Replace the `accelerate launch` header with the plain form and keep the same script arguments as §3.6.1:
```bash
accelerate launch --mixed_precision="bf16" \
scripts/sparkdiffusion/train_distill_wan2.1.py \
<the same --flag block as §3.6.1>
```
### 3.7 Multi-Node Distributed Training
The launchers are single-node. For multiple machines, add the following to the `accelerate launch` line (all machines must share the same dataset / base weights / `output_dir` over NFS/shared storage):
```bash
--main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT \
--num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK
```
Without RDMA, also set `NCCL_IB_DISABLE=1` and `NCCL_P2P_DISABLE=1`.
---
## 4. Inference — Sparse Few-Step Student
The Stage-2 product is a few-step, CFG-free sparse student. Being single-expert, inference loads **one** checkpoint (there is no high-noise expert). Build the wan2.1 sparse class `SparkWanTransformer3DModel`, set the training top-k, then load the distilled weights — the order matters because `set_student_topk` (re)creates the `attn_op` kernels that receive the trained RoLa params:
```python
from videox_fun.models import SparkWanTransformer3DModel
transformer = SparkWanTransformer3DModel.from_pretrained(
os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True, torch_dtype=weight_dtype,
)
transformer.set_student_topk(student_topk) # BEFORE load_state_dict
# transformer.load_state_dict(load_file(transformer_path), strict=False)
```
The Wan2.1 T2V entry [examples/wan2.1/predict_t2v.py](../../examples/wan2.1/predict_t2v.py) already runs few-step CFG-free sampling (`sampler_name="Flow"`, `num_inference_steps=4`, `guidance_scale=1.0`) under `WanPipeline`; adapt it by swapping the dense `WanTransformer3DModel` for the sparse build above. Set:
| Parameter | Description | Example Value |
|-----------|-------------|---------------|
| `GPU_memory_mode` | Memory mode | `model_full_load` |
| `config_path` | Wan2.1 config | `config/wan2.1/wan_civitai.yaml` |
| `model_name` | Wan2.1-T2V-1.3B base (transformer / VAE / UMT5) | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
| `student_topk` | **Must equal** the training `--student_topk` | `0.10` |
| `transformer_path` | Sparse student checkpoint (single expert) | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
| `sampler_name` | `Flow` matches the few-step CFG-free student | `Flow` |
| `num_inference_steps` | The student's few-step count | `4` |
| `guidance_scale` | Distilled student is CFG-free | `1.0` |
| `sample_size` / `video_length` / `fps` | `[height, width]` / frames / fps | `[480, 832]` / `81` / `16` |
```bash
python examples/wan2.1/predict_t2v.py
```
> 💡 Load order is `from_pretrained → set_student_topk → load_state_dict`; reversing it re-zeroes the trained RoLa weights. A sparse checkpoint is confirmed when the `attn_op.*` keys load with **0** unexpected keys.
> 📌 The bundled [examples/sparkdiffusion/predict_t2v.py](../../examples/sparkdiffusion/predict_t2v.py) targets the Wan2.2 two-expert MoE student (`SparkWan2_2Transformer3DModel` + `Wan2_2Pipeline`); the single-expert Wan2.1 student uses `SparkWanTransformer3DModel` + `WanPipeline` with one checkpoint.
> 📌 **Stage 3** of the paper (the FP8/INT8 fused quantized inference kernel) is intentionally not integrated in this repository.
---
## 5. Additional Resources
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
- **Paper**: SparkDiffusion, *Mitigating the High-Sparsity Trap* (arXiv:2609.23153)
- **Sparse student model**: `videox_fun/models/spark_transformer3d.py` (`SparkWanTransformer3DModel`)
- **RoLa sparse attention kernel**: `videox_fun/models/attention_utils.py`
- **Wan2.2 MoE counterpart**: [README_TRAIN_DISTILL_wan2.2.md](./README_TRAIN_DISTILL_wan2.2.md)
scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1_zh-CN.md
+452
View File
@@ -0,0 +1,452 @@
# SparkDiffusion 稀疏蒸馏训练指南（Wan2.1-T2V-1.3B）
本文档给出 Wan2.1-T2V-1.3B 上 **SparkDiffusion** 的完整流程：一条两阶段管线，先暖启一个**稀疏（RoLa）学生**，再把它**蒸馏**成少步、CFG-free 的生成器。内容涵盖环境配置、数据准备、两个启动脚本（Stage-1 暖启 + Stage-2 DMD 蒸馏）、以及稀疏少步推理。Wan2.1-T2V-1.3B 是**单专家、覆盖全噪声区间**的模型，因此本文是 Wan2.2 MoE 指南（[README_TRAIN_DISTILL_wan2.2_zh-CN.md](./README_TRAIN_DISTILL_wan2.2_zh-CN.md)）的精简版：没有 MoE 分界、没有高/低噪专家拆分、也没有 init-rollout 接力。
> **什么是 SparkDiffusion？**
>
> SparkDiffusion（arXiv:2609.23153，《Mitigating the High-Sparsity Trap》）把**补偿式稀疏注意力**（RoLa）与 **DMD** 少步蒸馏配方结合起来。Wan2.1-T2V-1.3B 是覆盖整个噪声区间的单专家 diffusion transformer，因此学生在**整个 rf-time 区间 `[0, 1]` 上一次性**完成暖启与蒸馏——没有可供拆分的 MoE 分界。
>
> - **Stage 1 — 稀疏暖启**（[`train_finetune_sparse_wan2.1.py`](./train_finetune_sparse_wan2.1.py)）：把 RoLa 稀疏自注意力插入 dense 主干，用标准的 step-local flow-matching 目标短暂微调，使蒸馏从一个可用的粗先验出发，而不是从零初始化的稀疏分支出发。该阶段刻意保持**短**。
> - **Stage 2 — DMD 蒸馏**（[`train_distill_wan2.1.py`](./train_distill_wan2.1.py)）：对着 **dense** teacher 把稀疏学生蒸馏成少步生成器。它是最普通的分布匹配蒸馏（DMD），含三个同驻网络——可训练的稀疏 `generator`、冻结的 dense `real_score` teacher、可训练的 dense `fake_score` critic——在离散的少步网格上 rollout。其上可叠加一个可选的 **CrossDistill** PCM 多样性项（`--loss_scale_div > 0`，见 §3.3.1），用于保持样本多样性。
>
> | 阶段 | 启动脚本 | 产物 |
> |-------|----------|--------|
> | **1 暖启** | [`train_finetune_sparse_wan2.1.sh`](./train_finetune_sparse_wan2.1.sh) | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola` |
> | **2 蒸馏** | [`train_distill_wan2.1.sh`](./train_distill_wan2.1.sh) | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill` |
>
> 按顺序运行两个启动脚本：Stage-2 的 generator 由 Stage-1 暖启产物播种，而单个蒸馏 checkpoint 在推理时即为完整的少步学生。
---
## 目录
- [一、环境配置](#一环境配置)
- [二、数据准备](#二数据准备)
- [三、训练](#三训练)
- [3.1 下载预训练模型](#31-下载预训练模型)
- [3.2 Stage 1 — 稀疏暖启](#32-stage-1--稀疏暖启)
- [3.3 Stage 2 — DMD 蒸馏](#33-stage-2--dmd-蒸馏)
- [3.4 通用训练参数](#34-通用训练参数)
- [3.5 FSDP 后端与显存说明](#35-fsdp-后端与显存说明)
- [3.6 其他后端](#36-其他后端)
- [3.7 多机分布式训练](#37-多机分布式训练)
- [四、推理 — 稀疏少步学生](#四推理--稀疏少步学生)
- [五、更多资源](#五更多资源)
---
## 一、环境配置
**方式一：使用 requirements.txt**
```bash
pip install -r requirements.txt
```
**方式二：手动安装**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 💡 RoLa 稀疏注意力依赖仓库内的 Triton block-sparse 内核（`videox_fun/models/attention_utils.py`），因此稀疏学生需要一个可用的 Triton 安装。
**方式三：使用 Docker**
```bash
# 拉取镜像
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# 进入镜像
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 二、数据准备
SparkDiffusion 使用普通的「视频-字幕」数据（与标准 Wan2.1 训练格式一致）。
### 2.1 快速测试数据集
```bash
# 下载官方示例数据集
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 metadata.json 格式
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
- `file_path`：视频路径（相对于 `DATASET_NAME`，或绝对路径）。
- `text`：字幕 / 提示词。
- `type`：固定为 `"video"`。
- `width` / `height`：建议提供以启用 bucket 训练；若缺失可用 `scripts/process_json_add_width_and_height.py` 补齐。
### 2.3 相对路径与绝对路径
```bash
# 相对路径
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# 绝对路径（数据在 NAS/OSS 上、跨机共享）
export DATASET_NAME=""
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
```
---
## 三、训练
完整顺序为 **Stage-1 暖启 → Stage-2 DMD 蒸馏**。请手动按此顺序运行两个启动脚本。
### 3.1 下载预训练模型
```bash
# 创建模型目录
mkdir -p models/Diffusion_Transformer
# 下载 Wan2.1 T2V 官方权重
modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
```
基座权重会播种 Stage-2 的全部三个网络（稀疏 `generator` 还会额外用 Stage-1 暖启产物初始化）。其目录结构为 `config/wan2.1/wan_civitai.yaml` 读取的 civitai 单文件形态：transformer 位于模型根目录（`transformer_subpath: ./`），与 `Wan2.1_VAE.pth`、UMT5 编码器 `models_t5_umt5-xxl-enc-bf16.pth` 以及 `google/umt5-xxl` 分词器并列。调度器 `shift` 为 `5.0`，由 config 的 `scheduler_kwargs` 提供；蒸馏启动脚本不再覆盖它。
### 3.2 Stage 1 — 稀疏暖启
Stage 1 插入 RoLa 稀疏自注意力，并用 step-local flow-matching 目标短暂暖启。这是一个**短**阶段——短暂暖启即足以为蒸馏建立可用先验，因此启动脚本把它限制在 `--max_train_steps=2000`。**不要**用该阶段自身的采样质量来评判它；清晰度是在 Stage 2 建立的。
用 [scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh](./train_finetune_sparse_wan2.1.sh) 运行：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--max_train_steps=2000 \
--checkpointing_steps=200 \
--learning_rate=1e-4 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--random_hw_adapt \
--enable_bucket \
--uniform_sampling \
--trainable_modules "attn_op" \
--student_topk=0.10 \
--train_mode="normal" \
--resume_from_checkpoint="latest"
```
或直接：
```bash
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh
```
Stage-1 关键参数：
| 旗标 | 含义 |
|------|---------|
| `--student_topk` | 稀疏注意力中每个 query 保留的 key block 比例（稀疏度 = `1 - topk`）。`0.10` ≈ 90% 稀疏。**必须与** Stage 2 一致。 |
| `--trainable_modules` | `"attn_op"` = Stage A（只训练稀疏分支）；`"."` = Stage B（稀疏网全量微调，用 `--transformer_path` 载入 Stage A）。 |
| `--max_train_steps` | 暖启预算；保持简短。 |
完成后暖启产物为 `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`。
### 3.3 Stage 2 — DMD 蒸馏
Stage 2 对着 dense teacher 把稀疏学生蒸馏成少步、CFG-free 的生成器。`--sparse_student` 把 `generator` 变为 `SparkWanTransformer3DModel`，而 teacher 与 critic 保持 dense。`--generator_transformer_path` 载入 Stage-1 暖启 checkpoint。作为单一全区间专家，学生在整个 rf-time 区间 `[0, 1]` 上完成蒸馏——没有分界段、也没有 init-rollout 接力。这是最普通的 DMD 基线：generator 在离散的少步网格上 rollout（`--denoising_step_indices_list 1000 750 500 250` 配 `--flow_euler_rollout`），并通过 `--gen_update_interval` 让 generator/critic 交替更新。在这个离散 DMD 宿主之上，可选的 CrossDistill PCM 多样性项（`--loss_scale_div > 0`，见下方 §3.3.1）额外引入一个 phase-0 边界对齐损失以保持样本多样性。另有两个可选开关（见 §3.3.2）会替换该基线的一部分：`--tdm`（官方 SparkDiffusion TDM 学生损失）与 `--joint_backward_timesteps`（用 rf 节点直接定义少步网格）；但启动脚本中两者都未启用，故默认运行仍是 plain DMD。
用 [scripts/sparkdiffusion/train_distill_wan2.1.sh](./train_distill_wan2.1.sh) 运行：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Stage-1 暖启产物。可用环境变量覆盖；下方 --student_topk 必须与 Stage 1 一致。
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--random_hw_adapt \
--enable_bucket \
--trainable_modules "." \
--train_mode="normal" \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--flow_euler_rollout \
--randomize_step_indices \
--gen_update_interval 5 \
--div_mode pcm_random \
--loss_scale_div 5.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 5.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--sparse_student \
--student_topk=0.10 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
```
或直接：
```bash
bash scripts/sparkdiffusion/train_distill_wan2.1.sh
```
> ⚠️ **蒸馏阶段没有步数上限。** 它会一直运行直到你手动停止（kill 掉 `train_distill_wan2.1.py` 进程）；checkpoint 每 `--checkpointing_steps=200` 步写一次。
#### 3.3.1 CrossDistill 多样性（PCM）
设置 `--loss_scale_div > 0` 会在离散 DMD 宿主之上叠加官方 SparkDiffusion 的 **CrossDistill** PCM 多样性项，用于抵消纯分布匹配可能带来的多样性损失。它是一个自包含的加性损失，不改动 DMD rollout：
- **作用**——一个 phase-0 边界对齐 MSE：学生从高噪 anchor 一步跳到 phase-0 边界（`--div_rf_boundary`，即官方 `backward_timesteps[0]`）的结果，必须与「teacher K 步 Euler rollout（`--dcm_skipping_interval_steps` 个子步）+ 学生自身的 target 跳转」相匹配。它跑在**自己的连续 rf 网格**上（`--dcm_total_steps` 细分、按 `--teacher_timestep_shift` 平移），与离散 DMD 网格解耦。
- **anchor**——配对的**真实视频 latent**（忠实官方 `_compute_pcm_div_loss`，其在真实 `x0` 上蒸馏），仅在多样性开启时按需编码。
- **合并**——`generator_loss = dmd_loss + loss_scale_div * div_loss`，用**单次** `accelerator.backward`（FSDP 每个 wrapped module 每轮只允许一次 backward）；teacher rollout 与学生 target 在带梯度的学生前向之前于 `no_grad` 下运行。
- **scale 与调参**——`--loss_scale_div` 处于**逐元素均值**空间（与 mean 归约的 `dmd_loss` 对齐）；官方在 sum 空间用 `10.0`，故推荐起点为 `5.0`。首轮请观察日志中的 `train_div_loss` 与 `train_dmd_loss` 并调整，使多样性项不压过 DMD；`--div_teacher_guidance` 是多样性 teacher 的 CFG 尺度，独立于 `--real_guidance_scale`。设 `--loss_scale_div 0` 可退回完全一致的 plain-DMD 基线。
#### 3.3.2 TDM 损失与直接 rf 网格（可选）
两个可选开关会替换上方 plain-DMD 基线的一部分；**`train_distill_wan2.1.sh` 中两者都未启用**，故默认运行仍是 plain DMD。
- **`--tdm`**——把 plain-DMD 学生损失替换为 **TDM**（Trajectory Distribution Matching，官方 SparkDiffusion 学生损失）。相对 DMD 有三处差异：(1) rollout 末段不塌缩到 `x0`，而是再走一步 rf-Euler 到 `s_{i-1}`（下一个网格节点，耗尽时为 `0.0`），故锚点为 `G_x_{t_{i-1}}`；(2) `D_time` 被严格 clamp 在 `s_{i-1}` 之上；(3) fake/teacher 的 `x0` 预测经 rf PF-ODE 解析步映回 `s_{i-1}`，surrogate 在 `x_{t_{i-1}}` 域比较。当 `s_{i-1}==0`（网格耗尽）时**精确**退化为 DMD。需 `--flow_euler_rollout`；与 `--dfd` / `--decoupled_dmd` 互斥；与多样性项兼容。
- **`--joint_backward_timesteps`**——降序 rf 节点（位于 `[0,1]`），**直接**定义少步网格，替代 `--denoising_step_indices_list`。Wan2.1 是单专家（不做 MoE 分界切分），故整份列表都用上：网格变为 `[1.0] + [k for k in knots if 0<k<1]`，rf 值直接充当 Euler sigma（`sigma == rf == timestep/1000`）。官方 Wan2.1 rCM 节点：`0.933781 0.852895 0.608979`。省略时，网格退回 `--denoising_step_indices_list`。
### 3.4 通用训练参数
| 参数 | 说明 | 示例值 |
|-----------|-------------|---------------|
| `--config_path` | Wan2.1 config（单专家，无 `boundary`） | `config/wan2.1/wan_civitai.yaml` |
| `--pretrained_model_name_or_path` | Wan2.1-T2V-1.3B 基座（transformer + VAE + UMT5） | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
| `--train_data_dir` / `--train_data_meta` | 数据集根目录 / metadata JSON | 见 §2 |
| `--train_mode` | SparkDiffusion 仅支持 `normal`（T2V） | `normal` |
| `--student_topk` | 稀疏注意力保留比例（稀疏度 = 1 − topk）；Stage 1 与 Stage 2 必须一致 | `0.10` |
| `--sparse_student` | （Stage 2）把 generator 建成稀疏学生 | - |
| `--generator_transformer_path` | （Stage 2）载入 generator 的 Stage-1 暖启 checkpoint | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` |
| `--denoising_step_indices_list` | generator rollout 的离散少步网格；validation 复用同一网格（`num_inference_steps` = 其长度） | `1000 750 500 250` |
| `--flow_euler_rollout` | 让 generator 在 flow/速度空间用确定性 Euler 步 rollout（LightX2V 式），不做 x0 重新加噪 | - |
| `--tdm` | （可选）用 TDM（官方 SparkDiffusion 学生损失）替代 plain DMD；当 `s_{i-1}==0` 时精确退化为 DMD。需 `--flow_euler_rollout`，与 `--dfd` / `--decoupled_dmd` 互斥 | - |
| `--joint_backward_timesteps` | （可选）降序 rf 节点，直接定义少步网格，替代 `--denoising_step_indices_list`（单专家：整份列表，不做 MoE 切分） | `0.933781 0.852895 0.608979` |
| `--uniform_sampling` | 在训练调度上均匀采样 DMD 噪声档位 | - |
| `--randomize_step_indices` | 每步抖动网格索引以增强鲁棒性 | - |
| `--gen_update_interval` | 每 N 步更新一次 generator（critic 每步都更新） | `5` |
| `--real_guidance_scale` | 冻结 dense teacher 的 CFG 尺度（不覆盖，走脚本默认值） | `6.0` |
| `--div_mode` | CrossDistill 多样性模式：`pcm_random`（随机 phase-0 偏移）或 `pcm_fixed` | `pcm_random` |
| `--loss_scale_div` | CrossDistill PCM 多样性项权重（逐元素均值空间）；`0` 关闭并退回 plain DMD | `5.0` |
| `--dcm_total_steps` | 多样性项连续 rf 网格的细分数（`dw = 1 / N`） | `48` |
| `--dcm_skipping_interval_steps` | phase 0 内 teacher Euler 子步数 `K` | `1` |
| `--teacher_timestep_shift` | 多样性网格的 rf timestep 平移（与 Wan2.1 scheduler shift 一致） | `5.0` |
| `--div_rf_boundary` | phase-0 右边界 `rf_b`（官方 `backward_timesteps[0]`） | `0.933781` |
| `--div_teacher_guidance` | 多样性 teacher rollout 的 CFG 尺度（独立于 `--real_guidance_scale`） | `5.0` |
| `--learning_rate` / `--learning_rate_critic` | 生成器 / critic 学习率 | `2e-06` / `2e-06` |
| `--checkpointing_steps` | 每 N 步保存 checkpoint | `200` |
| `--max_grad_norm` | 梯度裁剪阈值 | `0.05` |
| `--resume_from_checkpoint` | 续训路径或 `"latest"` | `latest` |
> 💡 少步网格由 `--denoising_step_indices_list` 设定（是对 `--train_sampling_steps=1000` 调度的索引）；validation 复用同一网格，因此 `num_inference_steps` 等于其长度（此处为 `4`）。
### 3.5 FSDP 后端与显存说明
Stage 2 需让**三个** transformer 同驻（稀疏 generator + dense teacher + dense critic）。每个 1.3B 远比 Wan2.2 的 14B 专家轻量，但优化器态仍占主导，因此蒸馏启动脚本保留 FSDP `FULL_SHARD` 与 `--gradient_checkpointing`：
- `--fsdp_sharding_strategy "FULL_SHARD"`——在多卡间分片 param/grad/optimizer。
- `--fsdp_backward_prefetch "BACKWARD_POST"`——降低 backward 期间的中位驻留。
- `--gradient_checkpointing`——在 backward 中重算激活。
- **OOM 的杠杆是分辨率、而非帧数。** 显存随 `seq_len` 增长；启动脚本以 `--image_sample_size=640 --video_sample_size=640 --token_sample_size=640` 配 `--random_hw_adapt --enable_bucket`、在 `--video_sample_n_frames=81` 下分桶。若 OOM，请降低 sample size（例如 `640 → 480`），而不是减少帧数。
- 帧数须满足 `(N − 1) % 4 == 0`。
> 💡 1.3B 下三个 transformer 通常无需 CPU offload 即可放下；若 OOM，在 `accelerate launch` 行加上 `--fsdp_offload_params true` 把 param/grad/optimizer offload 到 CPU（更慢）。`--low_vram` 只把冻结侧停到 CPU；它是对分片的补充，而非替代。
### 3.6 其他后端
Stage-2 启动脚本默认使用 FSDP `FULL_SHARD`。同一个 [`train_distill_wan2.1.py`](./train_distill_wan2.1.py) 也可在 DeepSpeed Zero-2 下运行，或不使用任何分片后端；在 1.3B 规模下这些方案比 Wan2.2 的 14B 专家更可行。请保持脚本参数与 §3.3 完全一致——只改动 `accelerate launch` 头部。
#### 3.6.1 使用 DeepSpeed-Zero-2 进行训练
DeepSpeed Zero-2 只分片优化器态与梯度、**不**分片参数。对三个同驻的 1.3B transformer 而言，它在默认分辨率下即可从容放下：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/sparkdiffusion/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--random_hw_adapt \
--enable_bucket \
--trainable_modules "." \
--train_mode="normal" \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--flow_euler_rollout \
--randomize_step_indices \
--gen_update_interval 5 \
--div_mode pcm_random \
--loss_scale_div 5.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 5.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--sparse_student \
--student_topk=0.10 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
```
#### 3.6.2 不使用 DeepSpeed 与 FSDP 训练
完全去掉分片后端——三个同驻的 1.3B transformer 可放进单张高显存 GPU，不过多卡仍推荐 FSDP。把 `accelerate launch` 头部替换为朴素形式，脚本参数与 §3.6.1 保持一致：
```bash
accelerate launch --mixed_precision="bf16" \
scripts/sparkdiffusion/train_distill_wan2.1.py \
<与 §3.6.1 相同的 --flag 参数块>
```
### 3.7 多机分布式训练
启动脚本为单机。多机时在 `accelerate launch` 行追加下列参数（所有机器须通过 NFS/共享存储共享同一份数据集 / 基座权重 / `output_dir`）：
```bash
--main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT \
--num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK
```
无 RDMA 时，另设 `NCCL_IB_DISABLE=1` 与 `NCCL_P2P_DISABLE=1`。
---
## 四、推理 — 稀疏少步学生
Stage-2 产物是一个少步、CFG-free 的稀疏学生。作为单专家模型，推理只载入**一个** checkpoint（没有高噪专家）。构建 wan2.1 稀疏类 `SparkWanTransformer3DModel`，设置训练时的 top-k，再载入蒸馏权重——顺序很关键，因为 `set_student_topk` 会（重）建接收训练好 RoLa 参数的 `attn_op` 内核：
```python
from videox_fun.models import SparkWanTransformer3DModel
transformer = SparkWanTransformer3DModel.from_pretrained(
os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True, torch_dtype=weight_dtype,
)
transformer.set_student_topk(student_topk) # 必须在 load_state_dict 之前
# transformer.load_state_dict(load_file(transformer_path), strict=False)
```
Wan2.1 T2V 入口 [examples/wan2.1/predict_t2v.py](../../examples/wan2.1/predict_t2v.py) 已在 `WanPipeline` 下运行少步 CFG-free 采样（`sampler_name="Flow"`、`num_inference_steps=4`、`guidance_scale=1.0`）；把 dense 的 `WanTransformer3DModel` 换成上面的稀疏构建即可。需设置：
| 参数 | 说明 | 示例值 |
|-----------|-------------|---------------|
| `GPU_memory_mode` | 显存模式 | `model_full_load` |
| `config_path` | Wan2.1 config | `config/wan2.1/wan_civitai.yaml` |
| `model_name` | Wan2.1-T2V-1.3B 基座（transformer / VAE / UMT5） | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
| `student_topk` | **必须等于**训练时的 `--student_topk` | `0.10` |
| `transformer_path` | 稀疏学生 checkpoint（单专家） | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
| `sampler_name` | `Flow` 与少步 CFG-free 学生匹配 | `Flow` |
| `num_inference_steps` | 学生的少步步数 | `4` |
| `guidance_scale` | 蒸馏学生为 CFG-free | `1.0` |
| `sample_size` / `video_length` / `fps` | `[高, 宽]` / 帧数 / 帧率 | `[480, 832]` / `81` / `16` |
```bash
python examples/wan2.1/predict_t2v.py
```
> 💡 载入顺序为 `from_pretrained → set_student_topk → load_state_dict`；顺序颠倒会把训练好的 RoLa 权重重新清零。当 `attn_op.*` 键载入且 unexpected keys 为 **0** 时，即确认为稀疏 checkpoint。
> 📌 仓库自带的 [examples/sparkdiffusion/predict_t2v.py](../../examples/sparkdiffusion/predict_t2v.py) 面向 Wan2.2 双专家 MoE 学生（`SparkWan2_2Transformer3DModel` + `Wan2_2Pipeline`）；单专家的 Wan2.1 学生使用 `SparkWanTransformer3DModel` + `WanPipeline`，只载入一个 checkpoint。
> 📌 论文的 **Stage 3**（FP8/INT8 融合量化推理核）在本仓库中有意未接入。
---
## 五、更多资源
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
- **论文**：SparkDiffusion，《Mitigating the High-Sparsity Trap》（arXiv:2609.23153）
- **稀疏学生模型**：`videox_fun/models/spark_transformer3d.py`（`SparkWanTransformer3DModel`）
- **RoLa 稀疏注意力内核**：`videox_fun/models/attention_utils.py`
- **Wan2.2 MoE 对应文档**：[README_TRAIN_DISTILL_wan2.2_zh-CN.md](./README_TRAIN_DISTILL_wan2.2_zh-CN.md)
scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2.md
+495
View File
@@ -0,0 +1,495 @@
# SparkDiffusion Sparse Distillation Training Guide
This document provides the complete workflow for **SparkDiffusion** on Wan2.2-T2V-A14B: a two-stage pipeline that first warms up a **sparse (RoLa) student** and then **distills** it into a few-step, CFG-free generator. It covers environment setup, data preparation, the four launchers (Stage-1 warm-up + Stage-2 CrossDistill, per MoE expert), the one-shot orchestrator, and the `predict_t2v.py` inference acceptance test.
> **What is SparkDiffusion?**
>
> SparkDiffusion (arXiv:2609.23153, *Mitigating the High-Sparsity Trap*) combines a **compensating sparse attention** (RoLa) with a **CrossDistill** few-step recipe. Wan2.2-T2V-A14B is a two-expert MoE (a high-noise expert and a low-noise expert, split at the config `boundary: 0.875`); each expert is warmed up and distilled **on its own noise range**, never jointly.
>
> - **Stage 1 — sparse warm-up** ([`train_finetune_sparse_wan2.2.py`](./train_finetune_sparse_wan2.2.py)): insert the RoLa sparse self-attention into the dense backbone and briefly fine-tune it with the plain step-local flow-matching objective, so distillation starts from a usable coarse prior instead of a zero-initialized sparse branch. This stage is intentionally **short**.
> - **Stage 2 — CrossDistill** ([`train_distill_wan2.2.py`](./train_distill_wan2.2.py)): distill the sparse student into a few-step generator against the **dense** teacher. It is Distribution Matching Distillation (DMD) with three co-resident transformers — a trainable sparse `generator`, a frozen dense `real_score` teacher, and a trainable dense `fake_score` critic. The **high-noise** expert adds a PCM-style diversity term; the **low-noise** expert is div-free.
>
> | Stage | High-noise expert | Low-noise expert | Output |
> |-------|-------------------|------------------|--------|
> | **1 warm-up** | [`train_finetune_sparse_wan2.2_high.sh`](./train_finetune_sparse_wan2.2_high.sh) | [`train_finetune_sparse_wan2.2_low.sh`](./train_finetune_sparse_wan2.2_low.sh) | `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola` |
> | **2 distill** | [`train_distill_wan2.2_high.sh`](./train_distill_wan2.2_high.sh) | [`train_distill_wan2.2_low.sh`](./train_distill_wan2.2_low.sh) | `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola_distill` |
>
> The two experts are distilled separately and both are loaded at inference to form the full few-step MoE student.
---
## Table of Contents
- [1. Environment Setup](#1-environment-setup)
- [2. Data Preparation](#2-data-preparation)
- [3. Training](#3-training)
- [3.1 Download Pretrained Models](#31-download-pretrained-models)
- [3.2 Stage 1 — Sparse Warm-up](#32-stage-1--sparse-warm-up)
- [3.3 Stage 2 — CrossDistill](#33-stage-2--crossdistill)
- [3.4 Common Training Parameters](#34-common-training-parameters)
- [3.5 FSDP Backend and Memory Notes](#35-fsdp-backend-and-memory-notes)
- [3.6 Other Backends](#36-other-backends)
- [3.7 Multi-Node Distributed Training](#37-multi-node-distributed-training)
- [4. Inference — Sparse Few-Step Student](#4-inference--sparse-few-step-student)
- [5. Additional Resources](#5-additional-resources)
---
## 1. Environment Setup
**Method 1: Using requirements.txt**
```bash
pip install -r requirements.txt
```
**Method 2: Manual Installation**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 💡 The RoLa sparse attention relies on the in-repo Triton block-sparse kernel (`videox_fun/models/attention_utils.py`), so a working Triton install is required for the sparse student.
**Method 3: Using Docker**
```bash
# pull image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# enter image
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 2. Data Preparation
SparkDiffusion consumes ordinary video-caption data (the same format as the standard Wan2.2 training).
### 2.1 Quick Test Dataset
```bash
# Download official example dataset
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 metadata.json Format
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
- `file_path`: video path (relative to `DATASET_NAME`, or absolute).
- `text`: caption / prompt.
- `type`: fixed as `"video"`.
- `width` / `height`: recommended for bucket training; add them with `scripts/process_json_add_width_and_height.py` if absent.
### 2.3 Relative vs Absolute Paths
```bash
# relative paths
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# absolute paths (data on NAS/OSS shared across machines)
export DATASET_NAME=""
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
```
---
## 3. Training
The full order is **Stage-1 high → Stage-2 high → Stage-1 low → Stage-2 low**. Run the four launchers by hand, in that order.
### 3.1 Download Pretrained Models
```bash
# Create model directory
mkdir -p models/Diffusion_Transformer
# Download Wan2.2 T2V official weights
modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-T2V-A14B
```
The base checkpoint seeds all three Stage-2 transformers (the sparse `generator` is additionally initialized from the Stage-1 warm-up product). Its layout is `high_noise_model/`, `low_noise_model/`, the VAE and the UMT5 text encoder; the MoE boundary is read from `config/wan2.2/wan_civitai_t2v.yaml` (`boundary: 0.875`).
### 3.2 Stage 1 — Sparse Warm-up
Stage 1 inserts the RoLa sparse self-attention and briefly warms it up with the step-local flow-matching objective. It is a **short** stage — a brief warm-up is enough to establish a usable prior for distillation, so the launchers cap it at `--max_train_steps=2000`. Do **not** judge this stage by its own sample quality; the sharpness is established in Stage 2.
Run the high-noise expert with [scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh](./train_finetune_sparse_wan2.2_high.sh):
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--max_train_steps=2000 \
--checkpointing_steps=200 \
--learning_rate=1e-4 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--enable_bucket \
--uniform_sampling \
--trainable_modules "attn_op" \
--student_topk=0.05 \
--train_mode="normal" \
--boundary_type="high" \
--resume_from_checkpoint="latest"
```
Or simply:
```bash
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh
```
The low-noise expert uses the same recipe with `--boundary_type="low"` and its own `--output_dir`:
```bash
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.2_low.sh
```
Stage-1 keys:
| Flag | Meaning |
|------|---------|
| `--boundary_type` | `high` / `low` — selects the MoE expert to warm up (from the config boundary). |
| `--student_topk` | Ratio of key blocks kept per query in the sparse attention (sparsity = `1 - topk`). `0.05` ≈ 95% sparse. **Must match** the value used in Stage 2. |
| `--trainable_modules` | `"attn_op"` = Stage A (train only the sparse branch); `"."` = Stage B (full sparse finetune, load Stage A via `--transformer_path`). |
| `--max_train_steps` | Warm-up budget; keep it short. |
On completion the warm-up product is `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`.
### 3.3 Stage 2 — CrossDistill
Stage 2 distills the sparse student into a few-step, CFG-free generator against the dense teacher. `--sparse_student` turns the `generator` (and the `init_rollout` model) into `SparkWan2_2Transformer3DModel` while the teacher and critic stay dense. `--generator_transformer_path` loads the Stage-1 warm-up checkpoint. The host is a plain **discrete-grid DMD**: the few-step grid is derived in-script by splitting `--joint_backward_timesteps` at the config `boundary` (`0.875`) and keeping only this expert's segment (`--boundary_type`), then rolled out with rf-native Euler steps (`--flow_euler_rollout`), the generator and critic alternating via `--gen_update_interval`. On top of this host the **high** expert adds an optional CrossDistill PCM diversity term (`--loss_scale_div > 0`, see §3.3.1); the **low** expert is div-free. An opt-in `--tdm` switch (§3.3.2) replaces the plain-DMD student loss with the official SparkDiffusion **TDM** loss; it is not enabled in the published launchers, so the default run stays plain DMD.
Run the high-noise expert with [scripts/sparkdiffusion/train_distill_wan2.2_high.sh](./train_distill_wan2.2_high.sh):
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Stage-1 warm-up product.
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
--fsdp_offload_params true \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_n_frames=81 \
--fix_sample_size 480 832 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--enable_bucket \
--trainable_modules "." \
--train_mode="normal" \
--uniform_sampling \
--flow_euler_rollout \
--gen_update_interval 5 \
--joint_backward_timesteps 0.933781 0.875 0.608979 \
--div_mode pcm_random \
--loss_scale_div 10.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 12.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--boundary_type="high" \
--real_guidance_scale=5.0 \
--sparse_student \
--student_topk=0.05 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
```
Or simply:
```bash
bash scripts/sparkdiffusion/train_distill_wan2.2_high.sh
```
The low-noise expert ([scripts/sparkdiffusion/train_distill_wan2.2_low.sh](./train_distill_wan2.2_low.sh)) differs in three ways:
```bash
bash scripts/sparkdiffusion/train_distill_wan2.2_low.sh
```
| Difference | High expert | Low expert |
|------------|-------------|------------|
| `--boundary_type` | `high` | `low` |
| Diversity term | `--div_mode="pcm_random" --loss_scale_div=10.0` | div-free (no div flags) |
| `--init_rollout_path` | not used | warm-starts the rollout's boundary latent from the trained **high** student (`INIT_ROLLOUT_CKPT`) |
> ⚠️ The low-noise segment of the few-step trajectory starts at the MoE boundary, not at pure noise. Point `INIT_ROLLOUT_CKPT` at the latest high-distill checkpoint (`output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors`) so the frozen high student rolls the high-noise segment; leaving it empty makes the untrained low student roll that segment itself and the samples degrade.
> ⚠️ **The distillation phases have no step cap.** They run until you stop them manually (kill the `train_distill_wan2.2.py` processes); checkpoints are written every `--checkpointing_steps=200`.
#### 3.3.1 CrossDistill Diversity (PCM, high expert only)
On the **high**-noise expert, setting `--loss_scale_div > 0` layers the official SparkDiffusion **CrossDistill** PCM diversity term on top of the discrete-grid DMD host, to counter the diversity loss that pure distribution matching can cause. It is a self-contained additive loss that leaves the DMD rollout untouched (the **low** expert keeps `--loss_scale_div 0` and is div-free):
- **What it does** — a phase-0 boundary-alignment MSE: the student's one-step jump from a noisy anchor to the phase-0 boundary (`--div_rf_boundary`, the high segment's first node `0.933781`) must match a teacher K-step Euler rollout (`--dcm_skipping_interval_steps` sub-steps) followed by the student's own target jump. It runs on its **own continuous rf grid** (`--dcm_total_steps` subdivisions shifted by `--teacher_timestep_shift 12.0`, the Wan2.2 scheduler shift), decoupled from the discrete DMD grid.
- **Anchor** — the paired **real video latent** (faithful to the official `_compute_pcm_div_loss`, which distills on real `x0`), encoded on demand only when diversity is on.
- **Merge** — `generator_loss = dmd_loss + loss_scale_div * div_loss` in a **single** `accelerator.backward` (FSDP allows one backward per wrapped module per iteration); the teacher rollout and the student target run under `no_grad` before the graded student forward.
- **Scale & tuning** — `--loss_scale_div` is in **per-element-mean** space (aligned with the mean-reduced `dmd_loss`); the launchers start at `10.0` on the high expert. On the first run watch `train_div_loss` vs `train_dmd_loss` and adjust so diversity does not overwhelm DMD; `--div_teacher_guidance` is the diversity teacher's CFG scale, independent of `--real_guidance_scale`. Set `--loss_scale_div 0` to recover the exact plain-DMD baseline.
#### 3.3.2 TDM Loss (opt-in)
An opt-in `--tdm` switch replaces the plain-DMD student loss with **TDM** (Trajectory Distribution Matching, the official SparkDiffusion student loss); **it is not enabled in the published `train_distill_wan2.2_{high,low}.sh`**, so the default run stays plain DMD. Three deltas from DMD: (1) the last rollout segment does not collapse to `x0` but takes one more rf-Euler step to `s_{i-1}` (the next grid knot), so the anchor is `G_x_{t_{i-1}}`; (2) `D_time` is clamped strictly above `s_{i-1}`; (3) the fake/teacher `x0` preds are mapped back to `s_{i-1}` via the rf PF-ODE analytic step and the surrogate is compared in the `x_{t_{i-1}}` domain. When `s_{i-1}==0` (grid exhausted) it reduces **exactly** to DMD. Requires `--flow_euler_rollout`; mutually exclusive with `--dfd` / `--decoupled_dmd`; compatible with the diversity term on the high expert.
### 3.4 Common Training Parameters
| Parameter | Description | Example Value |
|-----------|-------------|---------------|
| `--config_path` | Wan2.2 config; carries the MoE `boundary` | `config/wan2.2/wan_civitai_t2v.yaml` |
| `--pretrained_model_name_or_path` | Wan2.2-T2V-A14B base (both experts + VAE + UMT5) | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` |
| `--train_data_dir` / `--train_data_meta` | Dataset root / metadata JSON | see §2 |
| `--boundary_type` | `high` / `low` MoE expert selector | `high` |
| `--train_mode` | SparkDiffusion supports `normal` (T2V) only | `normal` |
| `--student_topk` | Sparse attention keep-ratio (sparsity = 1 − topk); Stage 1 and Stage 2 must match | `0.05` |
| `--sparse_student` | (Stage 2) build generator + init_rollout as the sparse student | - |
| `--generator_transformer_path` | (Stage 2) Stage-1 warm-up checkpoint loaded into the generator | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` |
| `--uniform_sampling` | Sample the DMD noise level uniformly over the training schedule | - |
| `--flow_euler_rollout` | Roll the generator in flow/velocity space with deterministic rf-native Euler steps (no x0 re-noising) | - |
| `--tdm` | (opt-in) Use TDM — the official SparkDiffusion student loss — instead of plain DMD; reduces exactly to DMD when `s_{i-1}==0`. Requires `--flow_euler_rollout`, mutually exclusive with `--dfd` / `--decoupled_dmd` | - |
| `--gen_update_interval` | Update the generator once every N steps (the critic updates every step) | `5` |
| `--joint_backward_timesteps` | Joint few-step rf grid (descending), split at the config `boundary` into per-expert segments by `--boundary_type`; passed explicitly in the launchers (equals the built-in default), and validation reuses it | `0.933781 0.875 0.608979` |
| `--div_mode` | CrossDistill diversity mode: `pcm_random` (random phase-0 offset) or `pcm_fixed` (high expert only) | `pcm_random` |
| `--loss_scale_div` | Weight of the CrossDistill PCM diversity term (per-element-mean space); `0` disables it | `10.0` (high) / `0` (low) |
| `--dcm_total_steps` | Subdivision of the diversity term's continuous rf grid (`dw = 1 / N`) | `48` |
| `--dcm_skipping_interval_steps` | Teacher Euler sub-steps `K` within phase 0 | `1` |
| `--teacher_timestep_shift` | rf timestep shift for the diversity grid (matches the Wan2.2 scheduler shift) | `12.0` |
| `--div_rf_boundary` | Phase-0 right boundary `rf_b` (the high segment's first node) | `0.933781` |
| `--div_teacher_guidance` | CFG scale of the diversity teacher rollout (independent of `--real_guidance_scale`) | `5.0` |
| `--real_guidance_scale` | CFG scale of the frozen dense teacher | `5.0` |
| `--init_rollout_path` | (low expert) trained high student used to roll the rollout's boundary start latent | high-distill ckpt |
| `--learning_rate` / `--learning_rate_critic` | Generator / critic learning rate | `2e-06` / `2e-06` |
| `--checkpointing_steps` | Save checkpoint every N steps | `200` |
| `--max_grad_norm` | Gradient clipping threshold | `0.05` |
| `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` |
> 💡 The per-expert few-step grid is **auto-derived**: `--joint_backward_timesteps` (the joint rf grid, passed explicitly in the launchers) is split at the config `boundary` by `--boundary_type`, and `--init_rollout_timesteps` (the low expert's frozen high-segment relay) is filled from the knots above the boundary. Leave `--init_rollout_timesteps` at its default; validation reuses the same grid, so `num_inference_steps` = `len(--joint_backward_timesteps) + 1` (`4` here).
### 3.5 FSDP Backend and Memory Notes
Stage 2 keeps **three** 14B transformers co-resident (sparse generator + dense teacher + dense critic; the low expert also parks a frozen `init_rollout` high-noise relay, sharded and `no_grad`), all optimizer-state dominated, so a sharding backend is mandatory — FSDP `FULL_SHARD` here. The distill launchers combine these memory levers:
- `--fsdp_sharding_strategy "FULL_SHARD"` — shard params/grads/optimizer across GPUs.
- `--fsdp_backward_prefetch "BACKWARD_POST"` — lower median residency during backward.
- `--gradient_checkpointing` — recompute activations in backward.
- **Resolution, not frame count, is the OOM lever.** Each step rolls the few-step grid and backprops a grad-carrying 14B forward, so memory scales with `seq_len`; the launchers use `--fix_sample_size 480 832` with `--video_sample_n_frames=81`. If it OOMs, lower `--fix_sample_size` (e.g. to `320 512`) rather than dropping frames.
- Frame count must satisfy `(N − 1) % 4 == 0`.
> 💡 The distill launchers already pass `--fsdp_offload_params true` (mandatory at 14B: it offloads params/grads/optimizer to CPU, at some speed cost); drop it only if you have ample VRAM. `--low_vram` additionally parks the frozen side (`real_score` / `init_rollout` / text encoder) on CPU; it complements, but does not replace, sharding.
### 3.6 Other Backends
The Stage-2 launchers default to FSDP `FULL_SHARD`. The same [`train_distill_wan2.2.py`](./train_distill_wan2.2.py) also runs under DeepSpeed Zero-2 or without a sharding backend. All examples below use the **high**-noise expert; for the low expert flip `--boundary_type="low"`, drop the diversity flags (`--div_mode` … `--div_teacher_guidance`), add `--init_rollout_path`, and change `--output_dir` / `SPARSE_STUDENT_CKPT` as in §3.3.
#### 3.6.1 Training with DeepSpeed Zero-2
DeepSpeed Zero-2 shards optimizer states and gradients but **not** parameters. Since Stage 2 keeps three 14B transformers co-resident, Zero-2 only fits at low resolution and may OOM where FSDP `FULL_SHARD` (which also shards parameters) does not — prefer FSDP for the 14B experts.
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Stage-1 warm-up product.
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/sparkdiffusion/train_distill_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_n_frames=81 \
--fix_sample_size 480 832 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--enable_bucket \
--trainable_modules "." \
--train_mode="normal" \
--uniform_sampling \
--flow_euler_rollout \
--gen_update_interval 5 \
--joint_backward_timesteps 0.933781 0.875 0.608979 \
--div_mode pcm_random \
--loss_scale_div 10.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 12.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--boundary_type="high" \
--real_guidance_scale=5.0 \
--sparse_student \
--student_topk=0.05 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
```
#### 3.6.2 Training Without DeepSpeed and FSDP
**This approach is not recommended: with three co-resident 14B transformers and no sharding backend it easily OOMs.** It is provided only as a reference shell.
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Stage-1 warm-up product.
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/sparkdiffusion/train_distill_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_n_frames=81 \
--fix_sample_size 480 832 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--enable_bucket \
--trainable_modules "." \
--train_mode="normal" \
--uniform_sampling \
--flow_euler_rollout \
--gen_update_interval 5 \
--joint_backward_timesteps 0.933781 0.875 0.608979 \
--div_mode pcm_random \
--loss_scale_div 10.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 12.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--boundary_type="high" \
--real_guidance_scale=5.0 \
--sparse_student \
--student_topk=0.05 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
```
### 3.7 Multi-Node Distributed Training
The launchers are single-node. For multiple machines, add the following to the `accelerate launch` line (all machines must share the same dataset / base weights / `output_dir` over NFS/shared storage):
```bash
--main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT \
--num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK
```
Without RDMA, also set `NCCL_IB_DISABLE=1` and `NCCL_P2P_DISABLE=1`.
---
## 4. Inference — Sparse Few-Step Student
Acceptance is a standalone run of [examples/sparkdiffusion/predict_t2v.py](../../examples/sparkdiffusion/predict_t2v.py), which rebuilds each MoE expert as a sparse `SparkWan2_2Transformer3DModel` and loads the distilled checkpoints. Edit the top of the script:
| Parameter | Description | Example Value |
|-----------|-------------|---------------|
| `GPU_memory_mode` | Memory mode | `model_cpu_offload` |
| `model_name` | Wan2.2-T2V-A14B base (VAE / UMT5 / schedulers) | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` |
| `student_topk` | **Must equal** the training `--student_topk` | `0.05` |
| `transformer_path` | **Low**-noise sparse student checkpoint | `output_dir_sparkdiffusion_wan2.2_14b_low_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
| `transformer_high_path` | **High**-noise sparse student checkpoint | `output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
| `sampler_name` | `Flow` matches the few-step CFG-free student | `Flow` |
| `num_inference_steps` | The student's few-step count | `4` |
| `guidance_scale` | Distilled student is CFG-free | `1.0` |
| `sample_size` / `video_length` / `fps` | `[height, width]` / frames / fps | `[240, 336]` / `81` / `16` |
```bash
python examples/sparkdiffusion/predict_t2v.py
```
> 💡 Per-expert loading order is `from_pretrained → set_student_topk → load_state_dict`; reversing it re-zeroes the trained RoLa weights. An expert whose path is `None` stays dense — for the full SparkDiffusion student, set **both** paths (prefer each expert's Stage-2 distill product).
> 📌 **Stage 3** of the paper (the FP8/INT8 fused quantized inference kernel) is intentionally not integrated in this repository.
---
## 5. Additional Resources
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
- **Paper**: SparkDiffusion, *Mitigating the High-Sparsity Trap* (arXiv:2609.23153)
- **Sparse student model**: `videox_fun/models/spark_transformer3d.py` (`SparkWan2_2Transformer3DModel`)
- **RoLa sparse attention kernel**: `videox_fun/models/attention_utils.py`
scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2_zh-CN.md
+495
View File
@@ -0,0 +1,495 @@
# SparkDiffusion 稀疏蒸馏训练指南
本文档给出 Wan2.2-T2V-A14B 上 **SparkDiffusion** 的完整流程：一条两阶段管线，先暖启一个**稀疏（RoLa）学生**，再把它**蒸馏**成少步、CFG-free 的生成器。内容涵盖环境配置、数据准备、四个启动脚本（Stage-1 暖启 + Stage-2 CrossDistill，按 MoE 专家分别执行）、以及 `predict_t2v.py` 推理验收。
> **什么是 SparkDiffusion？**
>
> SparkDiffusion（arXiv:2609.23153，《Mitigating the High-Sparsity Trap》）把**补偿式稀疏注意力**（RoLa）与 **CrossDistill** 少步配方结合起来。Wan2.2-T2V-A14B 是双专家 MoE（一个高噪专家、一个低噪专家，按 config 的 `boundary: 0.875` 切分）；每个专家都在**各自的噪声区间**上单独暖启与蒸馏，绝不联合训练。
>
> - **Stage 1 — 稀疏暖启**（[`train_finetune_sparse_wan2.2.py`](./train_finetune_sparse_wan2.2.py)）：把 RoLa 稀疏自注意力插入 dense 主干，用标准的 step-local flow-matching 目标短暂微调，使蒸馏从一个可用的粗先验出发，而不是从零初始化的稀疏分支出发。该阶段刻意保持**短**。
> - **Stage 2 — CrossDistill**（[`train_distill_wan2.2.py`](./train_distill_wan2.2.py)）：对着 **dense** teacher 把稀疏学生蒸馏成少步生成器。它是分布匹配蒸馏（DMD），含三个同驻网络——可训练的稀疏 `generator`、冻结的 dense `real_score` teacher、可训练的 dense `fake_score` critic。**高噪**专家额外加一项 PCM 式多样性损失；**低噪**专家则不含多样性项。
>
> | 阶段 | 高噪专家 | 低噪专家 | 产物 |
> |-------|-------------------|------------------|--------|
> | **1 暖启** | [`train_finetune_sparse_wan2.2_high.sh`](./train_finetune_sparse_wan2.2_high.sh) | [`train_finetune_sparse_wan2.2_low.sh`](./train_finetune_sparse_wan2.2_low.sh) | `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola` |
> | **2 蒸馏** | [`train_distill_wan2.2_high.sh`](./train_distill_wan2.2_high.sh) | [`train_distill_wan2.2_low.sh`](./train_distill_wan2.2_low.sh) | `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola_distill` |
>
> 两个专家分别蒸馏，推理时同时载入，共同构成完整的少步 MoE 学生。
---
## 目录
- [一、环境配置](#一环境配置)
- [二、数据准备](#二数据准备)
- [三、训练](#三训练)
- [3.1 下载预训练模型](#31-下载预训练模型)
- [3.2 Stage 1 — 稀疏暖启](#32-stage-1--稀疏暖启)
- [3.3 Stage 2 — CrossDistill](#33-stage-2--crossdistill)
- [3.4 通用训练参数](#34-通用训练参数)
- [3.5 FSDP 后端与显存说明](#35-fsdp-后端与显存说明)
- [3.6 其他后端](#36-其他后端)
- [3.7 多机分布式训练](#37-多机分布式训练)
- [四、推理 — 稀疏少步学生](#四推理--稀疏少步学生)
- [五、更多资源](#五更多资源)
---
## 一、环境配置
**方式一：使用 requirements.txt**
```bash
pip install -r requirements.txt
```
**方式二：手动安装**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 💡 RoLa 稀疏注意力依赖仓库内的 Triton block-sparse 内核（`videox_fun/models/attention_utils.py`），因此稀疏学生需要一个可用的 Triton 安装。
**方式三：使用 Docker**
```bash
# 拉取镜像
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# 进入镜像
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 二、数据准备
SparkDiffusion 使用普通的「视频-字幕」数据（与标准 Wan2.2 训练格式一致）。
### 2.1 快速测试数据集
```bash
# 下载官方示例数据集
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 metadata.json 格式
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
- `file_path`：视频路径（相对于 `DATASET_NAME`，或绝对路径）。
- `text`：字幕 / 提示词。
- `type`：固定为 `"video"`。
- `width` / `height`：建议提供以启用 bucket 训练；若缺失可用 `scripts/process_json_add_width_and_height.py` 补齐。
### 2.3 相对路径与绝对路径
```bash
# 相对路径
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# 绝对路径（数据在 NAS/OSS 上、跨机共享）
export DATASET_NAME=""
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
```
---
## 三、训练
完整顺序为 **Stage-1 high → Stage-2 high → Stage-1 low → Stage-2 low**。请手动按此顺序运行四个启动脚本。
### 3.1 下载预训练模型
```bash
# 创建模型目录
mkdir -p models/Diffusion_Transformer
# 下载 Wan2.2 T2V 官方权重
modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-T2V-A14B
```
基座权重会播种 Stage-2 的全部三个网络（稀疏 `generator` 还会额外用 Stage-1 暖启产物初始化）。其目录结构为 `high_noise_model/`、`low_noise_model/`、VAE 以及 UMT5 文本编码器；MoE 分界从 `config/wan2.2/wan_civitai_t2v.yaml`（`boundary: 0.875`）读取。
### 3.2 Stage 1 — 稀疏暖启
Stage 1 插入 RoLa 稀疏自注意力，并用 step-local flow-matching 目标短暂暖启。这是一个**短**阶段——短暂暖启即足以为蒸馏建立可用先验，因此启动脚本把它限制在 `--max_train_steps=2000`。**不要**用该阶段自身的采样质量来评判它；清晰度是在 Stage 2 建立的。
用 [scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh](./train_finetune_sparse_wan2.2_high.sh) 运行高噪专家：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--max_train_steps=2000 \
--checkpointing_steps=200 \
--learning_rate=1e-4 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--enable_bucket \
--uniform_sampling \
--trainable_modules "attn_op" \
--student_topk=0.05 \
--train_mode="normal" \
--boundary_type="high" \
--resume_from_checkpoint="latest"
```
或直接：
```bash
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh
```
低噪专家用同样的配方，只把 `--boundary_type` 改为 `"low"` 并指定各自的 `--output_dir`：
```bash
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.2_low.sh
```
Stage-1 关键参数：
| 旗标 | 含义 |
|------|---------|
| `--boundary_type` | `high` / `low`——选择要暖启的 MoE 专家（由 config 分界决定）。 |
| `--student_topk` | 稀疏注意力中每个 query 保留的 key block 比例（稀疏度 = `1 - topk`）。`0.05` ≈ 95% 稀疏。**必须与** Stage 2 一致。 |
| `--trainable_modules` | `"attn_op"` = Stage A（只训练稀疏分支）；`"."` = Stage B（稀疏网全量微调，用 `--transformer_path` 载入 Stage A）。 |
| `--max_train_steps` | 暖启预算；保持简短。 |
完成后暖启产物为 `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`。
### 3.3 Stage 2 — CrossDistill
Stage 2 对着 dense teacher 把稀疏学生蒸馏成少步、CFG-free 的生成器。`--sparse_student` 把 `generator`（以及 `init_rollout` 模型）变为 `SparkWan2_2Transformer3DModel`，而 teacher 与 critic 保持 dense。`--generator_transformer_path` 载入 Stage-1 暖启 checkpoint。宿主是一个朴素的**离散网格 DMD**：少步网格在脚本内由 `--joint_backward_timesteps` 按 config 的 `boundary`（`0.875`）切分、只保留本专家的段（`--boundary_type`）而派生，随后用 rf-native 的 Euler 步进滚动（`--flow_euler_rollout`），generator 与 critic 通过 `--gen_update_interval` 交替更新。在此宿主之上，**高噪**专家额外叠加一个可选的 CrossDistill PCM 多样性项（`--loss_scale_div > 0`，见 §3.3.1）；**低噪**专家则不含多样性项。一个可选的 `--tdm` 开关（§3.3.2）把 plain-DMD 学生损失替换为官方 SparkDiffusion **TDM** 损失；它未在发布态启动脚本中启用，故默认运行仍是 plain DMD。
用 [scripts/sparkdiffusion/train_distill_wan2.2_high.sh](./train_distill_wan2.2_high.sh) 运行高噪专家：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Stage-1 暖启产物。
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
--fsdp_offload_params true \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_n_frames=81 \
--fix_sample_size 480 832 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--enable_bucket \
--trainable_modules "." \
--train_mode="normal" \
--uniform_sampling \
--flow_euler_rollout \
--gen_update_interval 5 \
--joint_backward_timesteps 0.933781 0.875 0.608979 \
--div_mode pcm_random \
--loss_scale_div 10.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 12.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--boundary_type="high" \
--real_guidance_scale=5.0 \
--sparse_student \
--student_topk=0.05 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
```
或直接：
```bash
bash scripts/sparkdiffusion/train_distill_wan2.2_high.sh
```
低噪专家（[scripts/sparkdiffusion/train_distill_wan2.2_low.sh](./train_distill_wan2.2_low.sh)）有三处不同：
```bash
bash scripts/sparkdiffusion/train_distill_wan2.2_low.sh
```
| 差异项 | 高噪专家 | 低噪专家 |
|------------|-------------|------------|
| `--boundary_type` | `high` | `low` |
| 多样性项 | `--div_mode="pcm_random" --loss_scale_div=10.0` | 无多样性项（不加 div 旗标） |
| `--init_rollout_path` | 不使用 | 用训练好的**高噪**学生暖启 rollout 的分界起点 latent（`INIT_ROLLOUT_CKPT`） |
> ⚠️ 少步轨迹的低噪段起点在 MoE 分界处、而非纯噪声。请把 `INIT_ROLLOUT_CKPT` 指向最新的高噪蒸馏 checkpoint（`output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors`），让冻结的高噪学生去滚动高噪段；若留空，则由未训练的低噪学生自己滚动该段，样本会退化。
> ⚠️ **蒸馏阶段没有步数上限。** 它们会一直运行直到你手动停止（kill 掉 `train_distill_wan2.2.py` 进程）；checkpoint 每 `--checkpointing_steps=200` 步写一次。
#### 3.3.1 CrossDistill 多样性项（PCM，仅高噪专家）
在**高噪**专家上，设置 `--loss_scale_div > 0` 会在离散网格 DMD 宿主之上叠加官方 SparkDiffusion 的 **CrossDistill** PCM 多样性项，用以抵消纯分布匹配可能带来的多样性损失。它是一个自包含的加性损失，不改动 DMD rollout（**低噪**专家保持 `--loss_scale_div 0`，不含多样性项）：
- **它做什么**——一个 phase-0 分界对齐 MSE：学生从带噪锚点一步跳到 phase-0 分界（`--div_rf_boundary`，即高噪段的首个节点 `0.933781`），必须与 teacher 的 K 步 Euler rollout（`--dcm_skipping_interval_steps` 个子步）再接学生自己的目标跳变相匹配。它运行在**自己独立的连续 rf 网格**上（`--dcm_total_steps` 等分、并按 `--teacher_timestep_shift 12.0`——Wan2.2 调度器的 shift——偏移），与离散的 DMD 网格解耦。
- **锚点**——配对的**真实视频 latent**（忠于官方 `_compute_pcm_div_loss`，其在真实 `x0` 上蒸馏），仅在开启多样性项时按需编码。
- **合并**——`generator_loss = dmd_loss + loss_scale_div * div_loss`，在**单次** `accelerator.backward` 中完成（FSDP 每个被包裹的模块每轮迭代只允许一次 backward）；teacher rollout 与学生目标在带梯度的学生前向之前于 `no_grad` 下运行。
- **尺度与调参**——`--loss_scale_div` 处于**逐元素均值**空间（与均值归约的 `dmd_loss` 对齐）；启动脚本在高噪专家上从 `10.0` 起步。首次运行时观察 `train_div_loss` 与 `train_dmd_loss`，调整到多样性不会压过 DMD；`--div_teacher_guidance` 是多样性 teacher 的 CFG 尺度，独立于 `--real_guidance_scale`。设 `--loss_scale_div 0` 即可恢复到完全一致的朴素 DMD 基线。
#### 3.3.2 TDM 损失（可选）
一个可选的 `--tdm` 开关把 plain-DMD 学生损失替换为 **TDM**（Trajectory Distribution Matching，官方 SparkDiffusion 学生损失）；**发布态的 `train_distill_wan2.2_{high,low}.sh` 中未启用它**，故默认运行仍是 plain DMD。相对 DMD 有三处差异：(1) rollout 末段不塌缩到 `x0`，而是再走一步 rf-Euler 到 `s_{i-1}`（下一个网格节点），故锚点为 `G_x_{t_{i-1}}`；(2) `D_time` 被严格 clamp 在 `s_{i-1}` 之上；(3) fake/teacher 的 `x0` 预测经 rf PF-ODE 解析步映回 `s_{i-1}`，surrogate 在 `x_{t_{i-1}}` 域比较。当 `s_{i-1}==0`（网格耗尽）时**精确**退化为 DMD。需 `--flow_euler_rollout`；与 `--dfd` / `--decoupled_dmd` 互斥；在高噪专家上与多样性项兼容。
### 3.4 通用训练参数
| 参数 | 说明 | 示例值 |
|-----------|-------------|---------------|
| `--config_path` | Wan2.2 config；携带 MoE `boundary` | `config/wan2.2/wan_civitai_t2v.yaml` |
| `--pretrained_model_name_or_path` | Wan2.2-T2V-A14B 基座（两个专家 + VAE + UMT5） | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` |
| `--train_data_dir` / `--train_data_meta` | 数据集根目录 / metadata JSON | 见 §2 |
| `--boundary_type` | `high` / `low` MoE 专家选择器 | `high` |
| `--train_mode` | SparkDiffusion 仅支持 `normal`（T2V） | `normal` |
| `--student_topk` | 稀疏注意力保留比例（稀疏度 = 1 − topk）；Stage 1 与 Stage 2 必须一致 | `0.05` |
| `--sparse_student` | （Stage 2）把 generator + init_rollout 建成稀疏学生 | - |
| `--generator_transformer_path` | （Stage 2）载入 generator 的 Stage-1 暖启 checkpoint | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` |
| `--uniform_sampling` | 在训练调度上均匀采样 DMD 的噪声水平 | - |
| `--flow_euler_rollout` | 在 flow/velocity 空间用确定性的 rf-native Euler 步进滚动 generator（不做 x0 再加噪） | - |
| `--tdm` | （可选）用 TDM（官方 SparkDiffusion 学生损失）替代 plain DMD；当 `s_{i-1}==0` 时精确退化为 DMD。需 `--flow_euler_rollout`，与 `--dfd` / `--decoupled_dmd` 互斥 | - |
| `--gen_update_interval` | 每 N 步更新一次 generator（critic 每步都更新） | `5` |
| `--joint_backward_timesteps` | 联合少步 rf 网格（降序），按 config 的 `boundary` 由 `--boundary_type` 切分成各专家的段；在启动脚本中显式传入（与内置默认值一致），validation 复用它 | `0.933781 0.875 0.608979` |
| `--div_mode` | CrossDistill 多样性模式：`pcm_random`（随机 phase-0 偏移）或 `pcm_fixed`（仅高噪专家） | `pcm_random` |
| `--loss_scale_div` | CrossDistill PCM 多样性项的权重（逐元素均值空间）；`0` 表示关闭 | `10.0`（高噪）/ `0`（低噪） |
| `--dcm_total_steps` | 多样性项连续 rf 网格的等分数（`dw = 1 / N`） | `48` |
| `--dcm_skipping_interval_steps` | phase 0 内 teacher 的 Euler 子步数 `K` | `1` |
| `--teacher_timestep_shift` | 多样性网格的 rf 时间步 shift（与 Wan2.2 调度器的 shift 一致） | `12.0` |
| `--div_rf_boundary` | phase-0 右边界 `rf_b`（高噪段的首个节点） | `0.933781` |
| `--div_teacher_guidance` | 多样性 teacher rollout 的 CFG 尺度（独立于 `--real_guidance_scale`） | `5.0` |
| `--real_guidance_scale` | 冻结 dense teacher 的 CFG 尺度 | `5.0` |
| `--init_rollout_path` | （低噪专家）用于滚动 rollout 分界起点 latent 的、已训练好的高噪学生 | 高噪蒸馏 ckpt |
| `--learning_rate` / `--learning_rate_critic` | 生成器 / critic 学习率 | `2e-06` / `2e-06` |
| `--checkpointing_steps` | 每 N 步保存 checkpoint | `200` |
| `--max_grad_norm` | 梯度裁剪阈值 | `0.05` |
| `--resume_from_checkpoint` | 续训路径或 `"latest"` | `latest` |
> 💡 各专家的少步网格是**自动推导**的：`--joint_backward_timesteps`（联合 rf 网格，在启动脚本中显式传入）按 config 的 `boundary` 由 `--boundary_type` 切分，`--init_rollout_timesteps`（低噪专家冻结的高噪段中继）由分界以上的节点填充。`--init_rollout_timesteps` 保持默认值即可；validation 复用同一网格，故 `num_inference_steps` = `len(--joint_backward_timesteps) + 1`（此处为 `4`）。
### 3.5 FSDP Backend 与显存说明
Stage 2 需让**三个** 14B 网络同驻（稀疏 generator + dense teacher + dense critic；低噪专家还会停放一个冻结的 `init_rollout` 高噪中继，已分片且 `no_grad`），且都以优化器态为主，因此必须使用分片后端——此处为 FSDP `FULL_SHARD`。蒸馏启动脚本组合了下列显存杠杆：
- `--fsdp_sharding_strategy "FULL_SHARD"`——在多卡间分片 param/grad/optimizer。
- `--fsdp_backward_prefetch "BACKWARD_POST"`——降低 backward 期间的中位驻留。
- `--gradient_checkpointing`——在 backward 中重算激活。
- **OOM 的杠杆是分辨率、而非帧数。** 每一步都会滚动少步网格并对一段带梯度的 14B 前向做反传，因此显存随 `seq_len` 增长；启动脚本使用 `--fix_sample_size 480 832` 配 `--video_sample_n_frames=81`。若 OOM，请降低 `--fix_sample_size`（例如降到 `320 512`），而不是减少帧数。
- 帧数须满足 `(N − 1) % 4 == 0`。
> 💡 蒸馏启动脚本已经传入 `--fsdp_offload_params true`（14B 下为必需：它把 param/grad/optimizer offload 到 CPU，会牺牲一些速度）；仅在显存充裕时才可去掉。`--low_vram` 还会额外把冻结侧（`real_score` / `init_rollout` / 文本编码器）停到 CPU；它是对分片的补充，而非替代。
### 3.6 其他后端
Stage-2 启动脚本默认使用 FSDP `FULL_SHARD`。同一个 [`train_distill_wan2.2.py`](./train_distill_wan2.2.py) 也可在 DeepSpeed Zero-2 下运行，或不使用任何分片后端。下面所有示例都以**高噪**专家为例；低噪专家请参照 §3.3 把 `--boundary_type` 改为 `"low"`、去掉多样性旗标（`--div_mode` …… `--div_teacher_guidance`）、加上 `--init_rollout_path`，并修改 `--output_dir` / `SPARSE_STUDENT_CKPT`。
#### 3.6.1 使用 DeepSpeed-Zero-2 进行训练
DeepSpeed Zero-2 只分片优化器态与梯度、**不**分片参数。由于 Stage 2 需让三个 14B 网络同驻，Zero-2 仅在低分辨率下放得下，且可能在 FSDP `FULL_SHARD`（它还会分片参数）不会 OOM 的地方 OOM——14B 专家请优先选用 FSDP。
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Stage-1 暖启产物。
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/sparkdiffusion/train_distill_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_n_frames=81 \
--fix_sample_size 480 832 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--enable_bucket \
--trainable_modules "." \
--train_mode="normal" \
--uniform_sampling \
--flow_euler_rollout \
--gen_update_interval 5 \
--joint_backward_timesteps 0.933781 0.875 0.608979 \
--div_mode pcm_random \
--loss_scale_div 10.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 12.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--boundary_type="high" \
--real_guidance_scale=5.0 \
--sparse_student \
--student_topk=0.05 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
```
#### 3.6.2 不使用 DeepSpeed 与 FSDP 训练
**该方案并不被推荐：三个 14B 网络同驻且没有分片后端，很容易造成显存不足**。这里仅提供训练 Shell 用于参考。
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Stage-1 暖启产物。
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/sparkdiffusion/train_distill_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_n_frames=81 \
--fix_sample_size 480 832 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--enable_bucket \
--trainable_modules "." \
--train_mode="normal" \
--uniform_sampling \
--flow_euler_rollout \
--gen_update_interval 5 \
--joint_backward_timesteps 0.933781 0.875 0.608979 \
--div_mode pcm_random \
--loss_scale_div 10.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 12.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--boundary_type="high" \
--real_guidance_scale=5.0 \
--sparse_student \
--student_topk=0.05 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
```
### 3.7 多机分布式训练
启动脚本为单机。多机时在 `accelerate launch` 行追加下列参数（所有机器须通过 NFS/共享存储共享同一份数据集 / 基座权重 / `output_dir`）：
```bash
--main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT \
--num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK
```
无 RDMA 时，另设 `NCCL_IB_DISABLE=1` 与 `NCCL_P2P_DISABLE=1`。
---
## 四、推理 — 稀疏少步学生
验收方式为独立运行 [examples/sparkdiffusion/predict_t2v.py](../../examples/sparkdiffusion/predict_t2v.py)，它会把每个 MoE 专家重建为稀疏 `SparkWan2_2Transformer3DModel` 并载入蒸馏 checkpoint。编辑脚本顶部：
| 参数 | 说明 | 示例值 |
|-----------|-------------|---------------|
| `GPU_memory_mode` | 显存模式 | `model_cpu_offload` |
| `model_name` | Wan2.2-T2V-A14B 基座（VAE / UMT5 / 调度器） | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` |
| `student_topk` | **必须等于**训练时的 `--student_topk` | `0.05` |
| `transformer_path` | **低噪**稀疏学生 checkpoint | `output_dir_sparkdiffusion_wan2.2_14b_low_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
| `transformer_high_path` | **高噪**稀疏学生 checkpoint | `output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
| `sampler_name` | `Flow` 与少步 CFG-free 学生匹配 | `Flow` |
| `num_inference_steps` | 学生的少步步数 | `4` |
| `guidance_scale` | 蒸馏学生为 CFG-free | `1.0` |
| `sample_size` / `video_length` / `fps` | `[高, 宽]` / 帧数 / 帧率 | `[240, 336]` / `81` / `16` |
```bash
python examples/sparkdiffusion/predict_t2v.py
```
> 💡 每个专家的载入顺序为 `from_pretrained → set_student_topk → load_state_dict`；顺序颠倒会把训练好的 RoLa 权重重新清零。路径为 `None` 的专家保持 dense——要得到完整的 SparkDiffusion 学生，请设置**两个**路径（优先用各专家的 Stage-2 蒸馏产物）。
> 📌 论文的 **Stage 3**（FP8/INT8 融合量化推理核）在本仓库中有意未接入。
---
## 五、更多资源
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
- **论文**：SparkDiffusion，《Mitigating the High-Sparsity Trap》（arXiv:2609.23153）
- **稀疏学生模型**：`videox_fun/models/spark_transformer3d.py`（`SparkWan2_2Transformer3DModel`）
- **RoLa 稀疏注意力内核**：`videox_fun/models/attention_utils.py`
scripts/sparkdiffusion/train_distill_wan2.1.py
+3008
View File
File diff suppressed because it is too large Load Diff
scripts/sparkdiffusion/train_distill_wan2.1.sh
+77
View File
@@ -0,0 +1,77 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Stage-1 RoLa sparse warm-up product (train_finetune_sparse_wan2.1.sh -> _1.3b_sparse_rola/checkpoint-400).
# --student_topk below MUST match the Stage-1 value (0.10, i.e. sparsity 0.9).
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
# SparkDiffusion Stage-2 (CrossDistill) for the Wan2.1 single-expert T2V model: distill the sparse student into a
# few-step, CFG-free generator against the dense teacher. --sparse_student turns the generator into
# SparkWanTransformer3DModel while teacher + critic stay dense; --generator_transformer_path loads the Stage-1
# sparse warm-up ckpt (no prefix stripping, unexpected keys must be 0). Wan2.1-T2V-1.3B is a single full-range
# expert, so the student distills over the whole rf-time range [0, 1] (no MoE boundary, no init-rollout relay).
# The host is a plain discrete-grid DMD (--denoising_step_indices_list 1000 750 500 250 + --flow_euler_rollout, mean
# reduction). CrossDistill layers the official PCM diversity term on top as a self-contained additive loss
# (--div_mode pcm_random --loss_scale_div 5.0): a phase-0 boundary-alignment MSE anchored on the paired real video
# latent, rolled on its own continuous rf grid (--dcm_total_steps 48 --dcm_skipping_interval_steps 1,
# --teacher_timestep_shift 5.0 = the Wan2.1 scheduler shift, --div_rf_boundary 0.933781 = phase-0 right edge,
# --div_teacher_guidance 5.0). --loss_scale_div is in per-element-mean space (the official uses 10.0 in sum space);
# set it to 0 to recover the plain-DMD baseline. Memory: --gradient_checkpointing keeps the three nets
# (student/teacher/critic) off the activation OOM path (1.3B single expert needs no CPU offload; add
# --fsdp_offload_params true only if it OOMs). No --max_train_steps: runs until you kill the procs.
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--flow_euler_rollout \
--randomize_step_indices \
--gen_update_interval 5 \
--div_mode pcm_random \
--loss_scale_div 5.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 5.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--trainable_modules "." \
--train_mode="normal" \
--sparse_student \
--student_topk=0.10 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
scripts/sparkdiffusion/train_distill_wan2.2.py
+3137
View File
File diff suppressed because it is too large Load Diff
scripts/sparkdiffusion/train_distill_wan2.2_high.sh
+80
View File
@@ -0,0 +1,80 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
# SparkDiffusion Stage-2 (CrossDistill) for the HIGH-noise expert of Wan2.2-T2V-A14B (an MoE T2V model):
# distill the sparse student into a few-step, CFG-free generator against the dense teacher. --sparse_student
# turns the generator (and the init-rollout model) into SparkWan2_2Transformer3DModel while teacher + critic
# stay dense; --generator_transformer_path loads the Stage-1 sparse warm-up ckpt (no prefix stripping,
# unexpected keys must be 0). --boundary_type="high" restricts the student to the high-noise rf segment: the
# few-step grid is derived IN-SCRIPT by splitting --joint_backward_timesteps at the config boundary (0.875) and
# keeping the knots >= boundary, so the high expert rolls [1.0 -> 0.933781 -> 0.875] with rf-native Euler steps
# (--flow_euler_rollout) -- no contiguous backward simulation. The host is a plain discrete-grid DMD
# (--uniform_sampling over that grid, mean reduction, generator/critic alternating via --gen_update_interval 5).
# CrossDistill layers the official PCM diversity term on the HIGH expert only (--div_mode pcm_random
# --loss_scale_div 10.0): a phase-0 boundary-alignment MSE anchored on the paired real-video latent, rolled on
# its own continuous rf grid (--dcm_total_steps 48 --dcm_skipping_interval_steps 1, --teacher_timestep_shift 12.0
# = the Wan2.2 scheduler shift, --div_rf_boundary 0.933781 = the high-segment first node, --div_teacher_guidance
# 5.0). Set --loss_scale_div 0 to recover the plain-DMD baseline. Memory: --fsdp_offload_params true +
# --gradient_checkpointing are mandatory (three 14B nets are optimizer-state dominated); resolution
# (--fix_sample_size), not frame count, is the OOM lever. No --max_train_steps: runs until you kill the procs.
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
--fsdp_offload_params true \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=1 \
--video_sample_n_frames=81 \
--fix_sample_size 480 832 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--flow_euler_rollout \
--gen_update_interval 5 \
--joint_backward_timesteps 0.933781 0.875 0.608979 \
--div_mode pcm_random \
--loss_scale_div 10.0 \
--dcm_total_steps 48 \
--dcm_skipping_interval_steps 1 \
--teacher_timestep_shift 12.0 \
--div_rf_boundary 0.933781 \
--div_teacher_guidance 5.0 \
--boundary_type="high" \
--real_guidance_scale=5.0 \
--trainable_modules "." \
--train_mode="normal" \
--sparse_student \
--student_topk=0.05 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
scripts/sparkdiffusion/train_distill_wan2.2_low.sh
+77
View File
@@ -0,0 +1,77 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Stage-1 RoLa sparse warm-up product for the LOW expert (train_finetune_sparse_wan2.2_low.sh ->
# _low_sparse_rola/checkpoint-400).
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_low_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
# Warm-start of the low expert's rollout from the trained HIGH distilled student
# (train_distill_wan2.2_high.sh -> _high_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors).
# Empty = cold start from pure noise; set it to the latest HIGH distill checkpoint for a proper relay.
# --init_rollout_path is fault-tolerant (auto-finds the safetensors and strips generator/transformer/model
# prefixes, so a distill checkpoint dir also works).
export INIT_ROLLOUT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill/checkpoint-4000/diffusion_pytorch_model.safetensors"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
# SparkDiffusion Stage-2 (CrossDistill) for the LOW-noise expert of Wan2.2-T2V-A14B: the same sparse-student
# discrete-grid DMD as train_distill_wan2.2_high.sh, with two low-specific differences. (1) --boundary_type="low"
# keeps only the --joint_backward_timesteps knots BELOW the config boundary (0.875), so the trainable low student
# rolls [0.875 -> 0.608979] with rf-native Euler steps (--flow_euler_rollout). Because that segment starts at the
# MoE boundary rather than pure noise, --init_rollout_path warm-starts it: a FROZEN copy of the trained HIGH
# student (INIT_ROLLOUT_CKPT) first rolls the high segment [1.0 -> 0.933781 -> 0.875] to produce the boundary
# latent, then the low student takes over. Leaving INIT_ROLLOUT_CKPT empty = cold start (the untrained low student
# rolls the high segment itself and the samples degrade). (2) The low expert is div-free (no --div_mode /
# --loss_scale_div): the PCM diversity term is gated on --loss_scale_div > 0 and lives on the HIGH expert only.
# Memory levers are identical to train_distill_wan2.2_high.sh (--fsdp_offload_params true + --gradient_checkpointing;
# lower --fix_sample_size, not frame count, on OOM). No --max_train_steps: runs until you kill the procs.
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
--fsdp_offload_params true \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=1 \
--video_sample_n_frames=81 \
--fix_sample_size 480 832 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_low_sparse_rola_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--flow_euler_rollout \
--gen_update_interval 5 \
--joint_backward_timesteps 0.933781 0.875 0.608979 \
--boundary_type="low" \
--real_guidance_scale=5.0 \
--init_rollout_path=$INIT_ROLLOUT_CKPT \
--trainable_modules "." \
--train_mode="normal" \
--sparse_student \
--student_topk=0.05 \
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
--resume_from_checkpoint="latest"
scripts/sparkdiffusion/train_finetune_sparse_wan2.1.py
+1969
View File
File diff suppressed because it is too large Load Diff
scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh
+55
View File
@@ -0,0 +1,55 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
# SparkDiffusion Stage-1 (sparse warm-up) for the Wan2.1 single-expert T2V model: insert the compensating
# sparse self-attention (RoLa: blocks.*.self_attn.attn_op.{proj_q,proj_k,gate_proj,gate_bias}) into the dense
# Wan2.1 backbone and briefly warm it up with the plain step-local flow-matching objective, so Stage-2
# distillation starts from a usable coarse prior instead of a zero-initialized sparse branch. Wan2.1-T2V-1.3B is
# a single full-range expert (no MoE boundary), so the student warms up over the whole rf-time range [0, 1].
# --student_topk=0.10 (~90% sparsity) must match the value used later in train_distill_wan2.1.sh. Stage A trains
# only "attn_op"; for Stage B (full sparse finetune) set --trainable_modules "." plus --transformer_path=<stage-A
# ckpt> and a new --output_dir.
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--max_train_steps=2000 \
--checkpointing_steps=200 \
--learning_rate=1e-4 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=1.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--trainable_modules "attn_op" \
--student_topk=0.10 \
--train_mode="normal" \
--resume_from_checkpoint="latest"
scripts/sparkdiffusion/train_finetune_sparse_wan2.2.py
+2017
View File
File diff suppressed because it is too large Load Diff
scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh
+57
View File
@@ -0,0 +1,57 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
# SparkDiffusion Stage-1 (sparse warm-up) for the HIGH-noise expert: insert the compensating sparse
# self-attention (RoLa: blocks.*.self_attn.attn_op.{proj_q,proj_k,gate_proj,gate_bias}) into the dense
# Wan2.2 backbone and briefly warm it up with the plain step-local flow-matching objective, so Stage-2
# distillation starts from a usable coarse prior instead of a zero-initialized sparse branch.
# --boundary_type="high" picks the high-noise expert from the config boundary (config/wan2.2/wan_civitai_t2v.yaml,
# boundary:0.875); the low expert is warmed up separately by train_finetune_sparse_wan2.2_low.sh (Sec 4.3: each MoE
# expert group is warmed up and distilled on its own noise range). --student_topk=0.05 (~95% sparsity) must
# match the value used later in train_distill_wan2.2_high.sh. Stage A trains only "attn_op"; for Stage B (full sparse
# finetune) set --trainable_modules "." plus --transformer_path=<stage-A ckpt> and a new --output_dir.
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--max_train_steps=2000 \
--checkpointing_steps=200 \
--learning_rate=1e-4 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=1.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--trainable_modules "attn_op" \
--student_topk=0.05 \
--train_mode="normal" \
--boundary_type="high" \
--resume_from_checkpoint="latest"
scripts/sparkdiffusion/train_finetune_sparse_wan2.2_low.sh
+53
View File
@@ -0,0 +1,53 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
# SparkDiffusion Stage-1 (sparse warm-up) for the LOW-noise expert: same recipe as
# train_finetune_sparse_wan2.2_high.sh, only --boundary_type flips to "low" so the low-noise expert is picked from
# the config boundary (config/wan2.2/wan_civitai_t2v.yaml, boundary:0.875) and warmed up on its own
# noise range (Sec 4.3: each MoE expert group is warmed up and distilled separately). --student_topk=0.05
# must match train_distill_wan2.2_low.sh. Its checkpoint feeds train_distill_wan2.2_low.sh as the low sparse student.
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.2.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--max_train_steps=2000 \
--checkpointing_steps=200 \
--learning_rate=1e-4 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_low_sparse_rola" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=1.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--trainable_modules "attn_op" \
--student_topk=0.05 \
--train_mode="normal" \
--boundary_type="low" \
--resume_from_checkpoint="latest"
scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md
+2
View File
@@ -172,6 +172,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -211,6 +212,7 @@ bash scripts/wan2.1_causal_forcing/train_ar_diffusion.sh
| `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_ar_diffusion` |
| `--gradient_checkpointing` | Enable activation checkpointing | - |
| `--mixed_precision` | `fp16` / `bf16` | `bf16` |
| `--fsdp_reduce_fp32` | FP32 gradient reduction under FSDP mixed precision; forward and parameter gathering stay low precision. CLI default is off; enabled in the launcher | enabled |
| `--adam_weight_decay` | AdamW weight decay | 3e-2 |
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md
+2
View File
@@ -172,6 +172,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -211,6 +212,7 @@ bash scripts/wan2.1_causal_forcing/train_ar_diffusion.sh
| `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_ar_diffusion` |
| `--gradient_checkpointing` | 启用激活重计算 | - |
| `--mixed_precision` | `fp16` / `bf16` | `bf16` |
| `--fsdp_reduce_fp32` | FSDP 混合精度下使用 FP32 梯度归约，前向与参数聚合仍保持低精度。CLI 默认关闭，launcher 已启用 | 启用 |
| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md
+2
View File
@@ -183,6 +183,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_ccd" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
@@ -228,6 +229,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh
| `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_ccd` |
| `--gradient_checkpointing` | Enable activation checkpointing | - |
| `--mixed_precision` | `fp16` / `bf16` | `bf16` |
| `--fsdp_reduce_fp32` | FP32 gradient reduction under FSDP mixed precision; forward and parameter gathering stay low precision. CLI default is off; enabled in the launcher | enabled |
| `--adam_weight_decay` | AdamW weight decay | 0.0 |
| `--adam_beta1` | AdamW beta1 (CCD uses 0.0) | 0.0 |
| `--adam_beta2` | AdamW beta2 | 0.999 |
scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md
+2
View File
@@ -183,6 +183,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_ccd" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
@@ -228,6 +229,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh
| `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_ccd` |
| `--gradient_checkpointing` | 启用激活重计算 | - |
| `--mixed_precision` | `fp16` / `bf16` | `bf16` |
| `--fsdp_reduce_fp32` | FSDP 混合精度下使用 FP32 梯度归约，前向与参数聚合仍保持低精度。CLI 默认关闭，launcher 已启用 | 启用 |
| `--adam_weight_decay` | AdamW 权重衰减 | 0.0 |
| `--adam_beta1` | AdamW beta1（CCD 使用 0.0） | 0.0 |
| `--adam_beta2` | AdamW beta2 | 0.999 |
scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md
+45 -5
View File
@@ -26,6 +26,7 @@ This document provides the complete workflow for **Causal-Forcing Stage 3 — Di
- [5.1 Quick Start](#51-quick-start)
- [5.2 Key Parameters](#52-key-parameters)
- [5.3 DMD-Specific Parameters](#53-dmd-specific-parameters)
- [5.4 Optional One-Forcing GAN and Precision Controls](#54-optional-one-forcing-gan-and-precision-controls)
- [6. Use the Trained Checkpoint](#6-use-the-trained-checkpoint)
- [7. Additional Resources](#7-additional-resources)
@@ -103,7 +104,7 @@ The Stage 2 CCD checkpoint (generator/critic init) is loaded via `--ode_transfor
## 4. Prepare Training Data
DMD uses a **TextDataset** (`train_mode="normal"`) — it only needs prompts, not video data, because the generator creates its own training samples via autoregressive rollout. However, a `metadata.json` with prompts is still required.
With `train_mode="normal"`, teacher forcing disabled, and all GAN/R1/R2 weights set to zero, DMD uses a **TextDataset**: only prompts are required because the generator creates training samples via autoregressive rollout. Enabling any GAN/R1/R2 weight switches to **ImageVideoDataset** and requires captioned real videos; see Section 5.4.
### 4.1 Quick Demo Dataset
@@ -118,7 +119,7 @@ Same as Stage 1/2. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](./README_TRA
### 4.3 metadata.json Format
DMD only uses the `"text"` field from each entry. The `file_path` and other fields are ignored when `--train_mode="normal"` (TextDataset mode).
In prompt-only TextDataset mode, DMD only uses the `"text"` field from each entry. GAN/R1/R2 or teacher forcing requires video metadata with `file_path`, `type="video"`, and `text`, as in Stage 1/2.
```json
[
@@ -131,7 +132,7 @@ DMD only uses the `"text"` field from each entry. The `file_path` and other fiel
]
```
> **Note**: You can reuse any video metadata.json — only the `text` field matters for DMD.
> **Note**: Existing video metadata can be reused. Only `text` is read in prompt-only mode; GAN training also reads the video files.
---
@@ -179,6 +180,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_dmd" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
@@ -212,7 +214,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--real_score_pretrained_model_name_or_path` | 14B real-score teacher for DMD | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
| `--config_path` | Model config YAML | `config/wan2.1/wan_civitai.yaml` |
| `--train_data_dir` | Data root directory | `""` |
| `--train_data_meta` | Path to `metadata.json` (prompts only) | `datasets/X-Fun-Videos-Demo/metadata.json` |
| `--train_data_meta` | Metadata path: prompts for plain DMD, captioned videos for GAN/R1/R2 | `datasets/X-Fun-Videos-Demo/metadata.json` |
| `--ode_transformer_path` | Stage 2 CCD checkpoint for generator/critic init | `$STAGE2_CKPT` |
| `--train_batch_size` | Per-GPU batch size | 1 |
| `--gradient_accumulation_steps` | Gradient accumulation steps | 1 |
@@ -264,11 +266,48 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--context_noise` | Context noise level for KV cache update | 0 |
| `--use_teacher_forcing` | Enable teacher forcing (pass clean_x to transformer) | - |
| `--teacher_forcing_prob` | Probability of applying teacher forcing per step (1.0 = always) | 1.0 |
| `--train_mode` | Training mode: `normal` (TextDataset, prompt-only) or `i2v` | `normal` |
| `--train_mode` | `normal` (prompt-only unless teacher forcing or GAN/R1/R2 is enabled) or `i2v` | `normal` |
| `--resume_from_checkpoint` | Resume from checkpoint. Use `"latest"` to auto-select | `"latest"` |
---
### 5.4 Optional One-Forcing GAN and Precision Controls
These options reuse the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms in `scripts/wan2.1_self_forcing/train_distill.py`. All GAN/R1/R2 weights default to **0**, and both precision flags default to **off** in argparse. The AR diffusion, CCD, and DMD launchers enable `--fsdp_reduce_fp32`; GAN/R1/R2 and `--score_timestep_clamp` remain off. AR diffusion and CCD do not gain GAN losses.
The discriminator shares the **fake-score backbone**, with one learnable register token per selected feature layer. Its head is optimized by the critic optimizer, not the generator optimizer. Generator loss is `DMD + weighted GAN`; critic loss is `denoising + weighted GAN + R1 + R2`.
| Parameter | Description | Default |
|-----------|-------------|---------|
| `--gan_g_weight` | Non-saturating generator GAN loss weight | `0.0` |
| `--gan_d_weight` | Discriminator GAN loss weight | `0.0` |
| `--gan_feature_layers` | Fake-score block indices used by the head | `21 29` |
| `--gan_block_ffn_dim` | Head cross-attention FFN dimension | `2048` |
| `--gan_block_num_heads` | Head count; must divide fake-score hidden dimension | `12` |
| `--relativistic_discriminator` | Use paired real/fake logit differences | off |
| `--r1_weight` | APT-style finite-difference penalty on noisy real latents | `0.0` |
| `--r1_sigma` | R1 perturbation standard deviation | `0.01` |
| `--r2_weight` | Normalized finite-difference penalty on noisy generated latents | `0.0` |
| `--r2_sigma` | R2 perturbation standard deviation | `0.01` |
| `--fsdp_reduce_fp32` | FSDP FP32 gradient reduction; parameter gathering stays low precision | off (enabled in the launcher) |
| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off |
For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` for the generator and `softplus(-D_real) + softplus(D_fake)` for the discriminator. Relativistic losses instead use `softplus(-(D_fake - D_real))` and `softplus(-(D_real - D_fake))`.
R1 follows the source's APT form: `r1_weight * mean((D_real_perturbed - D_real)^2)`, **without dividing by sigma**. R2 follows `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`. These are finite-difference penalties, not input-gradient double-backward regularizers. Real/fake samples share timestep and noise; perturbations use independent noise.
To enable GAN, pass positive generator/discriminator weights and, optionally, R1/R2 weights to `train_causal_dmd.py`. Requirements:
- Captioned real videos, not a prompt-only JSON. Use the Stage 1/2 video metadata schema and enough frames for `score_num_frames` and the selected causal block geometry.
- `--gradient_accumulation_steps=1`; DeepSpeed is not supported by the GAN branch.
- Valid feature-layer indices for the selected fake-score model. The default `21 29` layers and `12` heads fit Wan2.1-T2V-1.3B; other backbones may need different settings.
- Enabling only `gan_g_weight` does not supply a discriminator classification-learning objective. Choose the discriminator weight deliberately; nonzero values are not a guaranteed quality improvement.
Trainable fake-score and head master weights stay **FP32**, while forward uses the configured autocast/FSDP precision. Logits, softplus losses, and finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient communication precision, not the forward dtype or model architecture. The clamp changes score-training timesteps, not the generator rollout or validation schedule.
Resume GAN checkpoints with the same GAN/head configuration. The head is saved with fake-score, including FP32 head tensors, and is recreated before loading. It is not part of the exported generator, so inference needs no discriminator or new GAN flags.
---
## 6. Use the Trained Checkpoint
The Stage 3 DMD checkpoint is the **final model** in the Causal-Forcing pipeline. Use it for inference:
@@ -300,4 +339,5 @@ python examples/wan2.1_causal_forcing/predict_t2v.py
## 7. Additional Resources
- **Causal-Forcing Paper**: https://github.com/thu-ml/Causal-Forcing
- **One-Forcing (optional GAN mechanism)**: https://github.com/Aurora-edu/One-Forcing
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md
+57 -17
View File
@@ -26,6 +26,7 @@
- [5.1 快速开始](#51-快速开始)
- [5.2 关键参数](#52-关键参数)
- [5.3 DMD 特有参数](#53-dmd-特有参数)
- [5.4 可选 One-Forcing GAN 与精度控制](#54-可选-one-forcing-gan-与精度控制)
- [六、使用训练好的 Checkpoint](#六使用训练好的-checkpoint)
- [七、更多资源](#七更多资源)
@@ -39,7 +40,7 @@ Stage 3 需要：
2. **Wan2.1-T2V-14B** 模型：作为 DMD real-score teacher。
```bash
# 示例：Stage 2 CCD 训练的 checkpoint
# Example: Stage 2 CCD checkpoint
export STAGE2_CKPT="output_dir_wan2.1_causal_forcing_ccd/checkpoint-5000/transformer/diffusion_pytorch_model.safetensors"
```
@@ -70,10 +71,10 @@ pip install deepspeed==0.17.0 numpy==1.26.4
**方式 3：使用 Docker**
```bash
# 拉取镜像
# Pull the image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# 进入容器
# Enter the container
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
@@ -87,13 +88,13 @@ DMD 需要**两个**预训练模型：
- **Wan2.1-T2V-14B**：非因果的 real-score teacher，用于计算真实分布得分。
```bash
# 创建模型目录
# Create the model directory
mkdir -p models/Diffusion_Transformer
# 下载 Wan2.1 T2V 1.3B（student 基础模型）
# Download the Wan2.1 T2V 1.3B student base model
modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
# 下载 Wan2.1 T2V 14B（DMD real-score teacher）
# Download the Wan2.1 T2V 14B real-score teacher
modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B
```
@@ -103,12 +104,12 @@ Stage 2 CCD checkpoint（生成器/判别器初始化）通过 `--ode_transforme
## 四、数据准备
DMD 使用 **TextDataset**（`train_mode="normal"`）— 只需要提示词，不需要视频数据，因为生成器通过自回归 rollout 自行创建训练样本。但仍需提供包含提示词的 `metadata.json`。
当 `train_mode="normal"`、未启用 teacher forcing 且 GAN/R1/R2 权重均为零时，DMD 使用 **TextDataset**：只需要提示词，生成器通过自回归 rollout 自行创建训练样本。任一 GAN/R1/R2 权重为正时会切换为 **ImageVideoDataset**，需要带描述的真实视频，详见 5.4。
### 4.1 快速测试数据集
```bash
# 下载官方示例数据集
# Download the official demo dataset
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
@@ -118,7 +119,7 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun
### 4.3 metadata.json 格式
DMD 仅使用每条记录的 `"text"` 字段。当 `--train_mode="normal"`（TextDataset 模式）时，`file_path` 和其他字段会被忽略。
仅提示词的 TextDataset 模式只读取每条记录的 `"text"` 字段。GAN/R1/R2 或 teacher forcing 需要与 Stage 1/2 相同的视频元数据，包括 `file_path`、`type="video"` 和 `text`。
```json
[
@@ -131,7 +132,7 @@ DMD 仅使用每条记录的 `"text"` 字段。当 `--train_mode="normal"`（Tex
]
```
> **说明**：你可以复用任何视频 metadata.json — DMD 只关心 `text` 字段。
> **说明**：可以复用现有视频元数据。仅提示词模式只读取 `text`，GAN 训练还会读取视频文件。
---
@@ -179,6 +180,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_dmd" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
@@ -212,7 +214,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--real_score_pretrained_model_name_or_path` | 14B real-score teacher | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
| `--config_path` | 模型配置 YAML | `config/wan2.1/wan_civitai.yaml` |
| `--train_data_dir` | 数据根目录 | `""` |
| `--train_data_meta` | `metadata.json` 路径（仅含提示词） | `datasets/X-Fun-Videos-Demo/metadata.json` |
| `--train_data_meta` | 元数据路径：普通 DMD 使用提示词，GAN/R1/R2 使用带描述的视频 | `datasets/X-Fun-Videos-Demo/metadata.json` |
| `--ode_transformer_path` | Stage 2 CCD checkpoint（生成器/判别器初始化） | `$STAGE2_CKPT` |
| `--train_batch_size` | 每 GPU batch 大小 | 1 |
| `--gradient_accumulation_steps` | 梯度累积步数 | 1 |
@@ -264,23 +266,60 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--context_noise` | KV 缓存更新的上下文噪声级别 | 0 |
| `--use_teacher_forcing` | 启用 teacher forcing（将 clean_x 传给 transformer） | - |
| `--teacher_forcing_prob` | 每步应用 teacher forcing 的概率（1.0 = 始终） | 1.0 |
| `--train_mode` | 训练模式：`normal`（TextDataset，仅提示词）或 `i2v` | `normal` |
| `--train_mode` | `normal`（未启用 teacher forcing 或 GAN/R1/R2 时仅提示词）或 `i2v` | `normal` |
| `--resume_from_checkpoint` | 恢复训练。使用 `"latest"` 自动选择最新 checkpoint | `"latest"` |
---
### 5.4 可选 One-Forcing GAN 与精度控制
这些选项复用 `scripts/wan2.1_self_forcing/train_distill.py` 中的 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 机制。GAN/R1/R2 权重默认均为 **0**，两个精度开关在 argparse 中默认均**关闭**。AR diffusion、CCD 和 DMD 的 launcher 均启用 `--fsdp_reduce_fp32`；GAN/R1/R2 和 `--score_timestep_clamp` 仍关闭。AR diffusion 和 CCD 不添加 GAN 损失。
判别器共享 **fake-score backbone**，每个选定特征层对应一个可学习 register token。判别器头由 critic optimizer 更新，不加入 generator optimizer。生成器损失为 `DMD + 加权 GAN`，critic 损失为 `denoising + 加权 GAN + R1 + R2`。
| 参数 | 说明 | 默认值 |
|------|------|--------|
| `--gan_g_weight` | 非饱和生成器 GAN 损失权重 | `0.0` |
| `--gan_d_weight` | 判别器 GAN 损失权重 | `0.0` |
| `--gan_feature_layers` | 判别器头使用的 fake-score block 索引 | `21 29` |
| `--gan_block_ffn_dim` | 判别器交叉注意力 FFN 维度 | `2048` |
| `--gan_block_num_heads` | 判别器头数，须整除 fake-score 隐藏维度 | `12` |
| `--relativistic_discriminator` | 使用配对真实/生成样本 logit 差值 | 关闭 |
| `--r1_weight` | 噪声真实 latent 上的 APT 形式有限差分正则权重 | `0.0` |
| `--r1_sigma` | R1 扰动标准差 | `0.01` |
| `--r2_weight` | 噪声生成 latent 上的归一化有限差分正则权重 | `0.0` |
| `--r2_sigma` | R2 扰动标准差 | `0.01` |
| `--fsdp_reduce_fp32` | FSDP 使用 FP32 梯度归约，参数仍以低精度聚合 | 关闭（launcher 已启用） |
| `--score_timestep_clamp` | 将 DMD、生成器 GAN 和 critic 时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关闭 |
设 logits 为 `D_fake` 和 `D_real`，普通 GAN 的生成器损失为 `softplus(-D_fake)`，判别器损失为 `softplus(-D_real) + softplus(D_fake)`。相对论模式分别使用 `softplus(-(D_fake - D_real))` 和 `softplus(-(D_real - D_fake))`。
R1 保持源实现的 APT 形式：`r1_weight * mean((D_real_perturbed - D_real)^2)`，**不除以 sigma**。R2 为 `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`。二者都是有限差分正则，不是需要 double backward 的输入梯度惩罚。真实/生成样本共用时间步和噪声，扰动噪声独立抽样。
开启 GAN 时，向 `train_causal_dmd.py` 传入正的生成器/判别器权重，并按需配置 R1/R2。要求：
- 使用带描述的真实视频，不能只提供提示词 JSON。元数据采用 Stage 1/2 视频格式，帧数须满足 `score_num_frames` 与因果块几何。
- `--gradient_accumulation_steps=1`；GAN 分支不支持 DeepSpeed。
- 特征层索引必须存在。默认 `21 29` 与 `12` 个头适用于 Wan2.1-T2V-1.3B，其他 backbone 可能需要调整。
- 仅开启 `gan_g_weight` 不会提供判别器的分类学习目标。须明确选择判别器权重，非零权重不保证质量提升。
可训练 fake-score 与判别器头的 master weights 保持 **FP32**，前向按配置使用 autocast/FSDP 混合精度。Logits、softplus 损失及有限差分归约使用 FP32。`--fsdp_reduce_fp32` 只改变梯度通信精度，不改变前向 dtype 或模型结构。时间步 clamp 只作用于 score 训练，不改变生成器 rollout 或 validation schedule。
恢复 GAN checkpoint 时应使用相同的 GAN/判别器头配置。判别器头随 fake-score 保存，包括 FP32 判别器头张量，并在加载前重建。导出的 generator 不包含该头，推理无需判别器或新增 GAN 参数。
---
## 六、使用训练好的 Checkpoint
Stage 3 DMD checkpoint 是 Causal-Forcing 流水线的**最终模型**，直接用于推理：
```python
# 在 examples/wan2.1_causal_forcing/predict_t2v.py 中
# In examples/wan2.1_causal_forcing/predict_t2v.py
transformer_path = "output_dir_wan2.1_causal_forcing_dmd/checkpoint-{N}/diffusion_pytorch_model.safetensors"
# DMD Stage 3 推理配置
guidance_scale = 1.0 # CFG 已烘焦到蒸馏权重中
num_inference_steps = 4 # 4 步 DMD
# DMD Stage 3 inference config
guidance_scale = 1.0 # CFG is baked into distilled weights
num_inference_steps = 4 # 4-step DMD
stochastic_sampling = True
num_frame_per_block = 3 # 分块生成
num_frame_per_block = 3 # Chunk-wise generation
```
或运行：
@@ -299,4 +338,5 @@ python examples/wan2.1_causal_forcing/predict_t2v.py
## 七、更多资源
- **Causal-Forcing 论文**：https://github.com/thu-ml/Causal-Forcing
- **One-Forcing（可选 GAN 机制）**：https://github.com/Aurora-edu/One-Forcing
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
scripts/wan2.1_causal_forcing/train_ar_diffusion.py
+13
View File
@@ -18,6 +18,7 @@ scripts/wan2.1_self_forcing/train_distill.py and Causal-Forcing
# See the License for the specific language governing permissions and
import argparse
import dataclasses
import gc
import logging
import math
@@ -314,6 +315,10 @@ def parse_args():
parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true",
help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
)
parser.add_argument(
"--logging_dir",
type=str,
@@ -596,6 +601,9 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
@@ -1254,6 +1262,8 @@ def main():
training_weights_table = training_weights_table.to(accelerator.device)
for epoch in range(first_epoch, args.num_train_epochs):
if global_step >= args.max_train_steps:
break
train_loss = 0.0
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
for step, batch in enumerate(train_dataloader):
@@ -1487,6 +1497,9 @@ def main():
if global_step >= args.max_train_steps:
break
if global_step >= args.max_train_steps:
break
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
with progress_bar.paused():
log_validation(
scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py
+14 -1
View File
@@ -18,6 +18,7 @@ scripts/wan2.1_self_forcing/train_distill.py and Causal-Forcing
# See the License for the specific language governing permissions and
import argparse
import dataclasses
import gc
import logging
import math
@@ -317,6 +318,10 @@ def parse_args():
parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true",
help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
)
parser.add_argument(
"--logging_dir",
type=str,
@@ -610,6 +615,9 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
@@ -769,7 +777,7 @@ def main():
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
torch_dtype=torch.float32,
)
ema_transformer3d.num_frame_per_block = args.num_frame_per_block
ema_transformer3d.independent_first_frame = args.independent_first_frame
@@ -1437,6 +1445,8 @@ def main():
torch.cuda.empty_cache()
for epoch in range(first_epoch, args.num_train_epochs):
if global_step >= args.max_train_steps:
break
train_loss = 0.0
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
for step, batch in enumerate(train_dataloader):
@@ -1737,6 +1747,9 @@ def main():
if global_step >= args.max_train_steps:
break
if global_step >= args.max_train_steps:
break
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
with progress_bar.paused():
log_validation(
scripts/wan2.1_causal_forcing/train_causal_dmd.py
+392 -91
View File
@@ -17,6 +17,7 @@
import argparse
import contextlib
import dataclasses
import gc
import json
import logging
@@ -74,10 +75,12 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512,
get_random_mask)
from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
WanTransformer3DModel,
WanTransformer3DModel_SelfForcing)
WanTransformer3DModel_SelfForcing,
add_wan_gan_discriminator)
from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
WanSelfForcingPipeline)
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.fsdp_ema import FSDPEMA
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
@@ -158,6 +161,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi
return pred_for_score, score_num_frames, True
def repeat_batch(value, repeats):
"""Repeat a batched tensor or per-sample list along the batch dimension."""
if value is None or repeats == 1:
return value
if torch.is_tensor(value):
return torch.cat([value] * repeats)
return list(value) * repeats
def filter_kwargs(cls, kwargs):
import inspect
sig = inspect.signature(cls.__init__)
@@ -494,6 +506,14 @@ def parse_args():
),
)
parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true",
help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
)
parser.add_argument(
"--score_timestep_clamp", action="store_true",
help="Clamp DMD, GAN and critic timesteps to [0.02, 0.98] * num_train_timesteps.",
)
parser.add_argument(
"--dataloader_num_workers",
type=int,
@@ -838,6 +858,47 @@ def parse_args():
help="Probability of applying teacher forcing per step (1.0 = always)"
)
parser.add_argument(
"--gan_g_weight", type=float, default=0.0,
help="One-Forcing: non-saturating generator GAN loss weight (0 disables). Positive GAN/R1/R2 weights require real videos.",
)
parser.add_argument(
"--gan_d_weight", type=float, default=0.0,
help="One-Forcing: discriminator GAN loss weight added to critic denoising loss (0 disables).",
)
parser.add_argument(
"--gan_feature_layers", nargs="+", type=int, default=[21, 29],
help="fake_score block indices feeding the discriminator head, one register token per layer.",
)
parser.add_argument(
"--gan_block_ffn_dim", type=int, default=2048,
help="FFN dimension of the discriminator cross-attention blocks.",
)
parser.add_argument(
"--gan_block_num_heads", type=int, default=12,
help="Number of heads of the discriminator cross-attention blocks.",
)
parser.add_argument(
"--relativistic_discriminator", action="store_true",
help="Use relativistic GAN losses based on real/fake logit differences.",
)
parser.add_argument(
"--r1_weight", type=float, default=0.0,
help="Finite-difference R1 penalty weight on noisy real latents (APT form, not divided by sigma).",
)
parser.add_argument(
"--r1_sigma", type=float, default=0.01,
help="Perturbation standard deviation for the R1 penalty.",
)
parser.add_argument(
"--r2_weight", type=float, default=0.0,
help="Finite-difference R2 penalty weight on noisy generated latents (0.5 * squared logit difference / sigma squared)."
)
parser.add_argument(
"--r2_sigma", type=float, default=0.01,
help="Perturbation standard deviation for the R2 penalty.",
)
args = parser.parse_args()
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
if env_local_rank != -1 and env_local_rank != args.local_rank:
@@ -848,6 +909,14 @@ def parse_args():
def main():
args = parse_args()
use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight))
if any(not math.isfinite(w) or w < 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight)):
raise ValueError("GAN and regularization weights must be finite and non-negative.")
if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1):
raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.")
if use_gan and (args.r1_weight > 0 and (not math.isfinite(args.r1_sigma) or args.r1_sigma <= 0)
or args.r2_weight > 0 and (not math.isfinite(args.r2_sigma) or args.r2_sigma <= 0)):
raise ValueError("Enabled R1/R2 penalties require finite, positive perturbation standard deviations.")
if args.report_to == "wandb" and args.hub_token is not None:
raise ValueError(
@@ -875,6 +944,8 @@ def main():
deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None
fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None
if use_gan and deepspeed_plugin is not None:
raise ValueError("The GAN branch does not support DeepSpeed.")
if deepspeed_plugin is not None:
zero_stage = int(deepspeed_plugin.zero_stage)
fsdp_stage = 0
@@ -896,6 +967,9 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
@@ -1111,28 +1185,55 @@ def main():
param.requires_grad = True
break
if use_gan:
if (not args.gan_feature_layers or any(i < 0 or i >= len(fake_score_transformer3d.blocks) for i in args.gan_feature_layers)
or args.gan_block_ffn_dim <= 0 or args.gan_block_num_heads <= 0
or fake_score_transformer3d.dim % args.gan_block_num_heads != 0):
raise ValueError("GAN feature layers must exist and the positive head count must divide the fake_score hidden dimension.")
# Keep trainable discriminator master weights in fp32; forward uses autocast.
add_wan_gan_discriminator(
fake_score_transformer3d,
args.gan_feature_layers,
args.gan_block_ffn_dim,
args.gan_block_num_heads,
torch.float32,
seed=args.seed or 0,
)
# Create EMA for the generator_transformer3d.
if args.use_ema:
if zero_stage == 3:
raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.")
ema_generator_transformer3d = EMAModel(
generator_transformer3d.parameters(),
model_cls=WanTransformer3DModel_SelfForcing,
model_config=generator_transformer3d.config,
)
if fsdp_plugin is not None:
if getattr(fsdp_plugin, "fsdp_version", 1) != 1:
raise NotImplementedError("FSDPEMA 仅支持 FSDP1。")
ema_module = WanTransformer3DModel_SelfForcing.from_config(generator_transformer3d.config)
ema_generator_transformer3d = FSDPEMA(
ema_module, source=generator_transformer3d,
accelerator=accelerator, fsdp_plugin=fsdp_plugin,
)
else:
ema_generator_transformer3d = EMAModel(
generator_transformer3d.parameters(),
model_cls=WanTransformer3DModel_SelfForcing,
model_config=generator_transformer3d.config,
)
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0 or zero_stage == 3:
if fsdp_plugin is not None or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
accelerate_state_dict = {k: v.to(dtype=weight_dtype) for k, v in accelerate_state_dict.items()}
accelerate_state_dict = {
k: v.to(dtype=torch.float32 if k.startswith("gan_head.") else weight_dtype)
for k, v in accelerate_state_dict.items()
}
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
@@ -1149,10 +1250,14 @@ def main():
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
if args.use_ema and models[0] is generator_transformer3d:
ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
models[0].save_pretrained(os.path.join(output_dir, "transformer"))
model = accelerator.unwrap_model(models[0])
save_dir = os.path.join(output_dir, "transformer")
if hasattr(model, "gan_head"):
from safetensors.torch import save_file
model.save_config(save_dir)
save_file(model.state_dict(), os.path.join(save_dir, "diffusion_pytorch_model.safetensors"), metadata={"format": "pt"})
else:
model.save_pretrained(save_dir)
if not args.use_deepspeed:
weights.pop()
@@ -1160,34 +1265,23 @@ def main():
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
if args.use_ema and models[0] is generator_transformer3d:
ema_path = os.path.join(input_dir, "transformer_ema")
_, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True)
load_model = WanTransformer3DModel_SelfForcing.from_pretrained(
input_dir, subfolder="transformer_ema",
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
)
load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config)
load_model.load_state_dict(ema_kwargs)
ema_generator_transformer3d.load_state_dict(load_model.state_dict())
ema_generator_transformer3d.to(accelerator.device)
del load_model
for i in range(len(models)):
# pop models so that they are not loaded again
model = models.pop()
# load diffusers style into model
load_model = WanTransformer3DModel.from_pretrained(
input_dir, subfolder="transformer",
low_cpu_mem_usage=True,
)
model.register_to_config(**load_model.config)
model.load_state_dict(load_model.state_dict())
del load_model
# GAN heads are recreated from the same CLI configuration before resume.
if hasattr(model, "gan_head"):
from safetensors.torch import load_file
model.load_state_dict(load_file(os.path.join(input_dir, "transformer", "diffusion_pytorch_model.safetensors")), strict=True)
else:
load_model = WanTransformer3DModel.from_pretrained(
input_dir,
subfolder="transformer",
low_cpu_mem_usage=True,
)
model.register_to_config(**load_model.config)
model.load_state_dict(load_model.state_dict())
del load_model
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
@@ -1201,6 +1295,32 @@ def main():
accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)
if args.use_ema:
def save_ema_hook(models, weights, output_dir):
if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process:
ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
def load_ema_hook(models, input_dir):
ema_path = os.path.join(input_dir, "transformer_ema")
if isinstance(ema_generator_transformer3d, FSDPEMA):
ema_generator_transformer3d.load_pretrained(ema_path)
else:
_, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True)
load_model = WanTransformer3DModel_SelfForcing.from_pretrained(
input_dir, subfolder="transformer_ema",
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=torch.float32,
)
load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config)
load_model.load_state_dict(ema_kwargs)
ema_generator_transformer3d.load_state_dict(load_model.state_dict())
ema_generator_transformer3d.to(accelerator.device)
del load_model
accelerator.register_save_state_pre_hook(save_ema_hook)
accelerator.register_load_state_pre_hook(load_ema_hook)
if args.gradient_checkpointing:
generator_transformer3d.enable_gradient_checkpointing()
fake_score_transformer3d.enable_gradient_checkpointing()
@@ -1267,8 +1387,8 @@ def main():
fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters()))
fake_trainable_params_optim = [
{'params': [], 'lr': args.learning_rate},
{'params': [], 'lr': args.learning_rate / 2},
{'params': [], 'lr': args.learning_rate_critic},
{'params': [], 'lr': args.learning_rate_critic / 2},
]
in_already = []
for name, param in fake_score_transformer3d.named_parameters():
@@ -1281,7 +1401,7 @@ def main():
high_lr_flag = True
fake_trainable_params_optim[0]['params'].append(param)
if accelerator.is_main_process:
print(f"Set {name} to lr : {args.learning_rate}")
print(f"Set {name} to lr : {args.learning_rate_critic}")
break
if high_lr_flag:
continue
@@ -1290,9 +1410,15 @@ def main():
in_already.append(name)
fake_trainable_params_optim[1]['params'].append(param)
if accelerator.is_main_process:
print(f"Set {name} to lr : {args.learning_rate / 2}")
print(f"Set {name} to lr : {args.learning_rate_critic / 2}")
break
if use_gan:
fake_trainable_params_optim[0]['params'].extend(
param for name, param in fake_score_transformer3d.named_parameters()
if name.startswith("gan_head.") and name not in in_already
)
if args.use_came:
optimizer = optimizer_cls(
trainable_params_optim,
@@ -1334,7 +1460,7 @@ def main():
args.random_hw_adapt = False
# Get the dataset
if args.train_mode != "normal" or args.use_teacher_forcing:
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
train_dataset = ImageVideoDataset(
args.train_data_meta, args.train_data_dir,
video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames,
@@ -1362,7 +1488,7 @@ def main():
return length_to_frame_num
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing:
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan:
aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
batch_sampler_generator = torch.Generator().manual_seed(args.seed)
batch_sampler = AspectRatioBatchImageVideoSampler(
@@ -1672,7 +1798,7 @@ def main():
)
fake_score_lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
optimizer=critic_optimizer,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * accelerator.num_processes,
)
@@ -1798,6 +1924,8 @@ def main():
vae_stream_2 = None
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
min_score_timestep = int(0.02 * noise_scheduler.config.num_train_timesteps)
max_score_timestep = int(0.98 * noise_scheduler.config.num_train_timesteps)
def randomize_denoising_step_indices(
denoising_step_indices_list,
@@ -1852,12 +1980,16 @@ def main():
return result
for epoch in range(first_epoch, args.num_train_epochs):
if global_step >= args.max_train_steps:
break
train_dmd_loss = 0.0
# Number of generator backward contributions since the last log flush; the
# generator only backprops every gen_update_interval batches, so its metrics
# must be averaged by contribution count, not by gradient_accumulation_steps.
train_gen_log_count = 0
train_denoising_loss = 0.0
train_gan_g_loss = 0.0
train_gan_d_loss = 0.0
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
for step, batch in enumerate(train_dataloader):
# Data batch sanity check
@@ -1878,7 +2010,7 @@ def main():
save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.mp4", rescale=True)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
if args.train_mode != "normal" or args.use_teacher_forcing:
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
# Convert images to latent space
pixel_values = batch["pixel_values"].to(weight_dtype)
@@ -1942,8 +2074,9 @@ def main():
temp_n_frames = (temp_n_frames - 1) // sample_n_frames_bucket_interval + 1
pixel_values = pixel_values[:, :temp_n_frames, :, :]
mask_pixel_values = mask_pixel_values[:, :temp_n_frames, :, :]
mask = mask[:, :temp_n_frames, :, :]
if args.train_mode != "normal":
mask_pixel_values = mask_pixel_values[:, :temp_n_frames, :, :]
mask = mask[:, :temp_n_frames, :, :]
# Keep all node same token length to accelerate the traning when resolution grows.
if args.keep_all_node_same_token_length:
@@ -1966,8 +2099,9 @@ def main():
actual_video_length = (actual_video_length - 1) // sample_n_frames_bucket_interval + 1
pixel_values = pixel_values[:, :actual_video_length, :, :]
mask_pixel_values = mask_pixel_values[:, :actual_video_length, :, :]
mask = mask[:, :actual_video_length, :, :]
if args.train_mode != "normal":
mask_pixel_values = mask_pixel_values[:, :actual_video_length, :, :]
mask = mask[:, :actual_video_length, :, :]
if args.low_vram:
torch.cuda.empty_cache()
@@ -1990,7 +2124,7 @@ def main():
pixel_values_bs = pixel_values_bs.sample()
new_pixel_values.append(pixel_values_bs)
return torch.cat(new_pixel_values, dim = 0)
if args.use_teacher_forcing:
if args.use_teacher_forcing or use_gan:
clean_latents = _batch_encode_vae(pixel_values)
else:
clean_latents = None
@@ -2022,7 +2156,7 @@ def main():
clip_context.append(_clip_context)
clip_context = torch.cat(clip_context)
if args.use_teacher_forcing:
if args.use_teacher_forcing or use_gan:
target_shape = clean_latents.size()
else:
target_shape = mask_latents.size()
@@ -2426,17 +2560,7 @@ def main():
# Final output — slice generated frames (may be < num_frames for variable-length)
generator_pred_full = output_pred[:, :, :num_generated_frames]
# Gradient mask: first block gets no gradient when generating > min frames
# (matches Self-Forcing reference: model/base.py L182-L190)
min_num_frames_score = args.score_num_frames
need_gradient_mask = (num_generated_frames != min_num_frames_score)
gradient_mask = None
if need_gradient_mask:
gradient_mask = torch.ones_like(generator_pred_full, dtype=torch.bool)
if args.independent_first_frame:
gradient_mask[:, :, :1] = False
else:
gradient_mask[:, :, :args.num_frame_per_block] = False
need_gradient_mask = (num_generated_frames != args.score_num_frames)
# Slice for score computation: last score_num_frames frames
if num_generated_frames > args.score_num_frames:
@@ -2450,11 +2574,13 @@ def main():
generator_pred_for_score = generator_pred_full
score_num_frames = num_generated_frames
# Compute score_mask for DMD loss (matches Self-Forcing: dmd.py L199-204)
score_mask = None
if gradient_mask is not None:
mask_offset = num_generated_frames - score_num_frames
score_mask = gradient_mask[:, :, mask_offset:mask_offset + score_num_frames]
if need_gradient_mask:
score_mask = torch.ones_like(generator_pred_for_score, dtype=torch.bool)
if args.independent_first_frame:
score_mask[:, :, :1] = False
else:
score_mask[:, :, :args.num_frame_per_block] = False
# generator_pred = the sliced version for DMD loss
generator_pred = generator_pred_for_score
@@ -2530,7 +2656,7 @@ def main():
generator_noise_list = [generator_noise[i] for i in range(bsz)]
clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None
# Use block_mask for causal training (一次性处理整个视频)
# Use block_mask for causal training
generator_pred = generator_transformer3d(
x=generator_noise_list,
context=prompt_embeds,
@@ -2563,6 +2689,8 @@ def main():
# Common code for both KV cache and block mask training
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep)
generator_denoised_input = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
@@ -2668,17 +2796,53 @@ def main():
train_dmd_loss += avg_dmd_loss.item()
train_gen_log_count += 1
# --- One-Forcing Generator GAN Loss ---
generator_loss = dmd_loss
if args.gan_g_weight > 0:
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
gan_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
gan_timestep = gan_timestep.clamp(min_score_timestep, max_score_timestep)
gan_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng)
gan_inputs = [add_noise(generator_pred, gan_noise, gan_timestep)]
if args.relativistic_discriminator:
real_latents = clean_latents[:, :, :generator_pred.size(2)].to(generator_pred.dtype)
gan_inputs.append(add_noise(real_latents, gan_noise, gan_timestep))
num_gan_inputs = len(gan_inputs)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, gan_logits = fake_score_transformer3d(
x=torch.cat(gan_inputs).to(dtype=weight_dtype),
context=repeat_batch(prompt_embeds, num_gan_inputs),
t=gan_timestep.repeat(num_gan_inputs),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, num_gan_inputs) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, num_gan_inputs) if args.train_mode != "normal" else None,
gan_classify=True,
)
gan_logits = gan_logits.float().chunk(num_gan_inputs)
if args.relativistic_discriminator:
gan_g_loss = F.softplus(-(gan_logits[0] - gan_logits[1])).mean()
else:
gan_g_loss = F.softplus(-gan_logits[0]).mean()
gan_g_loss = gan_g_loss * args.gan_g_weight
train_gan_g_loss += accelerator.gather(gan_g_loss.detach().repeat(args.train_batch_size)).mean().item()
generator_loss = generator_loss + gan_g_loss
if args.low_vram:
real_score_transformer3d = real_score_transformer3d.to("cpu")
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
if args.gan_g_weight == 0:
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
torch.cuda.empty_cache()
accelerator.backward(dmd_loss)
accelerator.backward(generator_loss)
if accelerator.sync_gradients:
accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
if args.gan_g_weight > 0:
# Discard discriminator gradients from the generator GAN backward.
critic_optimizer.zero_grad()
if args.low_vram:
fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device)
@@ -2967,6 +3131,8 @@ def main():
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
critic_timestep = critic_timestep.clamp(min_score_timestep, max_score_timestep)
critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng)
fake_score_denoised_input = add_noise(
@@ -2975,15 +3141,25 @@ def main():
critic_timestep
)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
)
# Generated and real latents share critic_noise / critic_timestep. Prepare the
# real input and the R1/R2 perturbations once; they are consumed by the critic
# forward below and by the R1/R2 penalties after the main backward(s).
real_noisy = None
r1_perturbation = None
r2_perturbation = None
if use_gan:
real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype)
real_noisy = add_noise(real_latents, critic_noise, critic_timestep)
del real_latents
if args.r1_weight > 0:
r1_perturbation = args.r1_sigma * torch.randn(
real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng
)
if args.r2_weight > 0:
r2_perturbation = args.r2_sigma * torch.randn(
fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype,
device=fake_score_denoised_input.device, generator=torch_rng
)
def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
noise_pred = noise_pred.float()
@@ -2997,11 +3173,136 @@ def main():
final_loss = masked_loss.mean()
return final_loss
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
accelerator_fake_score_transformer3d.backward(denoising_loss)
# --- One-Forcing Discriminator GAN Loss ---
if use_gan and not args.relativistic_discriminator:
# Memory-optimized critic: the non-relativistic gan_d_loss =
# softplus(-real) + softplus(fake) is separable, so run fake and real as two
# sequential 1x-batch forwards and backward the fake pass first. Backwarding
# fake frees its gradient-checkpointing block inputs before the real forward,
# dropping the packed fake+real+R1+R2 forward peak (up to 4x batch) down to
# 1x batch here; R1/R2 re-forward their own 2x batch below. Both backward(s)
# accumulate into fake_score.grad, giving the identical update as the packed
# single-forward critic.
# -- Fake pass (1x batch): denoising output + fake discriminator logit --
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output, fake_gan_logits = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
gan_classify=True,
)
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
fake_logit = fake_gan_logits.float()
fake_d_loss = F.softplus(fake_logit).mean() * args.gan_d_weight
critic_loss = denoising_loss + fake_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
# Fake graph freed: the transformer blocks' stored inputs are released here.
fake_d_detached = fake_d_loss.detach()
del critic_loss, fake_gan_logits, fake_logit, fake_d_loss, fake_score_denoised_output
# -- Real pass (1x batch): discriminator logit only --
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, real_gan_logits = fake_score_transformer3d(
x=real_noisy,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
gan_classify=True,
)
real_logit = real_gan_logits.float()
real_d_loss = F.softplus(-real_logit).mean() * args.gan_d_weight
accelerator_fake_score_transformer3d.backward(real_d_loss)
# Real graph freed. Combine the two detached halves for logging.
gan_d_loss = fake_d_detached + real_d_loss.detach()
train_gan_d_loss += accelerator.gather(gan_d_loss.repeat(args.train_batch_size)).mean().item()
del fake_d_detached, real_gan_logits, real_logit, real_d_loss
else:
# Combined path: 2x-batch forward for relativistic GAN (both logits coupled by
# softplus(-(real - fake))), or 1x-batch forward when GAN is disabled.
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
if use_gan:
fake_score_denoised_output, gan_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0]
gan_logits = gan_logits.float().chunk(2)
else:
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
)
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
critic_loss = denoising_loss
if use_gan:
fake_logit, real_logit = gan_logits[0], gan_logits[1]
if args.relativistic_discriminator:
gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
else:
gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
gan_d_loss = gan_d_loss * args.gan_d_weight
train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
critic_loss = critic_loss + gan_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
del critic_loss
# R1/R2 gradient penalties: each re-forwards its own 2x batch after the main
# backward(s) and finite-differences the perturbed logit against the clean one,
# matching the packed-formula value while keeping the peak at 2x batch.
if use_gan and args.r1_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, r1_logits = fake_score_transformer3d(
x=torch.cat([real_noisy + r1_perturbation, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
r1_logits = r1_logits.float().chunk(2)
r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean()
accelerator_fake_score_transformer3d.backward(r1_loss)
del r1_loss
if use_gan and args.r2_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, r2_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
r2_logits = r2_logits.float().chunk(2)
r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean()
accelerator_fake_score_transformer3d.backward(r2_loss)
del r2_loss
if accelerator_fake_score_transformer3d.sync_gradients:
accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm)
critic_optimizer.step()
@@ -3019,10 +3320,15 @@ def main():
ema_generator_transformer3d.step(generator_transformer3d.parameters())
progress_bar.update(1)
global_step += 1
accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step)
logs = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}
if use_gan:
logs.update({"train_gan_g_loss": train_gan_g_loss / max(train_gen_log_count, 1), "train_gan_d_loss": train_gan_d_loss})
accelerator.log(logs, step=global_step)
train_dmd_loss = 0.0
train_gen_log_count = 0
train_denoising_loss = 0.0
train_gan_g_loss = 0.0
train_gan_d_loss = 0.0
if global_step % args.checkpointing_steps == 0:
if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
@@ -3092,6 +3398,9 @@ def main():
if global_step >= args.max_train_steps:
break
if global_step >= args.max_train_steps:
break
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
with progress_bar.paused():
if args.use_ema:
@@ -3121,12 +3430,6 @@ def main():
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
if args.use_ema:
# Copy EMA weights to generator on ALL processes before save_state (collective op).
ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
if accelerator.is_main_process:
generator_transformer3d = unwrap_model(generator_transformer3d)
if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
gc.collect()
torch.cuda.empty_cache()
@@ -3135,8 +3438,6 @@ def main():
fake_score_save_path = os.path.join(save_path, "fake_score")
accelerator.save_state(save_path)
accelerator_fake_score_transformer3d.save_state(fake_score_save_path)
if args.use_ema and accelerator.is_main_process:
ema_generator_transformer3d.save_pretrained(os.path.join(save_path, "transformer_ema"))
logger.info(f"Saved state to {save_path}")
accelerator.end_training()
scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL.md
+10 -22
View File
@@ -12,11 +12,9 @@ This document provides the complete workflow for **Flex-Forcing Stage 2 — Dist
> 2. **§3.2 Pyramid timestep chunking** — after finishing step `t+1`, the partition for step `t` is *refined* by inserting boundaries while preserving the existing ones. The step-`t` result is buffered over the whole original chunk, then each sub-chunk resumes autoregressively once the KV it needs is available. Coarse planning first, fine detail later.
> 3. **§3.3 K-Projection** — clean-context keys live in a different space from noisy queries at timestep `t`, so an identity-initialised, timestep-conditioned projection maps the cache into the current noisy latent space. It is applied on the fly, never mutates the cache, and trains together with the generator.
>
> Training draws a **fresh random partition every iteration** (chunk sizes 2–10), which is what lets one set of weights cover the whole causal↔bidirectional spectrum. §4.2 of the paper then reuses the same machinery for **any-order editing**.
>
> The Flex-Forcing pipeline has two stages:
>
> 1. **Stage 1 — ODE Regression** (`train_ode.py`): regress the teacher's ODE trajectory so the model becomes a competent few-step generator. This is where the random partitions are first seen.
> 1. **Stage 1 — ODE Regression** (`train_ode.py`): regress the teacher's ODE trajectory so the model becomes a competent few-step generator.
> 2. **Stage 2 — DMD Distillation** (`train_distill.py`, this README): compress to a **4-step** generator with a **Wan2.1-T2V-14B** real-score teacher, and add the pyramid timestep schedule.
>
> This README covers **Stage 2 only**. See [README_TRAIN_ODE.md](./README_TRAIN_ODE.md) for Stage 1.
@@ -210,7 +208,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
The launcher pins: 81 pixel frames = **21 latent frames** (a 5 s clip) at **480×832**, a **4-step** model (`--denoising_step_indices_list 1000 750 500 250`, jittered per iteration by `--randomize_step_indices`), `--flex_pyramid_levels=4` - which over 21 latent frames reproduces the paper's ladder `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` - partitions drawn in **2..10** (plus, on half the iterations, the whole-clip coarse layout of §3.2 - see below), `--flow_euler_rollout`, `diag_rank1` K-Projection, batch 8 (8 GPUs × `--gradient_accumulation_steps=1`).
The launcher pins: 81 pixel frames = **21 latent frames** (a 5 s clip) at **480×832**, a **4-step** model (`--denoising_step_indices_list 1000 750 500 250`, jittered per iteration by `--randomize_step_indices`), `--flex_pyramid_levels=4` - which over 21 latent frames reproduces the paper's ladder `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` - the *random* arm's partitions drawn in **2..10** (that arm is only 10% of iterations; another 60% start from the whole-clip coarse layout of §3.2 - see below), `--flow_euler_rollout`, `diag_rank1` K-Projection, batch 8 (8 GPUs × `--gradient_accumulation_steps=1`).
Output: `output_dir_wan2.1_flex_forcing_distill/`.
@@ -243,7 +241,7 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
| `--denoising_step_indices_list` | Denoising step indices (DMD core param). Mapped through `timesteps[train_sampling_steps - idx]`, so `1000 750 500 250` becomes t = `[1000, 937.5, 833.33, 625]` at `shift=5.0` - bit-identical to what `stochastic_sampling_timesteps(4, shift)` feeds the pipeline. | `1000 750 500 250` |
| `--randomize_step_indices` | Jitter every index but the first symmetrically within `--index_jitter_ratio` of its neighbouring gap each iteration. The expected schedule is unchanged and monotonicity is enforced by construction, so inference's fixed t stays inside every step's trained envelope. | off |
| `--index_jitter_ratio` | Jitter budget as a fraction of the neighbouring gap. At `1000 750 500 250` with 0.3, 20 000 draws give idx `[675,825] / [425,575] / [175,325]`, i.e. t `[912.2,959.3] / [787.0,871.2] / [514.7,706.5]`; the inference points 937.5 / 833.33 / 625 all fall inside. | 0.3 |
| `--flow_euler_rollout` | Advance the generator's self-rollout with a deterministic Euler ODE step in flow space (`x_next = x_t + (sigma_next - sigma_t) * v`, fp32) instead of converting to x0 and re-noising with fresh noise. Algebraically this *is* re-noising with the same implied noise - measured `8.9e-16` apart in fp64, versus up to `6.4` for a fresh draw - so it removes the rollout's stochasticity and matches normal flow-matching inference. The final step still converts to x0, which is what DMD is defined on. The critic re-rolls the generator and follows the same switch. | off |
| `--flow_euler_rollout` | Advance the generator's self-rollout with a deterministic Euler ODE step in flow space (`x_next = x_t + (sigma_next - sigma_t) * v`, fp32) instead of converting to x0 and re-noising with fresh noise. Algebraically this *is* re-noising with the same implied noise - identical to it in exact arithmetic, whereas drawing fresh noise is an independent sample - so it removes the rollout's stochasticity and matches normal flow-matching inference. The final step still converts to x0, which is what DMD is defined on. The critic re-rolls the generator and follows the same switch. | off |
| `--output_dir` | Output directory | `output_dir_wan2.1_flex_forcing_distill` |
| `--gradient_checkpointing` | Trade compute for memory | - |
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
@@ -256,24 +254,14 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
| Parameter | Paper | Description | Default |
|-----------|-------|-------------|---------|
| `--flex_forcing` | §3.1 | Instantiate `WanTransformer3DModel_FlexForcing` and validate with `WanFlexForcingPipeline`. Off = the inherited path, unchanged. | off |
| `--flex_chunk_min` | §3.1 | Smallest chunk drawn per iteration. `1` is legal but adds no coverage under a pyramid - refining a 2-frame chunk already yields 1-frame leaves - and it would let single-frame chunks appear mid-clip. | 2 |
| `--flex_chunk_max` | §3.1 | Largest chunk drawn per iteration. Set equal to `--flex_chunk_min` to pin one fixed layout. **Keep it below the latent frame count**: the whole-clip layout is already covered by the mixture below, so raising this to the frame count only spends draws on a layout you already have and thins out §3.1 (at 21 latent frames, `max=21` leaves 3649 distinct partitions versus 4882 at `max=10`). | 10 |
| `--flex_pyramid_levels` | §3.2 | `1` = one partition per iteration. `>1` = nest it into a coarse-to-fine ladder, one level per denoising step; both training paths roll it out. **Stage 2 only.** | 1 |
| `--flex_chunk_min` | §5 | Smallest chunk drawn by the *random* arm (10% of iterations, regardless of `--flex_pyramid_levels`); the other three arms pin their own layouts and split the remaining 90% equally. `1` is legal but adds no coverage under a pyramid - refining a 2-frame chunk already yields 1-frame leaves - and it would let single-frame chunks appear mid-clip. | 2 |
| `--flex_chunk_max` | §5 | Largest chunk drawn by the *random* arm. Set equal to `--flex_chunk_min` to pin that arm to one fixed layout. **Keep it below the latent frame count**: the whole-clip layout is its own arm, so raising this to the frame count only spends random-arm draws on a layout you already have and thins out every other one (exhaustive enumeration at 21 latent frames: 6441 distinct partitions at `max=10`, 6765 at `max=21` - but `max=21` makes `[21]` alone 1/19 of that arm and cuts `P([3]*7)` from `1/183708` to `1/1106560`). | 10 |
| `--flex_pyramid_levels` | §3.2 | `1` = one partition per iteration. `>1` = nest it into a coarse-to-fine ladder, one level per denoising step; both training paths roll it out. The ladder itself is **Stage 2 only** (Stage 1 regresses `x0` in one forward per chunk, so it has no ladder to build), but Stage 1 **must still carry the same value**: it decides whether the shared level-0 mixture's coarse arm exists. | 1 |
| `--flex_min_num_frame_per_block` | §3.2 | Block size the refinement stops at; `1` is fully causal at the leaves. | 1 |
| `--flex_self_generated_context` | §3.3 | Feeds the K-Projection the model's own previous-step `x0` prediction as context, so it enters the autograd graph under prompt-only training. Costs a doubled attention sequence; from denoising step 1 on. Inert under `--use_kv_cache_training` - that path passes no `flex_state`, so the K-Projection is never invoked and stays at its identity init. | off |
| `--num_frame_per_block` | — | The uniform fallback for the inherited (non-Flex) path; with Flex on it is also one of the three level-0 bands (`UNIFORM_BLOCK_PROB = 0.1` of iterations pin level 0 to its uniform partition), and it is what `log_validation` samples with at `--flex_pyramid_levels=1` - see the two bullets below. | 3 |
| `--flex_self_generated_context` | §3.3 | Feeds the K-Projection the model's own previous-step `x0` prediction as context, so it enters the autograd graph under prompt-only training. Costs a doubled attention sequence; from denoising step 1 on. Inert under `--use_kv_cache_training`, but **not because the K-Projection is idle** - that path already commits its own clean context into the KV cache and needs no extra `clean_x` half; Π itself **is** invoked there. | off |
| `--num_frame_per_block` | — | The uniform fallback for the inherited (non-Flex) path; with Flex on it is also reached inside level 0's four-arm mixture (the uniform arm is 30% of iterations and one of its three equal sub-arms pins level 0 to its uniform partition, so `[nfpb x k]` alone is 10%), and it is what `log_validation` samples with at `--flex_pyramid_levels=1`. | 3 |
| `--independent_first_frame` | — | Prepend a length-1 chunk (`[1, N, N, ...]`); honoured by the partition sampler. | off |
Seven interactions are worth stating explicitly:
- **Level 0 is a three-band mixture decided by one draw.** `denoise_mode="pyramid"` at inference builds its level 0 as the whole clip in one chunk (fully bidirectional planning), so a partition drawn only from 2..10 would leave that step untrained - and the uniform `--num_frame_per_block` layout is just as unreachable by chance alone (at 21 latent frames `[3]*7` has probability `1/183708`, i.e. 0.11 expected hits in 20 000 iterations). `train_distill.py` therefore splits level 0 with a single uniform draw into three bands, neither constant a CLI flag: `COARSE_GLOBAL_PROB = 0.5` for the whole clip as one chunk, `UNIFORM_BLOCK_PROB = 0.1` pinned to the uniform `--num_frame_per_block` partition, and the remaining `0.4` for §3.1's random draw; all three refine into the same ladder. Measured over 200 000 iterations at 21 latent frames with `--flex_pyramid_levels=4`: whole clip `49.999%`, `[3]*7` `10.034%`, random `39.968%`, **4078** distinct level-0 layouts, and the inference ladder `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` at `49.999%`. At `--flex_pyramid_levels=1` the whole-clip band is empty and the split becomes `10% / 90%`, which keeps that mode as the paper describes it. Peak memory is set by the densest case (level 0 = whole clip), so budget for it.
- The ladder is **clamped to the number of trained denoising steps**. A 4-level ladder under a 2-step schedule can never reach its last two levels, so `sample_flex_partitions` drops the unreachable tail rather than silently reporting it as trained. Keep `len(--denoising_step_indices_list) >= --flex_pyramid_levels`; the launcher's 4 indices and `--flex_pyramid_levels=4` are matched on purpose. The converse also wastes compute: more steps than levels leaves `install_flex_partition` clamped on its finest level for the extra steps.
- **`--flow_euler_rollout` keeps the prediction in flow space between steps, so §3.3's context needs its own x0.** `--flex_self_generated_context` feeds the K-Projection the previous step's *clean* prediction; under an Euler rollout that variable is a velocity, so both rollout loops convert a detached copy for the context while the rollout state itself advances by Euler. The generator and the critic convert identically, otherwise the critic would score a rollout that is not the one being trained.
- **`--randomize_step_indices` does not break the pyramid.** The jitter moves the *timesteps*, not the partitions, and `sample_flex_partitions` keys off `len(denoising_step_list)`, which the jitter preserves. Every level of the ladder still lines up with its denoising step.
- **`log_validation` renders the layout training renders.** With `--flex_pyramid_levels > 1` it passes `chunk_spec=None` + `denoise_mode="pyramid"`, so its ladder is identical to `predict_t2v.py`'s and measured in `49.999%` of iterations. It used to pin `chunk_spec` to `--num_frame_per_block`, which nested `[[3]*7, [2,1]*7, [1]*21]` - a **3-level** ladder under a 4-step schedule (refinement hits the all-ones fixed point) that the trainer drew **0 times in 20 000** iterations, so those clips said nothing about the model being trained. At `--flex_pyramid_levels=1` the conclusion reverses: the whole-clip band is empty there and `[21]` is never drawn, so validation uses the uniform `--num_frame_per_block` partition - the one band that is both fixed across checkpoints and genuinely trained (`10.066%` of iterations, against ~`1.3%` for the most common random partition). On that 3-level ladder, now covered `10.034%` of the time: step 4 reuses the finest level `[1]*21` through `install_flex_partition`'s index clamp (`min(step_index, len(partitions) - 1)`), so nothing runs off the end.
- §3.3 has **no launcher flag**: the config yaml chooses the variant through `transformer_additional_kwargs.flex_kproj_mode` (default `diag_rank1`; set it to `none` there to ablate Pi, `diag` for the cheaper head-only map). Its parameters are named `flex_kproj.proj_out.*` on purpose - that name is what lets `initialize_missing_parameters()` zero-fill them when a checkpoint predates §3.3, which is how identity init survives `low_cpu_mem_usage=True` - and they otherwise train like any other layer, following `--learning_rate`. The trainer forces them `requires_grad` whenever the built model carries them, so a narrowed `--trainable_modules` cannot silently freeze Pi at its identity. Under `--use_kv_cache_training` it is never invoked and stays at its identity init as well.
- Every rank must train the same layout, otherwise the FlexAttention mask and the `num_frame_per_block` derived from it disagree across the SP/FSDP group. The trainer therefore passes the drawn partition through `broadcast_chunk_sizes` before using it, and `randomize_denoising_step_indices` broadcasts its jittered indices from rank 0 for the same reason.
### 3.4 Training Validation
You can configure validation parameters to periodically generate test videos during training, allowing you to monitor training progress and model quality.
@@ -757,7 +745,7 @@ num_frame_per_block = 7 # clean-context commit granularity; None = one
This `num_frame_per_block` is the same knob as Self-Forcing's: the clean context is committed in uniform blocks of that many frames, in temporal order. It differs in exactly two ways - the trailing chunk absorbs the remainder so divisibility is not required (the inherited rollout asserts it), and the transformer's block width becomes `max(it, edit span width)`.
Any-order editing needs the full KV cache: both entries set `local_attn_size = -1`. A pyramid needs the same, plus `stochastic_sampling = True` — the buffered step between two levels *is* the schedule's own re-noising. The pipeline raises rather than silently degrading if either is violated.
Any-order editing needs the full KV cache: both entries set `local_attn_size = -1`. Any per-step ladder (`denoise_mode="pyramid"` or `"full_then_blocks"`) needs the same, plus `stochastic_sampling = True` — the buffered step between two levels *is* the schedule's own re-noising. The pipeline raises rather than silently degrading if either is violated.
---
@@ -765,7 +753,7 @@ Any-order editing needs the full KV cache: both entries set `local_attn_size = -
- **Flex-Forcing paper**: https://arxiv.org/abs/2607.03509
- **Stage 1**: [README_TRAIN_ODE.md](./README_TRAIN_ODE.md)
- **Shared chunking helpers**: `videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`, `sample_flexible_chunks`, `refine_partition`, `build_pyramid_partitions`, `validate_nested_partitions`, `broadcast_chunk_sizes`
- **Shared chunking helpers**: `videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`, `sample_flexible_chunks`, `sample_flex_level0` (the level-0 four-arm mixture, drawn by **both** trainers), `refine_partition`, `build_pyramid_partitions`, `validate_nested_partitions`, `broadcast_chunk_sizes`
- **Model / pipeline**: `videox_fun/models/wan_transformer3d_flex_forcing.py`, `videox_fun/pipeline/pipeline_wan_flex_forcing.py`
- **Upstream stages**: [README_TRAIN_ODE.md](../wan2.1_self_forcing/README_TRAIN_ODE.md), [README_TRAIN_DISTILL.md](../wan2.1_self_forcing/README_TRAIN_DISTILL.md)
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md
+9 -25
View File
@@ -12,11 +12,9 @@
> 2. **§3.2 金字塔时间步分块** —— 完成步 `t+1` 后，步 `t` 的划分是在**保留已有边界的前提下插入新边界**细化的。步 `t` 的结果先在整个原 chunk 上缓冲，随后每个子 chunk 在其所需 KV 可用后逐个自回归续写。先粗规划、后细刻画。
> 3. **§3.3 K-Projection** —— 干净上下文的 key 与时间步 `t` 的含噪 query 处在不同的空间，因此用一个恒等初始化、以时间步为条件的投影把 cache 映射进当前含噪 latent 空间。它即时施加、从不改写 cache，并与生成器一起训练。
>
> 训练时每次迭代都抽取**全新的随机划分**（chunk 大小 2–10），这正是让同一套权重覆盖整个"因果↔双向"谱系的原因。论文 §4.2 随后把同一套机制复用为**任意序编辑**。
>
> Flex-Forcing 流程分两个阶段：
>
> 1. **阶段 1——ODE 回归**（`train_ode.py`）：对教师模型的 ODE 轨迹做回归，让模型先成为一个合格的小步数生成器。随机分块在这一阶段就首次出现。
> 1. **阶段 1——ODE 回归**（`train_ode.py`）：对教师模型的 ODE 轨迹做回归，让模型先成为一个合格的小步数生成器。
> 2. **阶段 2——DMD 蒸馏**（`train_distill.py`，即本 README）：用 **Wan2.1-T2V-14B** 真实分数教师蒸馏出 **4 步**生成器，并加入金字塔时间步调度。
>
> 本 README 只覆盖**阶段 2**。阶段 1 见 [README_TRAIN_ODE_zh-CN.md](./README_TRAIN_ODE_zh-CN.md)。
@@ -210,7 +208,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒片段）、分辨率 **480×832**、**4 步**模型（`--denoising_step_indices_list 1000 750 500 250`，并由 `--randomize_step_indices` 每次迭代抖动）、`--flex_pyramid_levels=4` —— 在 21 个 latent 帧上正好复现论文的阶梯 `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` —— 划分抽取范围 **2..10**（另有半数的迭代以 §3.2 的整段粗粒度布局起步，见下）、`--flow_euler_rollout`、`diag_rank1` K-Projection、batch 8（8 卡 × `--gradient_accumulation_steps=1`）。
launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒片段）、分辨率 **480×832**、**4 步**模型（`--denoising_step_indices_list 1000 750 500 250`，并由 `--randomize_step_indices` 每次迭代抖动）、`--flex_pyramid_levels=4` —— 在 21 个 latent 帧上正好复现论文的阶梯 `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` —— 随机档的划分抽取范围 **2..10**（该档只占 10% 的迭代；另有 60% 以 §3.2 的整段粗粒度布局起步，见下）、`--flow_euler_rollout`、`diag_rank1` K-Projection、batch 8（8 卡 × `--gradient_accumulation_steps=1`）。
输出目录：`output_dir_wan2.1_flex_forcing_distill/`。
@@ -243,7 +241,7 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
| `--denoising_step_indices_list` | 去噪步索引（DMD 核心参数）。经 `timesteps[train_sampling_steps - idx]` 映射，所以 `shift=5.0` 下 `1000 750 500 250` 变成 t = `[1000, 937.5, 833.33, 625]` —— 与 `stochastic_sampling_timesteps(4, shift)` 喂给 pipeline 的值逐位一致。 | `1000 750 500 250` |
| `--randomize_step_indices` | 每次迭代对除首个索引以外的所有索引做对称抖动，幅度为相邻间隔的 `--index_jitter_ratio`。期望调度不变，单调性由构造保证，所以推理用的固定 t 始终落在每一步的训练包络内。 | 关 |
| `--index_jitter_ratio` | 抖动预算（相邻间隔的比例）。在 `1000 750 500 250` 且取 0.3 时，20000 次抽取给出 idx `[675,825] / [425,575] / [175,325]`，即 t `[912.2,959.3] / [787.0,871.2] / [514.7,706.5]`；推理点 937.5 / 833.33 / 625 全部落在其中。 | 0.3 |
| `--flow_euler_rollout` | 生成器自滚用 flow 空间里的确定性 Euler ODE 步推进（`x_next = x_t + (sigma_next - sigma_t) * v`，fp32），而不是转成 x0 再用新噪声重加噪。代数上这**就等于**用同一个隐含噪声重加噪 —— fp64 下实测相差 `8.9e-16`，而重采新噪声最大差 `6.4` —— 因此它消掉了 rollout 的随机性，与正常的 flow-matching 推理一致。最后一步仍然转 x0，因为 DMD 目标定义在 x0 上。critic 重滚生成器，跟随同一开关。 | 关 |
| `--flow_euler_rollout` | 生成器自滚用 flow 空间里的确定性 Euler ODE 步推进（`x_next = x_t + (sigma_next - sigma_t) * v`，fp32），而不是转成 x0 再用新噪声重加噪。代数上这**就等于**用同一个隐含噪声重加噪 —— 在精确算术下与之完全相同，而重采新噪声是一次独立抽样 —— 因此它消掉了 rollout 的随机性，与正常的 flow-matching 推理一致。最后一步仍然转 x0，因为 DMD 目标定义在 x0 上。critic 重滚生成器，跟随同一开关。 | 关 |
| `--output_dir` | 输出目录 | `output_dir_wan2.1_flex_forcing_distill` |
| `--gradient_checkpointing` | 用算力换显存 | - |
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
@@ -256,24 +254,14 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
| 参数 | 论文 | 说明 | 默认值 |
|-----------|-------|-------------|---------|
| `--flex_forcing` | §3.1 | 实例化 `WanTransformer3DModel_FlexForcing` 并用 `WanFlexForcingPipeline` 校验。关闭 = 走继承的原路径，行为不变。 | 关 |
| `--flex_chunk_min` | §3.1 | 每次迭代抽取的最小 chunk 大小。金字塔下 `1` 合法但不增加覆盖：2 帧 chunk 细化后本来就会产生 1 帧叶子，而它会让单帧 chunk 出现在片段中间。 | 2 |
| `--flex_chunk_max` | §3.1 | 每次迭代抽取的最大 chunk 大小。设为与 `--flex_chunk_min` 相等即固定一种布局。**必须小于 latent 帧数**：“整段一个 chunk”这个布局已经由下面的混合机制覆盖，把这个值抬到帧数只是把抽取次数花在一个你已经有的布局上，反而稀释 §3.1（21 个 latent 帧时，`max=21` 只剩 3649 种不同划分，而 `max=10` 有 4882 种）。 | 10 |
| `--flex_pyramid_levels` | §3.2 | `1` = 每次迭代一个划分。`>1` = 把它嵌套成由粗到细的阶梯，每个去噪步一层；两条训练路径都会展开。**仅阶段 2 有。** | 1 |
| `--flex_chunk_min` | §5 | **random 档**（占 10% 迭代，与 `--flex_pyramid_levels` 无关）抽取的最小 chunk 大小；其余三档各自钉定自己的布局、平分剩下的 90%。金字塔下 `1` 合法但不增加覆盖：2 帧 chunk 细化后本来就会产生 1 帧叶子，而它会让单帧 chunk 出现在片段中间。 | 2 |
| `--flex_chunk_max` | §5 | **random 档**抽取的最大 chunk 大小。设为与 `--flex_chunk_min` 相等即把该档钉成一种固定布局。**必须小于 latent 帧数**：“整段一个 chunk”本身就是独立的一档，把这个值抬到帧数只是把 random 档的抽取次数花在一个你已经有的布局上，反而稀释其他所有布局（21 个 latent 帧下穷举计数：`max=10` 有 6441 种不同划分、`max=21` 有 6765 种，但 `max=21` 会让单是 `[21]` 就占该档 1/19 的抽取，并把 `P([3]*7)` 从 1/183708 压到 1/1106560）。 | 10 |
| `--flex_pyramid_levels` | §3.2 | `1` = 每次迭代一个划分。`>1` = 把它嵌套成由粗到细的阶梯，每个去噪步一层；两条训练路径都会展开。阶梯本身**仅阶段 2 有**（阶段 1 每 chunk 一次前向回归 x0，没有阶梯可搭），但阶段 1 **仍然要带同一个值**：它决定共享的 level-0 混合里 coarse 档是否存在。 | 1 |
| `--flex_min_num_frame_per_block` | §3.2 | 细化终止的块大小（每块二分直到不超过它）；`1` 表示叶子级完全因果。 | 1 |
| `--flex_self_generated_context` | §3.3 | 把模型自己上一步的 `x0` 预测喂给 K-Projection 当上下文，使其在纯 prompt 训练下也能进入 autograd 图。代价是注意力序列翻倍；从去噪步 1 起生效。在 `--use_kv_cache_training` 下不生效 —— 该路径不传 `flex_state`，K-Projection 不会被调用，停在恒等初始化。 | 关 |
| `--num_frame_per_block` | — | 继承路径（未开 Flex）的均匀 fallback；开了 Flex 后它还是 level 0 三档混合中的一档（`UNIFORM_BLOCK_PROB = 0.1` 的迭代把 level 0 钉成它的均匀划分），且 `--flex_pyramid_levels=1` 时 `log_validation` 就用它采样 —— 见下面两条 bullet。 | 3 |
| `--flex_self_generated_context` | §3.3 | 把模型自己上一步的 `x0` 预测作为上下文喂给 K-Projection，因此在纯 prompt 训练下它会进入自动求导图。代价是注意力序列翻倍；从去噪第 1 步起生效。在 `--use_kv_cache_training` 下它不起作用，但**并不是因为 K-Projection 空闲** —— 那条路径已经把自己的干净上下文写进 KV cache，不需要额外的 `clean_x` 一半；Π 本身在那里**是**被调用的。 | 关 |
| `--num_frame_per_block` | — | 继承路径（未开 Flex）的均匀 fallback；开了 Flex 后它还是 level 0 四档混合中 uniform 那一档（该档占 30% 的迭代、内部再三等分，其一把 level 0 钉成它的均匀划分，故 `[nfpb x k]` 单独占 10%），且 `--flex_pyramid_levels=1` 时 `log_validation` 就用它采样。 | 3 |
| `--independent_first_frame` | — | 在划分前插入一个长度为 1 的 chunk（`[1, N, N, ...]`）；划分采样器会遵守它。 | 关 |
有七个交互值得明说：
- **level 0 是三档混合，一次抽样定档。** 推理时 `denoise_mode="pyramid"` 的 level 0 是整段一个 chunk（全双向的全局规划），只从 2..10 抽划分会让那一步从未被训练；而 `--num_frame_per_block` 那个均匀布局光靠随机抽样又几乎碰不到（21 个 latent 帧下 `[3]*7` 的概率是 `1/183708`，20000 次迭代期望 0.11 次）。所以 `train_distill.py` 用一个均匀抽样把 level 0 分成三档，两个常量都故意不做成 CLI 参数：`COARSE_GLOBAL_PROB = 0.5` 整段一个 chunk、`UNIFORM_BLOCK_PROB = 0.1` 钉成 `--num_frame_per_block` 的均匀划分、剩下 `0.4` 是 §3.1 的随机划分；三档都细化成同一条阶梯。实测 200000 次（21 个 latent 帧、`--flex_pyramid_levels=4`）：整段 `49.999%`、`[3]*7` `10.034%`、随机 `39.968%`，共 **4078** 种不同的 level-0 布局，推理那条阶梯 `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` 占 `49.999%`。`--flex_pyramid_levels=1` 时整段那一档为空，比例变成 `10% / 90%`，与论文描述一致。峰值显存由最密的那种情形（level 0 = 整段）决定，请据此预留。
- 阶梯会被**截断到训练的去噪步数**。2 步调度下配 4 级阶梯永远到不了最后两级，所以 `sample_flex_partitions` 会丢弃到不了的那部分，而不是默默把它们当作已训练。请保持 `len(--denoising_step_indices_list) >= --flex_pyramid_levels`；launcher 里的 4 个索引与 `--flex_pyramid_levels=4` 就是故意对齐的。反过来也浪费算力：步数多于级数时，多出来的那些步里 `install_flex_partition` 会被钉在最细的一级上。
- **`--flow_euler_rollout` 会把预测留在 flow 空间，所以 §3.3 的上下文需要自己那份 x0。** `--flex_self_generated_context` 喂给 K-Projection 的是上一步的**干净**预测；在 Euler rollout 下那个变量是速度，所以两个 rollout 循环都会另外转换一份 detached 副本给上下文用，而 rollout 状态本身仍按 Euler 推进。生成器与 critic 的转换必须完全一致，否则 critic 评的就不是正在被训的那个 rollout。
- **`--randomize_step_indices` 不会破坏金字塔。** 抖动动的是**时间步**而不是划分，而 `sample_flex_partitions` 只看 `len(denoising_step_list)`，抖动不会改变它。阶梯的每一级仍然对得上自己的去噪步。
- **`log_validation` 渲染的就是训练渲染的那个布局。** `--flex_pyramid_levels > 1` 时它传 `chunk_spec=None` + `denoise_mode="pyramid"`，阶梯与 `predict_t2v.py` 的完全相同，实测占 `49.999%` 的迭代。它原本把 `chunk_spec` 钉成 `--num_frame_per_block`，那会嵌套出 `[[3]*7, [2,1]*7, [1]*21]` —— 在 4 步调度下只有 **3 级**（细化到全 1 就撞上固定点），而当时那条阶梯 trainer 在 20000 次迭代里抽出 **0 次**，所以那些校验视频根本说明不了正在训练的模型。`--flex_pyramid_levels=1` 时结论反过来：整段那一档为空、`[21]` 一次也抽不到，所以校验用 `--num_frame_per_block` 的均匀划分 —— 它是唯一既跨 checkpoint 固定、又真被训练的一档（`10.066%` 的迭代，而最高频的随机划分只有约 `1.3%`）。顺带说明那条 3 级阶梯：它现在有 `10.034%` 的覆盖率，第 4 步由 `install_flex_partition` 的索引 clamp（`min(step_index, len(partitions) - 1)`）复用最细一级 `[1]*21`，不会越界。
- §3.3 **没有 launcher 参数**：变体由 config yaml 里的 `transformer_additional_kwargs.flex_kproj_mode` 决定（默认 `diag_rank1`；要消融就在那里改成 `none`，想更省就改 `diag`）。它的参数故意取名 `flex_kproj.proj_out.*` —— 正是这个名字让 `initialize_missing_parameters()` 在 checkpoint 早于 §3.3 时将其置零，从而在 `low_cpu_mem_usage=True` 下保住恒等初始化；除此之外它和其他层一样训，跟 `--learning_rate`。只要构建出的模型带 Pi，trainer 就强制其 `requires_grad`，所以收窄 `--trainable_modules` 也不会把 Pi 悄悄冻在恒等点。在 `--use_kv_cache_training` 下它同样不会被调用，始终停在恒等初始化。
- 每张卡必须训练同一种布局，否则 FlexAttention 掩码与其派生的 `num_frame_per_block` 在 SP/FSDP 组之间会不一致。trainer 在使用抽取的划分前会先经 `broadcast_chunk_sizes` 广播对齐；`randomize_denoising_step_indices` 同样从 rank 0 广播它抖动后的索引，道理一样。
---
### 3.4 训练验证
@@ -746,15 +734,11 @@ num_frame_per_block = 7 # 干净上下文的提交粒度；None = 整
这里的 `num_frame_per_block` 与 Self-Forcing 那个同名参数是一回事：干净上下文按这么多帧一块、依时间序提交。区别只有两点——尾块吸收余数所以不要求整除（继承的 rollout 会断言整除），且 transformer 上的块宽取 `max(它, 编辑区间宽度)`。
任意序编辑需要完整 KV cache：两个入口都设 `local_attn_size = -1`。金字塔还需要 `stochastic_sampling = True`——两级之间的缓冲步本身就是调度自己的重加噪。任何一项不满足时 pipeline 会直接报错，而不是默默降级。
任意序编辑需要完整 KV cache：两个入口都设 `local_attn_size = -1`。任何逐步换划分的 ladder（`denoise_mode="pyramid"` 或 `"full_then_blocks"`）受同一条约束，并额外要求 `stochastic_sampling = True`——两级之间的缓冲步本身就是调度自己的重加噪。任何一项不满足时 pipeline 会直接报错，而不是默默降级。
---
## 五、更多资源
- **Flex-Forcing 论文**：https://arxiv.org/abs/2607.03509
- **阶段 1**：[README_TRAIN_ODE_zh-CN.md](./README_TRAIN_ODE_zh-CN.md)
- **共享分块 helper**：`videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`、`sample_flexible_chunks`、`refine_partition`、`build_pyramid_partitions`、`validate_nested_partitions`、`broadcast_chunk_sizes`
- **模型 / pipeline**：`videox_fun/models/wan_transformer3d_flex_forcing.py`、`videox_fun/pipeline/pipeline_wan_flex_forcing.py`
- **上游阶段**：[README_TRAIN_ODE.md](../wan2.1_self_forcing/README_TRAIN_ODE.md)、[README_TRAIN_DISTILL.md](../wan2.1_self_forcing/README_TRAIN_DISTILL.md)
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
scripts/wan2.1_flex_forcing/README_TRAIN_ODE.md
+27 -8
View File
@@ -7,9 +7,9 @@ This document provides the complete workflow for **Flex-Forcing Stage 1 — ODE
> ODE regression is the **pre-training stage** of the two-stage Flex-Forcing pipeline. The pipeline is:
>
> 1. **Step 1 — Generate ODE pairs** (`generate_ode_pairs.py`): Use the **bidirectional teacher** (Wan2.1-T2V-1.3B) to perform full multi-step CFG denoising on a list of text prompts. Save the intermediate latents along the ODE trajectory together with the encoded prompt embeddings as `.safetensors` files.
> 2. **Step 2 — Train ODE regression** (`train_ode.py`): Load the generated ODE pairs and train a **flexible-chunking generator** to predict the clean endpoint `x0` at multiple sampled trajectory points, drawing a fresh random frame partition every iteration. The output checkpoint serves as the few-step initialization for the subsequent **Stage 2 DMD distillation** (`train_distill.py`, see [README_TRAIN_FLEX_DMD_DISTILL.md](./README_TRAIN_FLEX_DMD_DISTILL.md)).
> 2. **Step 2 — Train ODE regression** (`train_ode.py`): Load the generated ODE pairs and train a **flexible-chunking generator** to predict the clean endpoint `x0` at multiple sampled trajectory points, drawing a fresh frame partition every iteration from the **same four-arm mixture Stage 2 takes its level 0 from**, so that no layout Stage 2 distills on is one this stage never regressed. The output checkpoint serves as the few-step initialization for the subsequent **Stage 2 DMD distillation** (`train_distill.py`, see [README_TRAIN_FLEX_DMD_DISTILL.md](./README_TRAIN_FLEX_DMD_DISTILL.md)).
>
> Flex-Forcing itself replaces Self-Forcing's scalar `num_frame_per_block` with a **partition of the frame axis** `a = (a_0, ..., a_K)`: attention is bidirectional *inside* a chunk and autoregressive *across* chunks. `[1] * F` recovers pure autoregression, `[F]` recovers pure bidirectional attention, and everything in between is covered by one model. The random partitions are first seen in this stage.
> Flex-Forcing itself replaces Self-Forcing's scalar `num_frame_per_block` with a **partition of the frame axis** `a = (a_0, ..., a_K)`: attention is bidirectional *inside* a chunk and autoregressive *across* chunks. `[1] * F` recovers pure autoregression, `[F]` recovers pure bidirectional attention, and everything in between is covered by one model. Both extremes are drawn here by explicit arms rather than left to chance: a random draw reaches `[F]` at best `1/19` of the time (and never at all while `--flex_chunk_max` sits below the frame count), and `[1] * F` with probability `1/F!`.
---
@@ -208,9 +208,22 @@ After Step 1 completes and `datasets/ode_pairs_output/outputs.json` is generated
For each training sample the script:
1. Loads the saved sparse trajectory (5 points) and prompt embedding from one `.safetensors` file.
2. Draws the frame-axis partition and gives **one timestep per chunk** (`get_timestep_for_ode_flexible`) — this is what makes a chunk internally bidirectional. 90% of iterations draw a **fresh random partition** (chunk sizes `--flex_chunk_min`..`--flex_chunk_max`); the other `UNIFORM_BLOCK_PROB = 0.1` pin it to the uniform `--num_frame_per_block` layout, which a random draw would reach with probability `1/183708` at 21 latent frames and which is exactly what `log_validation` renders.
2. Draws the frame-axis partition and gives **one timestep per chunk** (`get_timestep_for_ode_flexible`) — this is what makes a chunk internally bidirectional.
3. Feeds the noisy latent and the per-frame timestep through the generator, converts the predicted flow into an `x0` prediction, and computes MSE loss against the **clean endpoint** of the trajectory.
Step 2's partition comes from `sample_flex_level0` in `videox_fun/utils/flex_chunking.py` — the **same four-arm mixture Stage 2 takes its level 0 from** (shown for 21 frames with `--num_frame_per_block=3`):
| Arm | Partition drawn | Stage 2 share |
| --- | --- | --- |
| coarse | `[21]` — one chunk over the whole clip, fully bidirectional | 30% |
| uniform | The uniform grid, itself three equiprobable sub-layouts: `[3]*7` / `[1]*21` / `[3]*7` planning followed by `[1]*21` | 30% |
| full_then_blocks | `[21]` (the ladder's level 0) | 30% |
| random | A fresh random partition, chunk sizes `--flex_chunk_min`..`--flex_chunk_max` | 10% |
**Why the draw is shared**: this stage's checkpoint is Stage 2's `--ode_transformer_path` initialisation, so every layout Stage 2 spends iterations on has to be one Stage 1 actually regressed.
Stage 1 keeps only the level-0 partition, so the two ladder-building arms collapse onto `[21]`: at `--flex_pyramid_levels>1` that is `[21]` 60% / `[3]*7` 20% / `[1]*21` 10% / random 10%; at `=1` there is no coarse arm and the other two take 45% each, i.e. `[21]` 45% / `[3]*7` 30% / `[1]*21` 15% / random 10%. `log_validation` pins `[3]*7`, which the uniform arm covers at 20% of iterations instead of the `1/183708` a purely random draw would give it.
### 4.1 Quick Start
The ready-to-use launcher is [train_ode.sh](./train_ode.sh):
@@ -257,6 +270,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--flex_forcing \
--flex_chunk_min=2 \
--flex_chunk_max=10 \
--flex_pyramid_levels=4 \
--num_frame_per_block=3
```
@@ -268,7 +282,9 @@ bash scripts/wan2.1_flex_forcing/train_ode.sh
> 💡 Because the ODE trajectory and prompt embeddings are pre-computed in Step 1, **no VAE / text encoder is invoked during ODE training** — training is fast, memory-efficient, and `train_data_dir` can be left empty when `outputs.json` already contains absolute paths.
> Stage 1 carries `--flex_forcing` and `--flex_chunk_min/max`, but **no** `--flex_pyramid_levels`: it regresses `x0` in a single forward per block, so there is no multi-step schedule to ladder over. The pyramid timestep schedule arrives in Stage 2.
> Stage 1 carries `--flex_forcing`, `--flex_chunk_min/max` **and** `--flex_pyramid_levels`, but builds no ladder: it regresses `x0` in a single forward per chunk, so only level 0 of the shared draw is used. `--flex_pyramid_levels` must still **match Stage 2's**, because it decides whether the coarse (whole-clip) arm exists and therefore what the other three arms weigh — leaving it at the default `1` while Stage 2 runs `4` would cut the `[F]` layout's share of Stage 1 from 60% to 45%, under-training exactly the layout Stage 2 starts most of its iterations from. The pyramid *timestep schedule* itself arrives in Stage 2.
>
> Stage 1 also does **not** pass `--use_kv_cache_training`, so it runs the one-shot block-mask path. That forward supplies neither `kv_cache` nor `clean_x`, which is what arms the K-Projection (§3.3), so `flex_kproj.*` receives **zero gradient** for the whole of Stage 1 and reaches Stage 2 at its identity initialisation. That is the documented cold start rather than a defect — but it does mean Stage 1's steps buy Π nothing. Add the flag if you want Π trained here too; Stage 1's KV branch is already partition-aware and passes one timestep per block, which is what `_first_frame_timestep` assumes.
Output: `output_dir_wan2.1_flex_forcing_ode_regression/`.
@@ -304,9 +320,9 @@ Output: `output_dir_wan2.1_flex_forcing_ode_regression/`.
| Parameter | Description | Example Value |
|-----------|-------------|---------------|
| `--train_sampling_steps` | Size of the scheduler grid that `denoising_step_indices_list` indexes into | 1000 |
| `--denoising_step_indices_list` | Discrete timestep indices used during ODE regression (corresponds to the 5 sparse points sampled in Step 1). Read off the scheduler's own grid rather than a `set_timesteps()` rebuild, so at `shift=8.0` they map to t = `[1000, 960, 888.89, 727.27]` - bit-identical to what `stochastic_sampling_timesteps` feeds validation and inference. The rebuild was off by up to `5.44`, because it ends its grid at a different `sigma_min`. | `1000 750 500 250` |
| `--denoising_step_indices_list` | Discrete timestep indices used during ODE regression (corresponds to the 5 sparse points sampled in Step 1). Read off the scheduler's own grid rather than a `set_timesteps()` rebuild, so at `shift=8.0` they map to t = `[1000, 960, 888.89, 727.27]` - bit-identical to what `stochastic_sampling_timesteps` feeds validation and inference. The rebuild ends its grid at a different `sigma_min`, so indexing it would land away from the t the pipeline actually denoises at. | `1000 750 500 250` |
| `--shift` | Shift for `FlowMatchEulerDiscreteScheduler` — **must match `--shift` used in Step 1** | 8.0 |
| `--num_frame_per_block` | The uniform fallback for the inherited path; the partition `UNIFORM_BLOCK_PROB = 0.1` of iterations pin, so that layout is genuinely trained rather than merely reachable; and the *fixed* partition `log_validation` samples with so checkpoints stay comparable | 3 |
| `--num_frame_per_block` | The uniform fallback for the inherited path; the width the uniform arm of the shared level-0 mixture pins (`[3]*7`, which the uniform arm gives 20% of iterations at `--flex_pyramid_levels>1`, so that layout is genuinely trained rather than merely reachable); the `{1, nfpb}` set the `full_then_blocks` arm draws its block width from; and the *fixed* partition `log_validation` samples with so checkpoints stay comparable | 3 |
| `--independent_first_frame` | First frame is independent (`[1, N, N, ...]` block pattern) | - |
| `--context_noise` | Context noise level (matches downstream distillation config) | 0 |
@@ -315,8 +331,9 @@ Output: `output_dir_wan2.1_flex_forcing_ode_regression/`.
| Parameter | Paper | Description | Default |
|-----------|-------|-------------|---------|
| `--flex_forcing` | §3.1 | Instantiate `WanTransformer3DModel_FlexForcing` and validate with `WanFlexForcingPipeline`. Off = the inherited path, unchanged. | off |
| `--flex_chunk_min` | §3.1 | Smallest chunk drawn per iteration. `1` is legal and lets single-frame (strictly causal) chunks appear mid-clip; `2` is the paper's and keeps the sampler's tail from degenerating into a sliver. | 2 |
| `--flex_chunk_max` | §3.1 | Largest chunk drawn per iteration. Set equal to `--flex_chunk_min` to pin one fixed layout. Unlike stage 2 there is no coarse mixture here, so raising this to the latent frame count is the **only** way stage 1 ever regresses the whole clip as one fully bidirectional chunk (about 5% of draws at 21 latent frames); below the frame count that layout is never drawn. | 10 |
| `--flex_chunk_min` | §5 | Smallest chunk drawn by the *random* arm (10% of iterations). `1` is legal and lets single-frame (strictly causal) chunks appear mid-clip; `2` is the paper's and keeps the sampler's tail from degenerating into a sliver. Unlike Stage 2, nothing here refines a coarse chunk down to 1-frame leaves, so `1` really does add coverage — though the fully causal `[1]*F` layout is a separate arm and is drawn either way. | 2 |
| `--flex_chunk_max` | §5 | Largest chunk drawn by the *random* arm. Set equal to `--flex_chunk_min` to pin that arm to one fixed layout. The whole-clip layout is its own arm, so this knob does not decide whether Stage 1 ever regresses the clip as one fully bidirectional chunk; raising it to the latent frame count only spends random-arm draws on a layout the coarse arm already covers — at 21 latent frames it makes `[21]` 1/19 of that arm and cuts `P([3]*7)` from `1/183708` to `1/1106560`. | 10 |
| `--flex_pyramid_levels` | §3.2 | **Must match Stage 2's.** Stage 1 builds no ladder, but the level-0 mixture is shared and this flag decides whether its coarse arm exists: `>1` gives 30% whole-clip / 30% uniform / 30% `full_then_blocks` / 10% random (the random arm holds a fixed `FLEX_RANDOM_PROB` and the other three split the rest equally), while `=1` divides the coarse arm's share between the two survivors - random stays at 10% - for 45 / 45 / 10. | 1 |
**Validation Parameters (Optional)**:
@@ -368,6 +385,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
--flex_forcing \
--flex_chunk_min=2 \
--flex_chunk_max=10 \
--flex_pyramid_levels=4 \
--num_frame_per_block=3
```
@@ -419,6 +437,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
--flex_forcing \
--flex_chunk_min=2 \
--flex_chunk_max=10 \
--flex_pyramid_levels=4 \
--num_frame_per_block=3
```
scripts/wan2.1_flex_forcing/README_TRAIN_ODE_zh-CN.md
+26 -7
View File
@@ -7,7 +7,7 @@
> ODE 回归是 Flex-Forcing 两阶段流程的 **预训练阶段**，整体流程分两步：
>
> 1. **第一步 — 生成 ODE 对**（`generate_ode_pairs.py`）：使用 **双向教师模型** Wan2.1-T2V-1.3B，对一组文本提示词执行完整的多步 CFG 去噪，将 ODE 轨迹上的中间 latent 与编码后的 prompt embedding 一起保存为 `.safetensors` 文件。
> 2. **第二步 — ODE 回归训练**（`train_ode.py`）：加载第一步生成的 ODE 对，训练一个 **灵活分块生成器**，在轨迹上随机抽样多个噪声等级，预测干净的终点 `x0`；每次迭代都会抽取一个帧划分（九成是全新一轮的随机划分，一成钉成 `--num_frame_per_block` 的均匀块）。训练得到的权重作为 **阶段 2 DMD 蒸馏**（`train_distill.py`，参见 [README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md](./README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md)）的少步初始化。
> 2. **第二步 — ODE 回归训练**（`train_ode.py`）：加载第一步生成的 ODE 对，训练一个 **灵活分块生成器**，在轨迹上随机抽样多个噪声等级，预测干净的终点 `x0`；每次迭代都会抽取一个帧划分，用的是**阶段 2 取 level 0 时的同一套四档混合**，因此阶段 2 要蒸馏的布局没有一种是本阶段没回归过的。训练得到的权重作为 **阶段 2 DMD 蒸馏**（`train_distill.py`，参见 [README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md](./README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md)）的少步初始化。
>
> Flex-Forcing 本身把 Self-Forcing 的标量 `num_frame_per_block` 换成**帧轴上的一个划分** `a = (a_0, ..., a_K)`：chunk（分块）内部是双向注意力，chunk 之间是自回归。`[1] * F` 退化为纯自回归，`[F]` 退化为纯双向注意力，中间的一切形态都由同一个模型覆盖。随机分块正是在本阶段首次引入的。
@@ -208,9 +208,22 @@ accelerate launch --multi_gpu --num_processes=8 --mixed_precision="bf16" \
每个训练样本上，训练脚本会：
1. 从一个 `.safetensors` 文件中加载稀疏的 5 点轨迹与 prompt embedding；
2. 抽取帧划分，并给**每个 chunk 配一个统一的时间步**（`get_timestep_for_ode_flexible`）——这正是"chunk 内部双向"的来源；其中 90% 的迭代抽取一个**全新的随机划分**（chunk 大小在 `--flex_chunk_min`..`--flex_chunk_max` 之间），另外 `UNIFORM_BLOCK_PROB = 0.1` 钉成 `--num_frame_per_block` 的均匀划分——21 个 latent 帧下随机抽取碰到 `[3]*7` 的概率只有 `1/183708`，而 `log_validation` 渲染的正是它；
2. 抽取帧划分，并给**每个 chunk 配一个统一的时间步**（`get_timestep_for_ode_flexible`）——这正是"chunk 内部双向"的来源；
3. 将带噪 latent 与逐帧时间步送入生成器，把输出的 flow 转换为 `x0` 预测，与轨迹的**干净终点**计算 MSE 损失。
第 2 步的划分来自 `videox_fun/utils/flex_chunking.py` 的 `sample_flex_level0`，与阶段 2 取 level 0 时用的是**同一套四档混合**（以 21 帧、`--num_frame_per_block=3` 为例）：
| 档 | 抽出的划分 | 阶段 2 占比 |
| --- | --- | --- |
| coarse | `[21]` —— 整段一个 chunk，完全双向 | 30% |
| uniform | 均匀网格，内部再三等分：`[3]*7` / `[1]*21` / `[3]*7` 规划后接 `[1]*21` | 30% |
| full_then_blocks | `[21]`（阶梯的 level 0） | 30% |
| random | 全新随机划分，chunk 大小在 `--flex_chunk_min`..`--flex_chunk_max` 之间 | 10% |
**为什么共用这次抽取**：本阶段的 checkpoint 就是阶段 2 的 `--ode_transformer_path` 初始化，所以阶段 2 会花迭代去跑的每一种布局，都必须是阶段 1 真正回归过的。
阶段 1 只取 level-0 那一层，两个搭阶梯的档因此都塌到 `[21]` 上：`--flex_pyramid_levels>1` 时 `[21]` 60% / `[3]*7` 20% / `[1]*21` 10% / 随机 10%；`=1` 时没有 coarse 档，其余两档各 45%，即 `[21]` 45% / `[3]*7` 30% / `[1]*21` 15% / 随机 10%。`log_validation` 固定用 `[3]*7`，均匀档给了它 20% 的迭代，而纯随机抽取只有 `1/183708`。
### 4.1 快速开始
直接复用启动脚本 [train_ode.sh](./train_ode.sh)：
@@ -257,6 +270,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--flex_forcing \
--flex_chunk_min=2 \
--flex_chunk_max=10 \
--flex_pyramid_levels=4 \
--num_frame_per_block=3
```
@@ -268,7 +282,9 @@ bash scripts/wan2.1_flex_forcing/train_ode.sh
> 💡 因为 ODE 轨迹与 prompt embedding 已在第一步预先计算完毕，**ODE 训练阶段不会再调用 VAE / 文本编码器**，训练速度快、显存占用低。当 `outputs.json` 中已经使用绝对路径时，`train_data_dir` 可以留空。
> 阶段 1 带 `--flex_forcing` 与 `--flex_chunk_min/max`，但**不带** `--flex_pyramid_levels`：金字塔是**多步**调度，而 ODE 阶段对每个 block 只做一次前向来回归 `x0`，没有可以逐级细化的调度。金字塔时间步在阶段 2 才出现。
> 阶段 1 带 `--flex_forcing`、`--flex_chunk_min/max` **和** `--flex_pyramid_levels`，但**不搭阶梯**：每个 chunk 只做一次前向来回归 `x0`，所以共用抽取里只有 level 0 被用到。`--flex_pyramid_levels` 仍必须**与阶段 2 一致**，因为它决定 coarse（整段）档是否存在、也就决定了其余三档各占多少 —— 阶段 2 跑 `4` 而这里留默认 `1`，会把 `[F]` 布局在阶段 1 的占比从 60% 砍到 45%，而 `[F]` 恰恰是阶段 2 大多数迭代起步的布局，于是被训得不足。阶梯的*时间步排布*本身要到阶段 2 才出现。
>
> 阶段 1 也**不传** `--use_kv_cache_training`，因此走的是 one-shot block-mask 路径。该次前向既不给 `kv_cache` 也不给 `clean_x`，而这两者正是启用 K-Projection（§3.3）的条件，所以 `flex_kproj.*` 在整个阶段 1 **梯度恒为 0**，交到阶段 2 时仍停在恒等初始化。这是文档写明的冷启动而非缺陷 —— 但确实意味着阶段 1 这些步数对 Π 毫无贡献。若希望 Π 也在本阶段被训起来，加上该 flag 即可：阶段 1 的 KV 分支已经是划分感知的，且按 block 传单一时间步，正好符合 `_first_frame_timestep` 的假设。
输出目录：`output_dir_wan2.1_flex_forcing_ode_regression/`。
@@ -304,9 +320,9 @@ bash scripts/wan2.1_flex_forcing/train_ode.sh
| 参数 | 说明 | 示例值 |
|------|------|-------|
| `--train_sampling_steps` | `denoising_step_indices_list` 所索引的调度器网格大小 | 1000 |
| `--denoising_step_indices_list` | ODE 回归使用的离散时间步索引（与第一步抽样的 5 个稀疏轨迹点对应）。直接读调度器自带的网格，而不是 `set_timesteps()` 重建的那张，所以 `shift=8.0` 下映射为 t = `[1000, 960, 888.89, 727.27]` —— 与 `stochastic_sampling_timesteps` 喂给校验和推理的值逐位一致。重建的那张最多偏 `5.44`，因为它的网格末端 `sigma_min` 不同。 | `1000 750 500 250` |
| `--denoising_step_indices_list` | ODE 回归使用的离散时间步索引（与第一步抽样的 5 个稀疏轨迹点对应）。直接读调度器自带的网格，而不是 `set_timesteps()` 重建的那张，所以 `shift=8.0` 下映射为 t = `[1000, 960, 888.89, 727.27]` —— 与 `stochastic_sampling_timesteps` 喂给校验和推理的值逐位一致。重建的那张网格末端 `sigma_min` 不同，索引它会落在 pipeline 实际去噪的 t 之外。 | `1000 750 500 250` |
| `--shift` | `FlowMatchEulerDiscreteScheduler` 的 shift —— **必须与第一步生成时使用的 `--shift` 一致** | 8.0 |
| `--num_frame_per_block` | 继承路径的均匀 fallback；`UNIFORM_BLOCK_PROB = 0.1` 的迭代会把划分钉成它，使该布局是真正被训练而不只是理论上可达；也是 `log_validation` 采样验证时用的固定划分，保证 checkpoint 之间可比较 | 3 |
| `--num_frame_per_block` | 继承路径的均匀 fallback；共用 level-0 混合里均匀档所钉的块宽（`[3]*7`，`--flex_pyramid_levels>1` 时均匀档给它 20% 的迭代，使该布局是真正被训练而不只是理论上可达）；`full_then_blocks` 档抽块宽时的 `{1, nfpb}` 取值集合；也是 `log_validation` 采样验证时用的固定划分，保证 checkpoint 之间可比较 | 3 |
| `--independent_first_frame` | 第一帧是否独立（`[1, N, N, ...]` 块模式） | - |
| `--context_noise` | 上下文噪声等级（与下游蒸馏配置匹配） | 0 |
@@ -315,8 +331,9 @@ bash scripts/wan2.1_flex_forcing/train_ode.sh
| 参数 | 论文 | 说明 | 默认值 |
|-----------|-------|-------------|---------|
| `--flex_forcing` | §3.1 | 实例化 `WanTransformer3DModel_FlexForcing` 并用 `WanFlexForcingPipeline` 校验。关闭 = 走继承的原路径，行为不变。 | 关 |
| `--flex_chunk_min` | §3.1 | 每次迭代抽取的最小 chunk 大小。`1` 合法，会让单帧（严格因果）chunk 出现在片段中间；`2` 是论文取值，也能避免采样器把尾部退化成 1 帧的碎块。 | 2 |
| `--flex_chunk_max` | §3.1 | 每次迭代抽取的最大 chunk 大小。设为与 `--flex_chunk_min` 相等即固定一种布局。与阶段 2 不同，这里没有粗粒度混合机制，所以把这个值抬到 latent 帧数是阶段 1 **唯一**能把“整段作为一个完全双向 chunk”回归出来的途径（21 个 latent 帧时约占 5% 的抽取）；低于帧数则该布局永远不会被抽到。 | 10 |
| `--flex_chunk_min` | §5 | **随机档**（占 10% 迭代）抽取的最小 chunk 大小。`1` 合法，会让单帧（严格因果）chunk 出现在片段中间；`2` 是论文取值，也能避免采样器把尾部退化成 1 帧的碎块。与阶段 2 不同，这里没有任何东西会把粗 chunk 细化成 1 帧叶子，所以 `1` 确实能增加覆盖——不过完全因果的 `[1]*F` 布局本身是独立一档，无论怎么设都会被抽到。 | 2 |
| `--flex_chunk_max` | §5 | **随机档**抽取的最大 chunk 大小。设为与 `--flex_chunk_min` 相等即把该档固定成一种布局。整段布局本身就是独立的一档，所以这个旋钮并不决定阶段 1 能否把“整段作为一个完全双向 chunk”回归出来；把它抬到 latent 帧数只是让随机档把抽取花在 coarse 档已经覆盖的布局上——21 个 latent 帧时它会让 `[21]` 占该档的 1/19，并把 `P([3]*7)` 从 `1/183708` 摊薄到 `1/1106560`。 | 10 |
| `--flex_pyramid_levels` | §3.2 | **必须与阶段 2 一致。**阶段 1 不搭阶梯，但 level-0 混合是共用的，而这个 flag 决定 coarse 档是否存在：`>1` 时为 30% 整段 / 30% 均匀 / 30% `full_then_blocks` / 10% 随机（随机档占固定的 `FLEX_RANDOM_PROB`，其余三档平分剩下的 90%）；`=1` 时 coarse 档的那份由两个幸存档平分——随机档恒为 10%——变成 45 / 45 / 10。 | 1 |
**验证参数（可选）**：
@@ -368,6 +385,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
--flex_forcing \
--flex_chunk_min=2 \
--flex_chunk_max=10 \
--flex_pyramid_levels=4 \
--num_frame_per_block=3
```
@@ -419,6 +437,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
--flex_forcing \
--flex_chunk_min=2 \
--flex_chunk_max=10 \
--flex_pyramid_levels=4 \
--num_frame_per_block=3
```
scripts/wan2.1_flex_forcing/train_distill.py
+112 -169
View File
@@ -77,16 +77,12 @@ from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
WanTransformer3DModel,
WanTransformer3DModel_FlexForcing,
WanTransformer3DModel_SelfForcing)
from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
WanFlexForcingPipeline,
WanSelfForcingPipeline)
from videox_fun.pipeline import (WanFlexForcingPipeline, WanI2VPipeline,
WanPipeline, WanSelfForcingPipeline)
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.flex_chunking import (broadcast_chunk_sizes,
build_full_then_blocks_partitions,
build_pyramid_partitions,
from videox_fun.utils.flex_chunking import (build_pyramid_partitions,
chunk_boundaries,
sample_flexible_chunks,
uniform_chunks,
sample_flex_level0,
validate_nested_partitions)
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
@@ -299,6 +295,7 @@ def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer
# the layout inference builds - otherwise the samples say
# nothing about the model being trained.
flex_kwargs = {}
val_stochastic_sampling = not args.flow_euler_rollout
if args.flex_forcing:
if args.flex_pyramid_levels > 1:
# 3.2: exactly what predict_t2v.py passes. `chunk_spec=None`
@@ -311,14 +308,6 @@ def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer
min_num_frame_per_block=args.flex_min_num_frame_per_block,
)
else:
# No ladder in training, and with no pyramid the
# whole-clip band is empty - `[21]` is never drawn at
# `flex_chunk_max` below the frame count, so "bidir"
# would validate out of distribution too. The uniform
# block layout is the one band that is both fixed
# across checkpoints and actually trained:
# FLEX_ARM_PROB of iterations use it, against
# ~1.3% for the most common random partition.
flex_kwargs = dict(
chunk_spec=args.num_frame_per_block,
denoise_mode="fixed",
@@ -336,6 +325,7 @@ def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer
num_frame_per_block = args.num_frame_per_block,
independent_first_frame = args.independent_first_frame,
context_noise = args.context_noise,
stochastic_sampling = val_stochastic_sampling,
**flex_kwargs,
).videos
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
@@ -379,21 +369,20 @@ def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=Non
return torch.clip(t.to(torch.int32), low, high - 1)
# Each training iteration draws its level-0 layout from four equally likely arms
# (FLEX_ARM_PROB = 25% of iterations each): the whole-clip planning chunk (3.2's
# coarse step, `[[F], ...]` at inference), the launcher's own uniform
# `num_frame_per_block`, the fixed "first step full, then block-major" ladder
# (`denoise_mode="full_then_blocks"`), and a random 3.1 partition. Equal shares
# keep the coarse, causal and every in-between layout covered while giving the
# full_then_blocks ladder enough iterations to train the block-major-at-high-noise
# trajectory the binary pyramid never reaches. Deliberately not a CLI flag - it is
# a property of the schedule mixture, not a knob the launcher keeps in sync.
FLEX_ARM_PROB = 0.25
# The level-0 four-arm mixture (30% of iterations each for the whole-clip
# planning chunk, the launcher's own uniform `num_frame_per_block` and the fixed
# "first step full, then block-major" ladder, plus 10% for a random §3.3/§5
# partition - see `FLEX_RANDOM_PROB`) now lives in
# `videox_fun.utils.flex_chunking.sample_flex_level0`, because `train_ode.py`
# draws from the SAME law: its checkpoint is this stage's `--ode_transformer_path`
# initialisation, so a layout this stage spends most of its iterations on must
# not be one stage 1 never regressed. What stays here is only the ladder built on
# top of the draw.
# How many steps of a launch report the partition they drew. Counted from this
# process rather than from `global_step`, so a resumed run still gets its own
# look at the layout. Two lines come out per step (the generator and the critic
# each draw one), which is enough to watch the three arms show up in their
# each draw one), which is enough to watch the arms show up in their
# stated shares; after that they are only noise in a multi-day run.
FLEX_LAYOUT_LOG_STEPS = 20
@@ -402,134 +391,74 @@ def sample_flex_partitions(args, num_frames, num_denoising_steps, torch_rng,
device, verbose=False):
"""Draw the Flex-Forcing partition ladder for one training iteration.
arXiv 2607.03509 3.1 asks for a *random* partition per rollout (chunk sizes
2..10) so a single model covers the whole causal-to-bidirectional spectrum,
and 3.2 then nests that partition into a coarse-to-fine ladder with one level
per denoising step. Returns ``None`` when Flex-Forcing is off, which leaves
the inherited uniform ``num_frame_per_block`` masks untouched.
arXiv 2607.03509 §3.3 asks for a *random* partition per rollout so a single
model covers the whole causal-to-bidirectional spectrum, §5 fixes that draw's
chunk sizes at 2..10, and §3.2 then nests the partition into a coarse-to-fine
ladder with one level per denoising step. Returns ``None`` when Flex-Forcing
is off, which leaves the inherited uniform ``num_frame_per_block`` masks
untouched.
The level-0 layout is a four-way mixture decided by a single uniform draw,
each arm taking ``FLEX_ARM_PROB`` (25%) of iterations: one chunk over the
whole clip (inference's coarse planning step), a single-level uniform block
grid, the fixed "first step full, then block-major" ladder
(`denoise_mode="full_then_blocks"`), and a random 2..10 partition. The coarse
and random arms are refined into the same binary pyramid, so a single set of
weights covers `denoise_mode="pyramid"` whether or not the caller pins
`chunk_spec`. The uniform and full_then_blocks arms bypass that refinement:
uniform emits a pure-Self-Forcing ladder in one of three layouts drawn from
The level-0 layout is a four-way mixture decided by a single uniform draw:
one chunk over the whole clip (inference's coarse planning step), a
single-level uniform block grid, the fixed "first step full, then
block-major" ladder (`denoise_mode="full_then_blocks"`), and a random 2..10
partition. The random arm takes ``FLEX_RANDOM_PROB`` (10%) and the other
three split what is left equally, i.e. 30% each. The coarse and random arms
are refined into the same binary pyramid, so a single set of weights covers
`denoise_mode="pyramid"` whether or not the caller pins `chunk_spec`. The
uniform and full_then_blocks arms bypass that refinement: uniform emits a
pure-Self-Forcing ladder in one of three layouts drawn from
`num_frame_per_block` (`[3x7]`) and fully causal (`[1x21]`) -- `[3x7]` alone,
`[1x21]` alone, or `[3x7]` planning followed by `[1x21]` for every remaining
step; full_then_blocks emits its two-level ladder whole, drawing the same
block width. With no pyramid the coarse arm is empty (a `[F]` level 0 would be
a degenerate single-chunk rollout), so the split is 25% uniform / 25%
full_then_blocks / 50% random.
step, 10% of iterations each; full_then_blocks emits its two-level ladder
whole, drawing the same block width. With no pyramid the coarse arm is empty
(a `[F]` level 0 would be a degenerate single-chunk rollout), so the two
survivors widen and the split becomes 45% uniform / 45% full_then_blocks /
10% random.
The draw itself is :func:`videox_fun.utils.flex_chunking.sample_flex_level0`,
shared with ``train_ode.py`` so that stage 1 - whose checkpoint this stage
loads via ``--ode_transformer_path`` - regresses the same level-0 law it is
then distilled on. What this function adds is only the *ladder*: how the drawn
layout nests into one level per denoising step.
"""
if not args.flex_forcing:
return None
# The arm selector must be IDENTICAL on every rank, not just its broadcast
# count. Only the coarse/random arms funnel through the shared refine tail
# below; the uniform and full_then_blocks arms each set the ladder directly
# from a different tensor, so a per-rank `u` would leave ranks in different
# arms building different ladders
# -> different walk forward counts -> FSDP all-gather desync (NCCL hang).
# Broadcast rank 0's draw so all ranks take the same arm; the arm's existing
# single broadcast then reconciles its within-arm randomness (the random base,
# the block width) to rank 0, making every ladder byte-identical.
u = torch.rand((), generator=torch_rng, device=device)
if dist.is_available() and dist.is_initialized():
dist.broadcast(u.reshape(1), src=0)
u = u.item()
pyramid = args.flex_pyramid_levels > 1
# Four equally likely arms; the coarse (whole-clip planning) one only when a
# pyramid is on, else its 25% folds into the random arm.
coarse_hi = FLEX_ARM_PROB if pyramid else 0.0
uniform_hi = coarse_hi + FLEX_ARM_PROB
ftb_hi = uniform_hi + FLEX_ARM_PROB
ladder = None
if u < coarse_hi:
# Coarse end of 3.1: reuse the ladder builder's own level-0 rule so the
# `independent_first_frame` handling cannot drift from inference's.
base = build_pyramid_partitions(
num_frames, num_levels=1, base_chunks=None,
independent_first_frame=args.independent_first_frame)[0]
arm = "whole clip, the 3.2 coarse planning layout"
elif u < uniform_hi:
# Pure Self-Forcing (no binary-pyramid refinement), three layouts drawn
# from `num_frame_per_block` (-> [3x7]) and fully causal (1 frame/block,
# -> [1x21]): [3x7] alone, [1x21] alone, and [3x7] planning followed by
# [1x21] for every remaining step (the coarse-block -> per-frame descent
# the pyramid never samples because it halves through [2,1x14] first).
# The three change the ladder DEPTH (hence the walk's forward count), so
# the choice must be rank-consistent; it is free here because `u` was
# already broadcast to rank 0, so slicing its position within the uniform
# band gives every rank the same sub-arm with no extra collective. Like
# every arm this issues EXACTLY ONE broadcast (the level-0 layout; the
# trailing [1x21] levels are deterministic from `num_frames`). Setting
# `ladder` here skips the refine tail. `uniform_chunks`, not
# `normalize_chunk_spec`: the latter encodes `independent_first_frame` as
# a leading 1 and would silently disagree with the bands around it.
band = (u - coarse_hi) / FLEX_ARM_PROB # [0, 1) in the uniform arm
sub = min(int(band * 3), 2) # 0 / 1 / 2, equiprobable
nfpb = max(1, int(args.num_frame_per_block))
iff = args.independent_first_frame
causal = uniform_chunks(num_frames, 1, independent_first_frame=iff)
if sub == 0:
base = broadcast_chunk_sizes(causal, device=device)
ladder = [base]
arm = "uniform 1-frame blocks, single level (fully causal)"
elif sub == 1:
base = broadcast_chunk_sizes(
uniform_chunks(num_frames, nfpb, independent_first_frame=iff),
device=device)
ladder = [base]
arm = f"uniform {nfpb}-frame blocks, single level"
else:
base = broadcast_chunk_sizes(
uniform_chunks(num_frames, nfpb, independent_first_frame=iff),
device=device)
ladder = ([base] + [list(causal)
for _ in range(max(0, num_denoising_steps - 1))])
arm = (f"uniform {nfpb}-frame plan -> fully causal [1x" f"{num_frames}] "
"for the remaining steps")
elif u < ftb_hi:
# "First step full, every later step block-major": a fixed two-level
# ladder that does NOT go through the binary pyramid refinement, so the
# block-major level is reached at high noise (the 2nd step) instead of as
# a deep refinement - the one trajectory the pyramid arms never sample.
# The block-major width is drawn between fully causal (1 frame / block)
# and the launcher's `num_frame_per_block`, so both the tight-AR and the
# coarse-block refinement of the whole-clip plan get trained. The block
# width is a per-rank draw, so it has to be reconciled with EXACTLY ONE
# broadcast - the same collective count every other arm issues (they each
# broadcast their single base) - or the ranks desync and NCCL deadlocks.
# Level 0 is just [F], identical on every rank and needing no sync, so we
# broadcast only the block-major level; rank 0's draw wins and every rank
# rebuilds the same two-level ladder locally.
block_choices = sorted({1, max(1, int(args.num_frame_per_block))})
pick = int(torch.rand((), generator=torch_rng, device=device).item()
* len(block_choices))
block = block_choices[min(pick, len(block_choices) - 1)]
block = min(block, int(num_frames))
ftb = build_full_then_blocks_partitions(
num_frames, block,
independent_first_frame=args.independent_first_frame)
ftb[1] = broadcast_chunk_sizes(ftb[1], device=device)
ladder = ftb[:num_denoising_steps]
arm = (f"full clip -> uniform {max(ftb[1])}-frame blocks, "
"the full_then_blocks ladder")
draw = sample_flex_level0(
num_frames,
num_frame_per_block=args.num_frame_per_block,
chunk_min=args.flex_chunk_min,
chunk_max=args.flex_chunk_max,
pyramid=pyramid,
generator=torch_rng,
device=device,
independent_first_frame=args.independent_first_frame)
base = draw["sizes"]
arm = draw["arm_text"]
if draw["arm"] in ("uniform_causal", "uniform_blocks"):
# Pure Self-Forcing grids: a single level, and they deliberately bypass
# the binary refinement below - refining `[3x7]` would halve it through
# `[2, 1x14]`, which is the pyramid's own trajectory, not this arm's.
ladder = [base]
elif draw["arm"] == "uniform_plan_then_causal":
# `[3x7]` planning, then fully causal for every remaining step: the
# coarse-block -> per-frame descent the pyramid never samples. The
# trailing `[1x21]` levels are deterministic from `num_frames`, so the
# arm's single broadcast (already issued) covered everything random.
ladder = ([base] + [list(draw["causal"])
for _ in range(max(0, num_denoising_steps - 1))])
elif draw["arm"] == "full_then_blocks":
# Two levels, taken whole. The walk clamps to the finest level for every
# later step, so a ladder shorter than the denoising schedule is the
# intended shape here, not wasted steps.
ladder = [list(base), list(draw["block_level"])][:num_denoising_steps]
else:
base = sample_flexible_chunks(
num_frames, min_chunk=args.flex_chunk_min, max_chunk=args.flex_chunk_max,
generator=torch_rng, device=device,
independent_first_frame=args.independent_first_frame)
arm = f"random {args.flex_chunk_min}..{args.flex_chunk_max}-frame blocks, the 3.1 spectrum"
if ladder is None:
# Every rank has to train the same layout: the FlexAttention mask, and the
# `coarse` / `random`: the two arms the binary pyramid refines. Every rank
# has to train the same layout - the FlexAttention mask, and the
# `num_frame_per_block` derived from it, must agree across the SP/FSDP
# group. Only the coarse and random arms reach here; the uniform and
# full_then_blocks arms above already set `ladder` (each with its single
# broadcast), so they skip this refine path.
base = broadcast_chunk_sizes(base, device=device)
# group - which `sample_flex_level0` already guaranteed by broadcasting
# rank 0's draw.
ladder = [base]
if pyramid:
ladder = build_pyramid_partitions(
@@ -601,16 +530,24 @@ def flex_self_context_enabled(args, clean_x, final_step_index):
"""Whether 3.3's K-Projection gets a self-generated clean context this round.
This is a block-mask-path concern only, and this flag is inert under
``--use_kv_cache_training``. The KV-cache rollout passes no ``flex_state``,
so ``_flex_project`` short-circuits and :math:`\\Pi` stays at its identity
initialisation there - which matches KV inference, which never supplies a
timestep either, so the two sides agree rather than one of them drifting.
The block-mask path is the one that exercises :math:`\\Pi`, and it builds no
KV cache, so with prompt-only data - ``clean_latents is None``, which the
prompt-only branch sets unconditionally - no clean context exists anywhere
in the sequence, :math:`\\Pi_{t\\leftarrow 0}` has nothing to project and
never enters the autograd graph: its parameters sit in an optimizer group
and still stay at the identity initialisation for the entire run.
``--use_kv_cache_training`` - but *not* because :math:`\\Pi` is idle there.
:class:`WanTransformer3DModel_FlexForcing` arms :math:`\\Pi` by itself
whenever a projection was built and the forward has a clean context to read
(``kv_cache is not None or clean_x is not None``), taking the noise level
from ``t``; a caller that passes no ``flex_state`` still projects. So
:math:`\\Pi` **is** invoked on the KV-cache path, in training and in
inference alike, and the exit step that carries gradient gives it real
gradient. All this flag does *not* do there is supply a ``clean_x`` half:
the KV rollout already commits its own clean context into the cache ("feed
denoised_pred directly"), so it has no use for one and the switch has no
lever to pull.
On the block-mask path the switch is what decides whether :math:`\\Pi` has
anything to project at all: that path builds no KV cache, so with
prompt-only data - ``clean_latents is None``, which the prompt-only branch
sets unconditionally - no clean context exists anywhere in the sequence,
:math:`\\Pi_{t\\leftarrow 0}` has nothing to project and never enters the
autograd graph: its parameters sit in an optimizer group and still stay at
the identity initialisation for the entire run.
The denoising loop already re-rolls the model under ``no_grad`` and converts
every non-final step to an x0 prediction, so the previous step's prediction
@@ -1137,21 +1074,23 @@ def parse_args():
"--flex_chunk_min",
type=int,
default=2,
help="Smallest chunk drawn when sampling a partition (paper 3.1 uses 2). "
"1 is allowed but adds no coverage: refining a 2-frame chunk already "
"yields 1-frame leaves."
help="Smallest chunk drawn by the level-0 mixture's *random* arm (paper "
"§5 uses 2). 1 is allowed but adds no coverage: refining a 2-frame "
"chunk already yields 1-frame leaves."
)
parser.add_argument(
"--flex_chunk_max",
type=int,
default=10,
help="Largest chunk drawn when sampling a partition (paper 3.1 uses 10; "
"setting it equal to --flex_chunk_min pins one fixed layout). Keep it "
help="Largest chunk drawn by the level-0 mixture's *random* arm (paper §5 "
"uses 10; setting it equal to --flex_chunk_min pins that arm to one "
"fixed layout). Keep it "
"below the latent frame count: the whole-clip layout is already "
"covered by the pyramid mixture, so raising max to the frame count "
"only spends draws on a layout you already have and thins out 3.1 - "
"at 21 latent frames max=21 leaves 3649 distinct partitions versus "
"4882 at max=10."
"covered by the coarse arm, so raising max to the frame count spends "
"draws on a layout you already have and thins every other one - at 21 "
"latent frames exhaustive enumeration gives 6441 distinct partitions "
"at max=10 and 6765 at max=21, but max=21 makes [21] alone 1/19 of "
"draws and cuts P([3]*7) from 1/183708 to 1/1106560."
)
parser.add_argument(
"--flex_pyramid_levels",
@@ -1178,10 +1117,14 @@ def parse_args():
"block-mask path trains on prompts alone: it builds no KV cache, so "
"without this there is no clean context at all and the K-Projection "
"never enters the autograd graph. Doubles the attention sequence. "
"Inert under --use_kv_cache_training: that path passes no flex_state, "
"so the K-Projection is never invoked on either side and stays at its "
"identity init. Also ignored when real --use_teacher_forcing data is "
"present, and with --flex_forcing off."
"Inert under --use_kv_cache_training for a different reason than it "
"looks like: that path already commits its own clean context into the "
"KV cache, so it needs no clean_x half. The K-Projection itself IS "
"active there - the Flex backbone arms it automatically whenever a "
"projection was built and kv_cache or clean_x is present - so this "
"flag is not evidence that the projection sits untrained under "
"--use_kv_cache_training. Also ignored when real --use_teacher_forcing "
"data is present, and with --flex_forcing off."
)
parser.add_argument(
"--independent_first_frame",
scripts/wan2.1_flex_forcing/train_ode.py
+112 -49
View File
@@ -63,11 +63,7 @@ from videox_fun.models import (AutoencoderKLWan, WanT5EncoderModel,
WanTransformer3DModel_FlexForcing,
WanTransformer3DModel_SelfForcing)
from videox_fun.pipeline import WanFlexForcingPipeline, WanSelfForcingPipeline
from videox_fun.utils.flex_chunking import (UNIFORM_BLOCK_PROB,
broadcast_chunk_sizes,
chunk_boundaries,
sample_flexible_chunks,
uniform_chunks)
from videox_fun.utils.flex_chunking import chunk_boundaries, sample_flex_level0
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
from videox_fun.utils.utils import save_videos_grid
@@ -234,12 +230,27 @@ def get_timestep_for_ode(
def get_timestep_for_ode_flexible(min_timestep, max_timestep, batch_size,
chunk_sizes, device, generator=None):
chunk_sizes, device, generator=None,
clean_index=None):
"""`get_timestep_for_ode` for a Flex-Forcing partition (arXiv 2607.03509 3.1).
Same contract - every frame of a chunk shares one timestep index - but the
chunk widths come from the partition instead of a single scalar, so the
reshape trick of the uniform version does not apply.
``clean_index`` is the index whose entry in ``train_denoising_step_list`` is
t=0, or ``None`` when the schedule has no clean entry. When given, each
sample is guaranteed at least one chunk that did *not* draw it. An all-clean
draw makes the loss mask ``(timestep != 0)`` empty, and the caller's
fallback ``mse(pred_x0, target)`` is then *exactly* zero rather than merely
small: those frames gather the clean endpoint itself and sigma is pinned to
0, so ``pred_x0 == target`` bitwise. The iteration contributes no gradient
at all. The probability is ``len(train_denoising_step_list) ** -K`` in the
number of chunks K, so it is dominated by the coarsest layouts: the
whole-clip ``[F]`` partition is a single chunk and on its own carries
``1 / len(train_denoising_step_list)`` of iterations. A single chunk also
has no later block to serve as clean context *for*, which is the only reason
the clean index exists, so redrawing it loses nothing.
"""
num_frames = sum(int(c) for c in chunk_sizes)
per_chunk = torch.randint(
@@ -248,6 +259,21 @@ def get_timestep_for_ode_flexible(min_timestep, max_timestep, batch_size,
device=device, dtype=torch.long,
generator=generator,
)
if clean_index is not None and per_chunk.shape[1] > 0 \
and max_timestep - min_timestep > 1:
all_clean = (per_chunk == clean_index).all(dim=1)
if bool(all_clean.any()):
rows = all_clean.nonzero(as_tuple=True)[0]
# One chunk per affected sample, redrawn uniformly over every index
# *except* the clean one: draw from the shrunken range, then shift
# past the hole. Works for any clean_index, not just the last one.
pick = torch.randint(
0, per_chunk.shape[1], [rows.numel()],
device=device, dtype=torch.long, generator=generator)
redraw = torch.randint(
min_timestep, max_timestep - 1, [rows.numel()],
device=device, dtype=torch.long, generator=generator)
per_chunk[rows, pick] = redraw + (redraw >= clean_index).long()
timestep = torch.empty(
batch_size, num_frames, device=device, dtype=per_chunk.dtype)
for chunk_idx, (start_frame, end_frame) in enumerate(chunk_boundaries(chunk_sizes)):
@@ -256,38 +282,48 @@ def get_timestep_for_ode_flexible(min_timestep, max_timestep, batch_size,
def sample_flex_partition(args, num_frames, torch_rng, device):
"""Draw this iteration's Flex-Forcing partition (3.1/3.3).
"""Draw this iteration's Flex-Forcing partition (§3.1 / §3.3).
Returns ``None`` when Flex-Forcing is off, so every caller keeps using the
uniform ``num_frame_per_block`` layout. 3.2's pyramid is a *multi-step*
schedule and the ODE stage regresses x0 in a single forward per block, so
there is nothing to ladder over here.
uniform ``num_frame_per_block`` layout.
``UNIFORM_BLOCK_PROB`` of iterations pin the launcher's own uniform
``num_frame_per_block`` layout instead of drawing a random one. Random draws
would almost never produce it - the uniform ``[3] * 7`` at 21 latent frames
has probability 1/183708, i.e. 0.11 expected hits in 20000 iterations - yet
`log_validation` renders exactly that partition, so without this band stage
1's samples would come from a layout it never regressed.
The draw is :func:`videox_fun.utils.flex_chunking.sample_flex_level0` - the
*same* four-arm mixture stage 2 uses for its level 0, random arm at
``FLEX_RANDOM_PROB`` (10%) and the three deterministic arms at 30% each.
Sharing it is the whole point: this stage's checkpoint is stage 2's
``--ode_transformer_path`` initialisation, so every layout stage 2 spends
iterations on has to be one stage 1 actually regressed. The two-band sampler
this replaced drew 10% uniform / 90% random - coincidentally the mirror image
of today's split, and beside the point, since what made it unusable was its
chunk support being exactly ``{--flex_chunk_min .. --flex_chunk_max}``: the
whole-clip ``[F]`` layout that 60% of stage 2's level-0 iterations start from
(at ``--flex_pyramid_levels=4``) had probability *zero* there (``max_chunk``
below the frame count makes it unreachable, and at most 1/19 of random draws
even with ``max_chunk`` raised to it), and so did every 1-frame chunk -
``[1] * F`` is ``1/F!`` under a random draw, i.e. unreachable at any setting
of these two knobs.
§3.2's pyramid is a *multi-step* schedule and this stage regresses x0 in a
single forward per chunk, so there is no ladder to build here and only
level 0 is used. ``--flex_pyramid_levels`` is still read, and must match
stage 2's, because it decides whether the coarse arm exists and therefore
how the other arms split the 90% the random one leaves behind.
``sample_flex_level0`` broadcasts rank 0's draw, so every rank regresses the
same layout; otherwise the FlexAttention mask and the per-chunk timesteps
would disagree across the group.
"""
if not args.flex_forcing:
return None
if torch.rand(
(), generator=torch_rng, device=device).item() < UNIFORM_BLOCK_PROB:
# `uniform_chunks` rather than `normalize_chunk_spec`: the latter takes no
# `independent_first_frame` argument (it encodes that as a leading 1), so
# it would silently drop the leading 1 the random branch below adds.
sizes = uniform_chunks(
num_frames, args.num_frame_per_block,
independent_first_frame=args.independent_first_frame)
else:
sizes = sample_flexible_chunks(
num_frames, min_chunk=args.flex_chunk_min, max_chunk=args.flex_chunk_max,
generator=torch_rng, device=device,
independent_first_frame=args.independent_first_frame)
# Every rank must regress the same layout, otherwise the FlexAttention mask
# and the per-chunk timesteps disagree across the group.
return broadcast_chunk_sizes(sizes, device=device)
return sample_flex_level0(
num_frames,
num_frame_per_block=args.num_frame_per_block,
chunk_min=args.flex_chunk_min,
chunk_max=args.flex_chunk_max,
pyramid=args.flex_pyramid_levels > 1,
generator=torch_rng,
device=device,
independent_first_frame=args.independent_first_frame)["sizes"]
def initialize_kv_cache_for_training(batch_size, num_frames, frame_seq_length,
@@ -566,20 +602,39 @@ def parse_args():
"--flex_chunk_min",
type=int,
default=2,
help="Smallest chunk drawn when sampling a partition (paper 3.1 uses 2). "
"1 is allowed but adds no coverage: refining a 2-frame chunk already "
"yields 1-frame leaves."
help="Smallest chunk drawn by the *random* arm of the level-0 mixture "
"(paper §5 uses 2). Unlike stage 2, nothing here refines a coarse "
"chunk down to 1-frame leaves, so 1 is a real setting: it lets the "
"random arm emit single-frame chunks itself. The fully causal "
"[1] * F layout is a separate arm and is drawn either way."
)
parser.add_argument(
"--flex_chunk_max",
type=int,
default=10,
help="Largest chunk drawn when sampling a partition (paper 3.1 uses 10; "
"setting it equal to --flex_chunk_min pins one fixed layout). Unlike "
"stage 2 there is no coarse mixture here, so raising this to the "
"latent frame count is the *only* way stage 1 ever sees the whole "
"clip as one fully bidirectional chunk - about 5% of draws at 21 "
"latent frames. Below the frame count that layout is never drawn."
help="Largest chunk drawn by the *random* arm of the level-0 mixture "
"(paper §5 uses 10; setting it equal to --flex_chunk_min pins that "
"arm to one fixed layout). The whole-clip layout is a separate arm, "
"so this knob does not decide whether stage 1 ever regresses the "
"clip as one fully bidirectional chunk. Raising it to the latent "
"frame count only spends random-arm draws on a layout the coarse arm "
"already covers - at 21 latent frames it makes [21] 1/19 of that arm "
"and thins every other one, cutting P([3]*7) from 1/183708 to "
"1/1106560."
)
parser.add_argument(
"--flex_pyramid_levels",
type=int,
default=1,
help="Must match stage 2's --flex_pyramid_levels. Stage 1 builds no ladder "
"(it regresses x0 in one forward per chunk, so only level 0 is used), "
"but the level-0 mixture is shared with stage 2 and this flag decides "
"whether its coarse arm exists: >1 gives 30%% whole-clip / 30%% "
"uniform / 30%% full_then_blocks / 10%% random, while =1 divides the "
"coarse arm's share between the two survivors - random stays at 10%% - "
"for 45/45/10. Leaving it at 1 while stage 2 "
"runs 4 would under-train exactly the [F] layout stage 2 starts most "
"of its iterations from."
)
parser.add_argument(
"--independent_first_frame",
@@ -1044,14 +1099,13 @@ def main():
# Compute denoising_step_list from the scheduler's own grid.
# Deliberately *not* calling `set_timesteps(...)`: it rebuilds the grid with
# a different `sigma_min`, so indexing it lands up to 5.4 away from the t the
# pipeline actually denoises at (measured 5.44 at shift=8.0, 3.71 at 5.0).
# The `__init__` grid maps index k to exactly t = 1 - k / num_train_timesteps,
# which is what `stochastic_sampling_timesteps` hardcodes, so training,
# validation and inference read the same numbers - |diff| = 0.0 at both
# shifts. `train_distill.py` relies on the same grid. The explicit `.to()`
# keeps these on the accelerator, which `set_timesteps(device=...)` used to
# do for us and which the per-frame `timestep` below requires.
# a different `sigma_min`, so indexing it lands away from the t the pipeline
# actually denoises at. The `__init__` grid maps index k to exactly
# t = 1 - k / num_train_timesteps, which is what
# `stochastic_sampling_timesteps` hardcodes, so training, validation and
# inference read the same numbers. `train_distill.py` relies on the same
# grid. The explicit `.to()` keeps these on the accelerator, which the
# per-frame `timestep` below requires.
denoising_step_list = noise_scheduler.timesteps.to(accelerator.device)[
args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)
]
@@ -1064,6 +1118,13 @@ def main():
denoising_step_list, torch.tensor([0], device=denoising_step_list.device)
])
num_denoising_steps = len(train_denoising_step_list)
# Index of the clean (t=0) entry, derived from the list itself rather than
# assumed to be the last one. Passed to `get_timestep_for_ode_flexible`,
# which keeps every iteration from drawing it for *all* chunks - an
# all-clean draw zeroes the loss mask and makes the fallback MSE exactly 0,
# i.e. a zero-gradient iteration. See that function for the probability.
_clean_hits = (train_denoising_step_list == 0).nonzero(as_tuple=True)[0]
clean_step_index = int(_clean_hits[-1]) if _clean_hits.numel() > 0 else None
if accelerator.is_main_process:
print(f"Denoising step list (inference): {denoising_step_list.tolist()}")
print(f"Denoising step list (training): {train_denoising_step_list.tolist()}")
@@ -1221,6 +1282,7 @@ def main():
0, num_denoising_steps, bsz, all_num_frames,
accelerator.device,
generator=torch_rng,
clean_index=clean_step_index,
) # [B, F]
else:
index = get_timestep_for_ode(
@@ -1350,6 +1412,7 @@ def main():
0, num_denoising_steps, bsz, flex_partition,
accelerator.device,
generator=torch_rng,
clean_index=clean_step_index,
) # [B, F]
else:
index = get_timestep_for_ode(
scripts/wan2.1_flex_forcing/train_ode.sh
+1
View File
@@ -39,4 +39,5 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--flex_forcing \
--flex_chunk_min=2 \
--flex_chunk_max=10 \
--flex_pyramid_levels=4 \
--num_frame_per_block=3
scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md
+23
View File
@@ -258,6 +258,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--vae_mini_batch` | VAE encoding mini-batch size | 1 |
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
| `--fsdp_reduce_fp32` | Optional: reduce gradients in fp32 under FSDP mixed precision (default reduces in bf16) | - |
| `--enable_bucket` | Enable bucket training, no cropping, group by resolution | - |
| `--random_hw_adapt` | Auto-scale images/videos to random sizes in `[min_size, max_size]` range | - |
| `--training_with_video_token_length` | Train based on token length, supports arbitrary resolutions | - |
@@ -275,12 +276,17 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| Parameter | Description | Example Value |
|-----|------|-------|
| `--denoising_step_indices_list` | Denoising step indices list (core distillation parameter) | `1000 750 500 250` |
| `--randomize_step_indices` | Randomize the denoising timestep indices during training | - |
| `--real_guidance_scale` | Real guidance scale for scoring | 6.0 |
| `--score_timestep_clamp` | Optional: clamp DMD/critic timesteps to `[0.02, 0.98] * num_train_timesteps` | - |
| `--fake_guidance_scale` | Fake guidance scale for scoring | 0.0 |
| `--gen_update_interval` | Generator update interval | 5 |
| `--negative_prompt` | Negative prompt for distillation | Chinese negative prompt |
| `--train_sampling_steps` | Training sampling steps | 1000 |
| `--ode_transformer_path` | Path to ODE-trained weights to load into generator transformer3d | `models/Diffusion_Transformer/Self-Forcing/checkpoints/ode_init.pt` |
| `--real_score_model_name_or_path` | Optional: separate model used as real score (defaults to `pretrained_model_name_or_path`) | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
| `--score_num_frames` | Latent frames in the score-computation window; `fake_score`/`real_score` always receive this many frames | 21 |
| `--min_length_prob_bias` | Probability bias toward sampling the minimum length (`score_num_frames`); `0.0` = uniform sampling | 0.0 |
**Self-Forcing-Specific Parameters**:
@@ -294,6 +300,23 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| `--use_teacher_forcing` | Enable teacher forcing training (pass clean_x to transformer) | - |
| `--teacher_forcing_prob` | Probability of applying teacher forcing per step | 1.0 |
**GAN / Discriminator Parameters (One-Forcing)**:
> All default to `0` / off, so a plain DMD run is unaffected. Setting **any** positive `--gan_g_weight` / `--gan_d_weight` / `--r1_weight` / `--r2_weight` attaches a discriminator head to `fake_score` and **requires real videos** in the training data.
| Parameter | Description | Example Value |
|-----|------|-------|
| `--gan_g_weight` | One-Forcing: weight of the non-saturating GAN loss added to the DMD generator loss (`0` disables) | 0.0 |
| `--gan_d_weight` | One-Forcing: weight of the GAN discriminator loss added to the `fake_score` denoising loss (`0` disables) | 0.0 |
| `--gan_feature_layers` | `fake_score` block indices whose outputs feed the discriminator head (one register token per layer) | `21 29` |
| `--gan_block_ffn_dim` | FFN dim of the discriminator cross-attention blocks | 2048 |
| `--gan_block_num_heads` | Number of heads of the discriminator cross-attention blocks | 12 |
| `--relativistic_discriminator` | Use the relativistic GAN loss (fake vs. real logit difference) | - |
| `--r1_weight` | Finite-difference R1 penalty on noisy real latents (APT form, not divided by sigma) | 0.0 |
| `--r1_sigma` | Perturbation std of the R1 penalty | 0.01 |
| `--r2_weight` | Finite-difference R2 penalty on noisy generated latents | 0.0 |
| `--r2_sigma` | Perturbation std of the R2 penalty | 0.01 |
**Sample Size Configuration Guide**:
- `video_sample_size` represents the resolution size of videos; when `random_hw_adapt` is True, it represents the minimum value between video and image resolutions.
- `image_sample_size` represents the resolution size of images; when `random_hw_adapt` is True, it represents the maximum value between video and image resolutions.
scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md
+23
View File
@@ -259,6 +259,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| `--adam_epsilon` | AdamW epsilon 值 | 1e-10 |
| `--vae_mini_batch` | VAE 编码时的迷你批次大小 | 1 |
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
| `--fsdp_reduce_fp32` | 可选：FSDP 混精下梯度以 fp32 做跨卡 reduce（默认 bf16） | - |
| `--enable_bucket` | 启用分桶训练，不裁剪图片/视频，按分辨率分组训练 | - |
| `--random_hw_adapt` | 自动缩放图片/视频到 `[min_size, max_size]` 范围内的随机尺寸 | - |
| `--training_with_video_token_length` | 根据 token 长度训练，支持任意分辨率 | - |
@@ -276,12 +277,17 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| 参数 | 说明 | 示例值 |
|-----|------|-------|
| `--denoising_step_indices_list` | 去噪步骤列表（蒸馏核心参数） | `1000 750 500 250` |
| `--randomize_step_indices` | 训练时随机化去噪 timestep 索引 | - |
| `--real_guidance_scale` | 用于评分的真实 guidance scale | 6.0 |
| `--score_timestep_clamp` | 可选：将 DMD/critic 的 timestep 截断到 `[0.02, 0.98] * num_train_timesteps` | - |
| `--fake_guidance_scale` | 用于评分的虚拟 guidance scale | 0.0 |
| `--gen_update_interval` | 生成器更新间隔 | 5 |
| `--negative_prompt` | 用于蒸馏的负向提示词 | 中文负向提示词 |
| `--train_sampling_steps` | 训练采样步数 | 1000 |
| `--ode_transformer_path` | ODE 训练权重路径，加载到 generator transformer3d 中 | `models/Diffusion_Transformer/Self-Forcing/checkpoints/ode_init.pt` |
| `--real_score_model_name_or_path` | 可选：单独指定 real score 模型（默认与 `pretrained_model_name_or_path` 相同） | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
| `--score_num_frames` | 评分窗口的 latent 帧数；`fake_score`/`real_score` 始终接收这么多帧 | 21 |
| `--min_length_prob_bias` | 采样最小长度（`score_num_frames`）的概率偏置；`0.0` 表示均匀采样 | 0.0 |
**Self-Forcing 特有参数**：
@@ -295,6 +301,23 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| `--use_teacher_forcing` | 启用 teacher forcing 训练（将 clean_x 传给 transformer） | - |
| `--teacher_forcing_prob` | 每步应用 teacher forcing 的概率 | 1.0 |
**GAN / 判别器参数（One-Forcing）**：
> 全部默认为 `0` / 关闭，因此普通 DMD 训练不受影响。只要 `--gan_g_weight` / `--gan_d_weight` / `--r1_weight` / `--r2_weight` 中**任一**为正值，就会在 `fake_score` 上挂一个判别器头，并且训练数据**必须包含真实视频**。
| 参数 | 说明 | 示例值 |
|-----|------|-------|
| `--gan_g_weight` | One-Forcing：加到 DMD 生成器损失上的非饱和 GAN 损失权重（`0` 关闭） | 0.0 |
| `--gan_d_weight` | One-Forcing：加到 `fake_score` 去噪损失上的 GAN 判别器损失权重（`0` 关闭） | 0.0 |
| `--gan_feature_layers` | 输出喂给判别器头的 `fake_score` block 索引（每层一个 register token） | `21 29` |
| `--gan_block_ffn_dim` | 判别器 cross-attention block 的 FFN 维度 | 2048 |
| `--gan_block_num_heads` | 判别器 cross-attention block 的头数 | 12 |
| `--relativistic_discriminator` | 使用 relativistic GAN 损失（fake 与 real logit 之差） | - |
| `--r1_weight` | 对含噪真实 latent 的有限差分 R1 惩罚（APT 形式，不除以 sigma） | 0.0 |
| `--r1_sigma` | R1 惩罚的扰动标准差 | 0.01 |
| `--r2_weight` | 对含噪生成 latent 的有限差分 R2 惩罚 | 0.0 |
| `--r2_sigma` | R2 惩罚的扰动标准差 | 0.01 |
**Sample Size 配置指南**：
- `video_sample_size` 表示视频的分辨率大小；当 `random_hw_adapt` 为 True 时，表示视频和图像分辨率的最小值。
- `image_sample_size` 表示图像的分辨率大小；当 `random_hw_adapt` 为 True 时，表示视频和图像分辨率的最大值。
scripts/wan2.1_self_forcing/README_TRAIN_ODE.md
+3
View File
@@ -277,6 +277,9 @@ bash scripts/wan2.1_self_forcing/train_ode.sh
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
| `--trainable_modules` | Trainable modules (`"."` = all) | `"."` |
| `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` |
| `--save_state` | Save the full accelerator state on checkpoint (required to resume training) | - |
| `--use_kv_cache_training` | Run block-by-block KV cache training that matches `pipeline_wan_self_forcing` inference; otherwise fall back to one-shot causal-mask ODE regression (baseline) | - |
| `--prob_full_zero_start` | Per-sample probability of forcing all frames in all blocks to timestep index `0` (pure-noise start) to bridge the train-inference gap; `0.0` disables | 0.0 |
**ODE-specific parameters** (must match Step 1 unless you understand the consequences):
scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md
+3
View File
@@ -277,6 +277,9 @@ bash scripts/wan2.1_self_forcing/train_ode.sh
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
| `--trainable_modules` | 可训练模块（`"."` 表示全量） | `"."` |
| `--resume_from_checkpoint` | 恢复训练路径或 `"latest"` | `latest` |
| `--save_state` | 在 checkpoint 中保存完整的 accelerator 状态（恢复训练所需） | - |
| `--use_kv_cache_training` | 运行逐块 KV 缓存训练，完全匹配 `pipeline_wan_self_forcing` 推理行为；否则回退到一次性的 causal-mask ODE 回归（baseline） | - |
| `--prob_full_zero_start` | 每个样本强制所有块的所有帧使用 timestep 索引 `0`（纯噪声起点）的概率，用于弥合训练-推理差距；`0.0` 表示关闭 | 0.0 |
**ODE 特有参数**（除非清楚后果，否则需与第一步保持一致）：
scripts/wan2.1_self_forcing/train_distill.py
+385 -26
View File
@@ -17,6 +17,7 @@
import argparse
import contextlib
import dataclasses
import gc
import json
import logging
@@ -74,10 +75,12 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512,
get_random_mask)
from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
WanTransformer3DModel,
WanTransformer3DModel_SelfForcing)
WanTransformer3DModel_SelfForcing,
add_wan_gan_discriminator)
from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
WanSelfForcingPipeline)
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.fsdp_ema import FSDPEMA
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
@@ -165,6 +168,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi
return pred_for_score, score_num_frames, True
def repeat_batch(value, repeats):
"""Repeat a batched tensor or per-sample list along the batch dimension."""
if value is None or repeats == 1:
return value
if torch.is_tensor(value):
return torch.cat([value] * repeats)
return list(value) * repeats
def filter_kwargs(cls, kwargs):
import inspect
sig = inspect.signature(cls.__init__)
@@ -492,6 +504,8 @@ def parse_args():
parser.add_argument(
"--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
)
parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.")
parser.add_argument("--ema_decay", type=float, default=0.9999, help="EMA decay applied per generator update (Causal-Forcing uses 0.99).")
parser.add_argument(
"--allow_tf32",
action="store_true",
@@ -523,6 +537,9 @@ def parse_args():
parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true", help="Reduce gradients in fp32 under FSDP mixed precision (params are still gathered in low precision)."
)
parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
parser.add_argument(
@@ -722,6 +739,12 @@ def parse_args():
default=None,
help=("If you want to load the ode-trained weight into generator transformer3d, input its path."),
)
parser.add_argument(
"--real_score_model_name_or_path",
type=str,
default=None,
help=("Path to a separate model (e.g. Wan2.1-T2V-14B) used as real score. Defaults to pretrained_model_name_or_path."),
)
parser.add_argument(
"--vae_path",
type=str,
@@ -782,6 +805,11 @@ def parse_args():
default=4.5,
help="The cfg scale for real score.",
)
parser.add_argument(
"--score_timestep_clamp",
action="store_true",
help="Clamp DMD and critic timesteps to [0.02, 0.98] * num_train_timesteps.",
)
parser.add_argument(
'--denoising_step_indices_list',
nargs='+',
@@ -842,6 +870,67 @@ def parse_args():
default=1.0,
help="Probability of applying teacher forcing per step (1.0 = always)"
)
parser.add_argument(
"--gan_g_weight",
type=float,
default=0.0,
help="One-Forcing: weight of the non-saturating GAN loss added to the DMD generator loss (0 disables). "
"Any positive GAN/R1/R2 weight attaches a discriminator head to fake_score and requires real videos."
)
parser.add_argument(
"--gan_d_weight",
type=float,
default=0.0,
help="One-Forcing: weight of the GAN discriminator loss added to the fake_score denoising loss (0 disables)."
)
parser.add_argument(
"--gan_feature_layers",
nargs="+",
type=int,
default=[21, 29],
help="fake_score block indices whose outputs feed the discriminator head (one register token per layer)."
)
parser.add_argument(
"--gan_block_ffn_dim",
type=int,
default=2048,
help="FFN dim of the discriminator cross-attention blocks."
)
parser.add_argument(
"--gan_block_num_heads",
type=int,
default=12,
help="Number of heads of the discriminator cross-attention blocks."
)
parser.add_argument(
"--relativistic_discriminator",
action="store_true",
help="Use the relativistic GAN loss (fake vs. real logit difference)."
)
parser.add_argument(
"--r1_weight",
type=float,
default=0.0,
help="Finite-difference R1 penalty on noisy real latents (APT form, not divided by sigma)."
)
parser.add_argument(
"--r1_sigma",
type=float,
default=0.01,
help="Perturbation std of the R1 penalty."
)
parser.add_argument(
"--r2_weight",
type=float,
default=0.0,
help="Finite-difference R2 penalty on noisy generated latents."
)
parser.add_argument(
"--r2_sigma",
type=float,
default=0.01,
help="Perturbation std of the R2 penalty."
)
args = parser.parse_args()
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
@@ -914,6 +1003,9 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
@@ -970,6 +1062,9 @@ def main():
args.mixed_precision = accelerator.mixed_precision
args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list]
use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight))
if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1):
raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.")
# Load scheduler, tokenizer and models.
noise_scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
@@ -1030,7 +1125,7 @@ def main():
low_cpu_mem_usage=True,
).to(weight_dtype)
real_score_transformer3d = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
os.path.join(args.real_score_model_name_or_path or args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
).to(weight_dtype)
@@ -1062,7 +1157,8 @@ def main():
state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()}
m, u = generator_transformer3d.load_state_dict(state_dict, strict=False)
m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False)
if args.real_score_model_name_or_path is None:
m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False)
m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
assert len(u) == 0
@@ -1116,6 +1212,33 @@ def main():
if trainable_module_name in name:
param.requires_grad = True
break
if use_gan:
add_wan_gan_discriminator(
fake_score_transformer3d, args.gan_feature_layers, args.gan_block_ffn_dim, args.gan_block_num_heads,
weight_dtype, seed=args.seed or 0,
)
# Create EMA for the generator_transformer3d.
if args.use_ema:
if zero_stage == 3:
raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.")
if fsdp_plugin is not None:
if getattr(fsdp_plugin, "fsdp_version", 1) != 1:
raise NotImplementedError("FSDPEMA only supports FSDP1.")
ema_module = WanTransformer3DModel_SelfForcing.from_config(generator_transformer3d.config)
ema_generator_transformer3d = FSDPEMA(
ema_module, source=generator_transformer3d,
accelerator=accelerator, fsdp_plugin=fsdp_plugin,
decay=args.ema_decay,
)
else:
ema_generator_transformer3d = EMAModel(
generator_transformer3d.parameters(),
decay=args.ema_decay,
model_cls=WanTransformer3DModel_SelfForcing,
model_config=generator_transformer3d.config,
)
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
@@ -1178,6 +1301,32 @@ def main():
accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)
if args.use_ema:
def save_ema_hook(models, weights, output_dir):
if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process:
ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
def load_ema_hook(models, input_dir):
ema_path = os.path.join(input_dir, "transformer_ema")
if isinstance(ema_generator_transformer3d, FSDPEMA):
ema_generator_transformer3d.load_pretrained(ema_path)
else:
_, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True)
load_model = WanTransformer3DModel_SelfForcing.from_pretrained(
input_dir, subfolder="transformer_ema",
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=torch.float32,
)
load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config)
load_model.load_state_dict(ema_kwargs)
ema_generator_transformer3d.load_state_dict(load_model.state_dict())
ema_generator_transformer3d.to(accelerator.device)
del load_model
accelerator.register_save_state_pre_hook(save_ema_hook)
accelerator.register_load_state_pre_hook(load_ema_hook)
if args.gradient_checkpointing:
generator_transformer3d.enable_gradient_checkpointing()
fake_score_transformer3d.enable_gradient_checkpointing()
@@ -1269,6 +1418,11 @@ def main():
if accelerator.is_main_process:
print(f"Set {name} to lr : {args.learning_rate_critic / 2}")
break
if use_gan:
fake_trainable_params_optim[0]['params'].extend(
param for name, param in fake_score_transformer3d.named_parameters()
if name.startswith("gan_head.") and name not in in_already
)
if args.use_came:
optimizer = optimizer_cls(
@@ -1311,7 +1465,7 @@ def main():
args.random_hw_adapt = False
# Get the dataset
if args.train_mode != "normal" or args.use_teacher_forcing:
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
train_dataset = ImageVideoDataset(
args.train_data_meta, args.train_data_dir,
video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames,
@@ -1339,7 +1493,7 @@ def main():
return length_to_frame_num
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing:
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan:
aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
batch_sampler_generator = torch.Generator().manual_seed(args.seed)
batch_sampler = AspectRatioBatchImageVideoSampler(
@@ -1682,6 +1836,9 @@ def main():
if args.train_mode != "normal":
clip_image_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if args.use_ema:
ema_generator_transformer3d.to(accelerator.device)
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
if overrode_max_train_steps:
@@ -1772,6 +1929,8 @@ def main():
vae_stream_2 = None
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
min_score_timestep = int(0.02 * noise_scheduler.config.num_train_timesteps)
max_score_timestep = int(0.98 * noise_scheduler.config.num_train_timesteps)
def randomize_denoising_step_indices(
denoising_step_indices_list,
@@ -1832,6 +1991,8 @@ def main():
# must be averaged by contribution count, not by gradient_accumulation_steps.
train_gen_log_count = 0
train_denoising_loss = 0.0
train_gan_g_loss = 0.0
train_gan_d_loss = 0.0
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
for step, batch in enumerate(train_dataloader):
# Data batch sanity check
@@ -1852,7 +2013,7 @@ def main():
save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.mp4", rescale=True)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
if args.train_mode != "normal" or args.use_teacher_forcing:
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
# Convert images to latent space
pixel_values = batch["pixel_values"].to(weight_dtype)
@@ -1964,7 +2125,7 @@ def main():
pixel_values_bs = pixel_values_bs.sample()
new_pixel_values.append(pixel_values_bs)
return torch.cat(new_pixel_values, dim = 0)
if args.use_teacher_forcing:
if args.use_teacher_forcing or use_gan:
clean_latents = _batch_encode_vae(pixel_values)
else:
clean_latents = None
@@ -1996,7 +2157,7 @@ def main():
clip_context.append(_clip_context)
clip_context = torch.cat(clip_context)
if args.use_teacher_forcing:
if args.use_teacher_forcing or use_gan:
target_shape = clean_latents.size()
else:
target_shape = mask_latents.size()
@@ -2504,7 +2665,7 @@ def main():
generator_noise_list = [generator_noise[i] for i in range(bsz)]
clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None
# Use block_mask for causal training (一次性处理整个视频)
# Use block_mask for causal training (process the whole video at once)
generator_pred = generator_transformer3d(
x=generator_noise_list,
context=prompt_embeds,
@@ -2537,6 +2698,8 @@ def main():
# Common code for both KV cache and block mask training
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep)
generator_denoised_input = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
@@ -2642,17 +2805,53 @@ def main():
train_dmd_loss += avg_dmd_loss.item()
train_gen_log_count += 1
# One-Forcing GAN generator loss: gradients flow through fake_score into generator_pred
generator_loss = dmd_loss
if args.gan_g_weight > 0:
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
gan_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
gan_timestep = gan_timestep.clamp(min_score_timestep, max_score_timestep)
gan_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng)
gan_inputs = [add_noise(generator_pred, gan_noise, gan_timestep)]
if args.relativistic_discriminator:
real_latents = clean_latents[:, :, :generator_pred.size(2)].to(generator_pred.dtype)
gan_inputs.append(add_noise(real_latents, gan_noise, gan_timestep))
num_gan_inputs = len(gan_inputs)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, gan_logits = fake_score_transformer3d(
x=torch.cat(gan_inputs).to(dtype=weight_dtype),
context=repeat_batch(prompt_embeds, num_gan_inputs),
t=gan_timestep.repeat(num_gan_inputs),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, num_gan_inputs) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, num_gan_inputs) if args.train_mode != "normal" else None,
gan_classify=True,
)
gan_logits = gan_logits.float().chunk(num_gan_inputs)
if args.relativistic_discriminator:
gan_g_loss = F.softplus(-(gan_logits[0] - gan_logits[1])).mean()
else:
gan_g_loss = F.softplus(-gan_logits[0]).mean()
gan_g_loss = gan_g_loss * args.gan_g_weight
train_gan_g_loss += accelerator.gather(gan_g_loss.detach().repeat(args.train_batch_size)).mean().item()
generator_loss = generator_loss + gan_g_loss
if args.low_vram:
real_score_transformer3d = real_score_transformer3d.to("cpu")
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
if args.gan_g_weight == 0:
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
torch.cuda.empty_cache()
accelerator.backward(dmd_loss)
accelerator.backward(generator_loss)
if accelerator.sync_gradients:
accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
if args.gan_g_weight > 0:
# Drop the fake_score grads produced by the generator GAN loss
critic_optimizer.zero_grad()
if args.low_vram:
fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device)
@@ -2941,6 +3140,8 @@ def main():
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
critic_timestep = critic_timestep.clamp(min_score_timestep, max_score_timestep)
critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng)
fake_score_denoised_input = add_noise(
@@ -2949,15 +3150,25 @@ def main():
critic_timestep
)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
)
# Generated and real latents share critic_noise / critic_timestep. Prepare the
# real input and the R1/R2 perturbations once; they are consumed by the critic
# forward below and by the R1/R2 penalties after the main backward(s).
real_noisy = None
r1_perturbation = None
r2_perturbation = None
if use_gan:
real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype)
real_noisy = add_noise(real_latents, critic_noise, critic_timestep)
del real_latents
if args.r1_weight > 0:
r1_perturbation = args.r1_sigma * torch.randn(
real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng
)
if args.r2_weight > 0:
r2_perturbation = args.r2_sigma * torch.randn(
fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype,
device=fake_score_denoised_input.device, generator=torch_rng
)
def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
noise_pred = noise_pred.float()
@@ -2971,11 +3182,132 @@ def main():
final_loss = masked_loss.mean()
return final_loss
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
accelerator_fake_score_transformer3d.backward(denoising_loss)
if use_gan and not args.relativistic_discriminator:
# Memory-optimized critic: the non-relativistic gan_d_loss =
# softplus(-real) + softplus(fake) is separable, so run fake and real as two
# sequential 1x-batch forwards and backward the fake pass first. Backwarding
# fake frees its gradient-checkpointing block inputs before the real forward,
# dropping the peak activation from 2x batch to 1x batch. Both backward(s)
# accumulate into fake_score.grad, giving the identical update as the combined
# 2x-batch pass. R1/R2 below stay as their own 2x-batch re-forwards.
# -- Fake pass (1x batch): denoising output + fake discriminator logit --
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output, fake_gan_logits = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
gan_classify=True,
)
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
fake_logit = fake_gan_logits.float()
fake_d_loss = F.softplus(fake_logit).mean() * args.gan_d_weight
critic_loss = denoising_loss + fake_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
# Fake graph freed: the transformer blocks' stored inputs are released here.
fake_d_detached = fake_d_loss.detach()
del critic_loss, fake_gan_logits, fake_logit, fake_d_loss, fake_score_denoised_output
# -- Real pass (1x batch): discriminator logit only --
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, real_gan_logits = fake_score_transformer3d(
x=real_noisy,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
gan_classify=True,
)
real_logit = real_gan_logits.float()
real_d_loss = F.softplus(-real_logit).mean() * args.gan_d_weight
accelerator_fake_score_transformer3d.backward(real_d_loss)
# Real graph freed. Combine the two detached halves for logging; gan_d_loss is
# kept alive because the progress-bar logging below reads it.
gan_d_loss = fake_d_detached + real_d_loss.detach()
train_gan_d_loss += accelerator.gather(gan_d_loss.repeat(args.train_batch_size)).mean().item()
del fake_d_detached, real_gan_logits, real_logit, real_d_loss
else:
# Combined path: 2x-batch forward for relativistic GAN (both logits coupled by
# softplus(-(real - fake))), or 1x-batch forward when GAN is disabled.
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
if use_gan:
fake_score_denoised_output, gan_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0]
gan_logits = gan_logits.float().chunk(2)
else:
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
)
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
critic_loss = denoising_loss
if use_gan:
fake_logit, real_logit = gan_logits[0], gan_logits[1]
if args.relativistic_discriminator:
gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
else:
gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
gan_d_loss = gan_d_loss * args.gan_d_weight
train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
critic_loss = critic_loss + gan_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
del critic_loss
if use_gan and args.r1_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, r1_logits = fake_score_transformer3d(
x=torch.cat([real_noisy + r1_perturbation, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
r1_logits = r1_logits.float().chunk(2)
r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean()
accelerator_fake_score_transformer3d.backward(r1_loss)
del r1_loss
if use_gan and args.r2_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, r2_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
r2_logits = r2_logits.float().chunk(2)
r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean()
accelerator_fake_score_transformer3d.backward(r2_loss)
del r2_loss
if accelerator_fake_score_transformer3d.sync_gradients:
accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm)
critic_optimizer.step()
@@ -2988,13 +3320,22 @@ def main():
# Checks if the accelerator has performed an optimization step behind the scenes
if accelerator.sync_gradients:
# EMA tracks the generator only on rounds where it actually ran optimizer.step();
# with gen_update_interval > 1 the critic-only rounds leave the generator unchanged.
if args.use_ema and generator_update:
ema_generator_transformer3d.step(generator_transformer3d.parameters())
progress_bar.update(1)
global_step += 1
accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step)
train_log = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}
if use_gan:
train_log.update({"train_gan_g_loss": train_gan_g_loss / max(train_gen_log_count, 1), "train_gan_d_loss": train_gan_d_loss})
accelerator.log(train_log, step=global_step)
train_dmd_loss = 0.0
train_gen_log_count = 0
train_denoising_loss = 0.0
train_gan_g_loss = 0.0
train_gan_d_loss = 0.0
if global_step % args.checkpointing_steps == 0:
if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
@@ -3038,6 +3379,10 @@ def main():
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
with progress_bar.paused():
if args.use_ema:
# Store the generator parameters temporarily and load the EMA parameters to perform inference.
ema_generator_transformer3d.store(generator_transformer3d.parameters())
ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
log_validation(
vae,
text_encoder,
@@ -3050,8 +3395,15 @@ def main():
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original generator parameters.
ema_generator_transformer3d.restore(generator_transformer3d.parameters())
logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
if use_gan:
logs["gan_d_loss"] = gan_d_loss.detach().item()
if args.gan_g_weight > 0:
logs["gan_g_loss"] = gan_g_loss.detach().item()
progress_bar.set_postfix(**logs)
if global_step >= args.max_train_steps:
@@ -3059,6 +3411,10 @@ def main():
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
with progress_bar.paused():
if args.use_ema:
# Store the generator parameters temporarily and load the EMA parameters to perform inference.
ema_generator_transformer3d.store(generator_transformer3d.parameters())
ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
log_validation(
vae,
text_encoder,
@@ -3071,6 +3427,9 @@ def main():
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original generator parameters.
ema_generator_transformer3d.restore(generator_transformer3d.parameters())
# Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever
# something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto
scripts/wan2.1_self_forcing/train_ode.py
+10
View File
@@ -15,6 +15,7 @@
# See the License for the specific language governing permissions and
import argparse
import dataclasses
import gc
import logging
import math
@@ -362,6 +363,10 @@ def parse_args():
parser.add_argument("--adam_weight_decay", type=float, default=3e-2, help="Weight decay to use.")
parser.add_argument("--adam_epsilon", type=float, default=1e-10, help="Epsilon value for the Adam optimizer.")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true",
help="Reduce gradients in fp32 under FSDP mixed precision (params are still gathered in low precision).",
)
parser.add_argument(
"--mixed_precision",
type=str,
@@ -581,6 +586,11 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(
fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32
)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
scripts/wan2.2/README_TRAIN_REWARD.md
+374
View File
@@ -0,0 +1,374 @@
# Wan2.2 Reward LoRA Training Guide
This document provides a complete workflow for **Reward Backpropagation** fine-tuning of Wan2.2, i.e. optimizing a LoRA so that the generated videos better align with human preferences. It covers environment configuration, data preparation, the reward-specific training arguments (single / multi reward, differentiable VAE decoding, backprop strategy), training validation, and inference testing.
> **Note**: Wan2.2 is a video generation model supporting text-to-video (T2V), image-to-video (I2V), and text-image-to-video (TI2V). The A14B variants adopt a dual-Transformer architecture (high-noise / low-noise models), while the TI2V-5B variant uses a single-Transformer architecture. Reward LoRA training runs a differentiable rollout (a short denoising trajectory) and back-propagates the reward model's loss into the trainable expert, so it needs more VRAM than standard LoRA fine-tuning. This guide covers training your own reward LoRA for any of the Wan2.2 variants.
We explore the Reward Backpropagation technique <sup>[1](#ref1) [2](#ref2)</sup> to optimize the generated videos by [Wan2.2](https://github.com/aigc-apps/VideoX-Fun) for better alignment with human preferences. The training script is [train_reward_lora.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/wan2.2/train_reward_lora.py); you can train your own reward LoRA and load it onto the corresponding base model as a plug-in.
> [!NOTE]
> Reward backpropagation tends to reward-hack: as the reward keeps rising the perceptual quality can plateau or degrade (adversarial artifacts in the background). Prefer a **video reward** (e.g. `VideoAlignReward`) or a **multi-reward** mix (e.g. `HPSv3Reward` + `VideoAlignReward`) over a pure image reward to mitigate the loss of temporal dynamism — see [3.4](#34-single-vs-multi-reward-training) and [5. Limitations](#5-limitations).
---
## Table of Contents
- [1. Environment Configuration](#1-environment-configuration)
- [2. Data Preparation](#2-data-preparation)
- [2.1 Quick Test Dataset](#21-quick-test-dataset)
- [2.2 Reward Model Weights](#22-reward-model-weights)
- [2.3 Dataset Structure and metadata.json](#23-dataset-structure-and-metadatajson)
- [3. Reward LoRA Training](#3-reward-lora-training)
- [3.1 Download Pre-trained Base Model](#31-download-pre-trained-base-model)
- [3.2 Quick Start (FSDP)](#32-quick-start-fsdp)
- [3.3 Reward-specific Parameter Explanation](#33-reward-specific-parameter-explanation)
- [3.4 Single vs Multi-Reward Training](#34-single-vs-multi-reward-training)
- [3.5 Differentiable Decoding and `--use_1steps_vae`](#35-differentiable-decoding-and---use_1steps_vae)
- [3.6 Training Validation](#36-training-validation)
- [4. Inference Testing](#4-inference-testing)
- [5. Limitations](#5-limitations)
- [6. More Resources](#6-more-resources)
- [References](#references)
---
## 1. Environment Configuration
**Method 1: Using requirements.txt**
```bash
pip install -r requirements.txt
```
**Method 2: Manual Dependency Installation**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> Some video reward models (e.g. `VideoAlignReward`, `HPSv3Reward`) are built on a Qwen2-VL backbone and require `transformers==4.51.3`. Install the matching version when using those rewards.
**Method 3: Using Docker**
When using Docker, please ensure that the graphics card driver and CUDA environment are correctly installed on your machine, then execute the following commands:
```
# pull image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# enter image
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 2. Data Preparation
### 2.1 Quick Test Dataset
Reward LoRA training runs a differentiable rollout from text prompts, so a **prompt-only** dataset is sufficient for the standard T2V reward path (`train_mode="normal"`), which is the recommended setup. For the I2V / inpaint path, a video dataset provides the reference conditioning.
```bash
# Download official example dataset
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 Reward Model Weights
The backbone / checkpoint source differs per reward:
- **`HPSReward` (HPS v2.1) and `MPSReward`**: the score heads are tiny and are **downloaded automatically on first run** into the torch hub cache (`~/.cache/torch/hub`); note that **HPS additionally needs two special config files patched into `hpsv2`'s bundled open_clip** (see below), while MPS needs no extra setup;
- **`VideoAlignReward` (VideoReward) and `HPSv3Reward`**: built on Qwen2-VL backbones and fairly large, so **stage them locally in advance** and point `--reward_fn_kwargs` at the local paths.
The dependencies and sources are summarized below (see [3.3](#33-reward-specific-parameter-explanation) for the `--reward_fn` values):
| Reward | Backbone / dependency | Score head / checkpoint | Manual download |
|--|--|--|--|
| `HPSReward` (`version="v2.1"`) | open_clip `ViT-H-14`, requires `pip install hpsv2` + the `factory.py` patch and BPE vocab | `HPS_v2.1_compressed.pt` (auto-downloaded) | No, override with `model_path` |
| `MPSReward` | `laion/CLIP-ViT-H-14-laion2B-s32B-b79K` (config/tokenizer auto-downloaded); MPS code is vendored under `videox_fun/reward/MPS/` | `MPS_overall.pth` (auto-downloaded) | No, override with `model_path` |
| `VideoAlignReward` | `Qwen/Qwen2-VL-2B-Instruct` | `KlingTeam/VideoReward` (contains `model_config.json` + weights) | **Yes**, `model_path` is required |
| `HPSv3Reward` | `Qwen/Qwen2-VL-7B-Instruct` | `HPSv3.safetensors` from `MizzenAI/HPSv3` | Recommended; auto-downloaded from HF when omitted |
**HPS v2.1 (`HPSReward`)** — besides `pip install hpsv2`, you must drop two special config files into `hpsv2`'s bundled open_clip: `create_model_and_transforms("ViT-H-14", ...)` uses arguments added by the HPS fork, so it needs the patched `factory.py`; `get_tokenizer("ViT-H-14")` needs the BPE vocab `bpe_simple_vocab_16e6.txt.gz`. Loading fails if either is missing. The `ViT-H-14` backbone and the `HPS_v2.1_compressed.pt` head download automatically to the torch hub cache on first run:
```bash
pip install hpsv2
# Patch hpsv2's bundled open_clip factory.py and add the BPE vocab file
site_packages=$(python -c "import site; print(site.getsitepackages()[0])")
wget -O $site_packages/hpsv2/src/open_clip/factory.py https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/package/patches/hpsv2_src_open_clip_factory_patches.py
wget -O $site_packages/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz https://github.com/tgxs002/HPSv2/raw/refs/heads/master/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz
```
```bash
--reward_fn="HPSReward" \
--reward_fn_kwargs='{"version": "v2.1"}'
```
**MPS (`MPSReward`)** — the MPS code is vendored under `videox_fun/reward/MPS/` and needs no patching; the `MPS_overall.pth` head downloads automatically from OSS, and the config + tokenizer of `laion/CLIP-ViT-H-14-laion2B-s32B-b79K` download automatically from HuggingFace (if HF is unreachable, point `processor_name_or_path` to a local mirror):
```bash
--reward_fn="MPSReward" \
--reward_fn_kwargs='{"max_reward": 1, "loss_scale": 1}'
```
**VideoReward (`VideoAlignReward`)** — needs the Qwen2-VL-2B backbone + a VideoReward head directory (containing `model_config.json` and weights). Download both in advance; `model_path` is required:
```bash
modelscope download --model Qwen/Qwen2-VL-2B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-2B-Instruct
huggingface-cli download KlingTeam/VideoReward --local-dir models/Personalized_Model/VideoReward
```
```bash
--reward_fn="VideoAlignReward" \
--reward_fn_kwargs='{"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "reward_dim": "TA+MQ", "fps": 16, "use_norm": true, "differentiable": true, "max_reward": 2.0}'
```
**HPSv3 (`HPSv3Reward`)** — needs the Qwen2-VL-7B backbone + `HPSv3.safetensors`. When `checkpoint_path` is omitted it downloads automatically from HF `MizzenAI/HPSv3`; staging it locally is recommended:
```bash
modelscope download --model Qwen/Qwen2-VL-7B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-7B-Instruct
modelscope download --model MizzenAI/HPSv3 --local_dir models/Personalized_Model/HPSv3
```
```bash
--reward_fn="HPSv3Reward" \
--reward_fn_kwargs='{"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}'
```
> With `--backprop`, `VideoAlignReward` / `HPSv3Reward` must set `"differentiable": true`, otherwise their `@torch.no_grad()` path cuts the gradient chain (see [3.3](#33-reward-specific-parameter-explanation)).
### 2.3 Dataset Structure and metadata.json
```
📦 datasets/
├── 📂 my_dataset/
│ ├── 📂 train/
│ │ ├── 📄 video001.mp4
│ │ ├── 📄 video002.mp4
│ │ └── 📄 ...
│ └── 📄 metadata.json
```
**Relative Path Format** (example):
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
**Key Field Descriptions**:
- `file_path`: Video path (relative or absolute path).
- `text`: Video description (English prompt).
- `type`: Data type, fixed as `"video"`.
- `width` / `height`: Video width and height (**highly recommended**, used for bucket training). You can use `scripts/process_json_add_width_and_height.py` to extract these fields.
**Relative vs Absolute Path**:
- Relative: `export DATASET_NAME="datasets/X-Fun-Videos-Demo/"` and `export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"`.
- Absolute: `export DATASET_NAME=""` and `export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"`.
---
## 3. Reward LoRA Training
### 3.1 Download Pre-trained Base Model
```bash
# Create model directory
mkdir -p models/Diffusion_Transformer
# Download Wan2.2 official weights
# T2V model (text-to-video, dual-Transformer)
modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-T2V-A14B
# or I2V model (image-to-video, dual-Transformer)
# modelscope download --model Wan-AI/Wan2.2-I2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-I2V-A14B
# or TI2V model (text-image-to-video, single-Transformer)
# modelscope download --model Wan-AI/Wan2.2-TI2V-5B --local_dir models/Diffusion_Transformer/Wan2.2-TI2V-5B
```
### 3.2 Quick Start (FSDP)
Reward LoRA training keeps the full denoising rollout and the reward model's forward graph in memory, so FSDP (`FULL_SHARD`) is recommended. With `--boundary_type="high"`, the *other* (frozen) expert is already sharded and resident on GPU and is reused during validation, so no second copy is re-loaded from disk.
**Wan2.2 T2V Reward LoRA Training Example (single HPS v2.1 reward, FSDP)**:
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
export NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False \
scripts/wan2.2/train_reward_lora.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=1 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir_wan2.2_reward_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--enable_bucket \
--boundary_type="high" \
--rank=128 \
--network_alpha=64 \
--target_name="q,k,v,ffn.0,ffn.2" \
--train_mode="normal" \
--guidance_scale=5.0 \
--num_inference_steps=40 \
--num_decoded_latents=3 \
--reward_fn="HPSReward" \
--reward_fn_kwargs='{"version": "v2.1"}' \
--backprop \
--backprop_strategy="tail" \
--backprop_num_steps=1 \
--low_vram
```
> The `train_reward_lora.sh` script in this directory provides a ready-to-run template. Adjust `MODEL_NAME`, dataset paths, `--reward_fn`, and `--boundary_type` to your setup.
### 3.3 Reward-specific Parameter Explanation
Besides the standard LoRA arguments (see [README_TRAIN_LORA.md](./README_TRAIN_LORA.md)), the reward trainer adds the following:
| Parameter | Description | Example Value |
|-----------|-------------|---------------|
| `--reward_fn` | Reward model class name(s). A **comma-separated list** enables multi-reward training; the differentiable losses are combined by a weighted sum. Available: `AestheticReward`, `HPSReward`, `PickScoreReward`, `MPSReward`, `HPSv3Reward`, `VideoAlignReward`. | `"HPSv3Reward,VideoAlignReward"` |
| `--reward_fn_kwargs` | Keyword args for the reward model(s). A **flat JSON dict** for a single reward; a **nested JSON dict keyed by reward name** for multiple rewards. | `'{"VideoAlignReward": {"max_reward": 2.0}, "HPSv3Reward": {"max_reward": 12}}'` |
| `--multi_reward_weights` | JSON dict of per-reward loss weights. Normalized to sum to 1; if omitted with multiple rewards, equal weights are used. Ignored for a single reward. | `'{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}'` |
| `--num_inference_steps` | Number of denoising steps in the differentiable rollout and validation. | 40 |
| `--num_decoded_latents` | Number of latent frames decoded back to pixels for the reward. The chunked causal VAE limits this to `<= 3` (gradient-checkpointing constraint) unless `--use_1steps_vae` is set. | 3 |
| `--num_reward_frames` | Randomly select this many decoded frames for reward evaluation. `None` = use all. Only applied to **image** rewards (see below). | 2 |
| `--use_1steps_vae` | Use the single-pass Wan2.2 VAE (`AutoencoderKLWan3_8_1Steps`) for differentiable decoding; numerically equivalent, decodes the whole latent in one forward and lifts the `<= 3` latent limit. | - |
| `--backprop` | Enable the reward-backpropagation training mode. | - |
| `--backprop_strategy` | Which rollout steps keep a gradient graph: `last`, `tail`, `uniform`, `random`. | `tail` |
| `--backprop_num_steps` | Number of backprop steps, used when `--backprop_strategy` is `tail`/`uniform`/`random`. | 1 |
| `--first_frame_mse_weight` | Weight for the first-frame latent MSE auxiliary loss in inpaint mode. | 1.0 |
**Image vs Video rewards (`is_image_reward`)**:
- **Image rewards** (`AestheticReward`, `HPSReward`, `PickScoreReward`, `MPSReward`, `HPSv3Reward`) score individual frames, so frame subsampling (`--num_reward_frames`) is safe.
- **Video rewards** (`VideoAlignReward`, `is_image_reward=False`) need the full temporal sequence to judge motion quality; collapsing frames would degrade the video to a near-static input and break the motion dimension. For these, all decoded frames are kept (count bounded by `--num_decoded_latents`).
- In inpaint mode (`train_mode="inpaint"`), the first frame is dropped before reward evaluation for every reward.
> **IMPORTANT**: when using `--backprop`, a video reward such as `VideoAlignReward` / `HPSv3Reward` must be constructed with `"differentiable": true` inside `--reward_fn_kwargs`; otherwise its `@torch.no_grad()` path cuts the gradient chain entirely.
### 3.4 Single vs Multi-Reward Training
A single reward uses `--reward_fn="HPSv3Reward"` with a flat `--reward_fn_kwargs`. To combine several objectives, pass a comma-separated list, a nested kwargs dict, and (optionally) weights. Per-reward loss and reward values are logged separately (e.g. `train/loss_HPSv3Reward`, `train/reward_VideoAlignReward`), while the optimized `loss = Σ wᵢ · lossᵢ`.
**Multi-reward example (HPSv3 + VideoAlign, T2V)** — the strongest alignment setup, since the video reward supplies the temporal (motion / consistency) signal that image rewards cannot:
```bash
--num_decoded_latents=13 \
--num_reward_frames=2 \
--use_1steps_vae \
--reward_fn="VideoAlignReward,HPSv3Reward" \
--reward_fn_kwargs='{"VideoAlignReward": {"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "fps": 16, "reward_dim": "TA+MQ", "differentiable": true, "use_norm": true, "max_reward": 2.0, "loss_scale": 1.0}, "HPSv3Reward": {"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}}' \
--multi_reward_weights='{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}' \
--backprop \
--backprop_strategy="tail" \
--backprop_num_steps=1
```
### 3.5 Differentiable Decoding and `--use_1steps_vae`
The default chunked causal VAE (`AutoencoderKLWan3_8`) applies gradient checkpointing per chunk, which limits differentiable decoding to the **first 3 latent frames** (`--num_decoded_latents <= 3`). Setting `--num_decoded_latents > 3` without `--use_1steps_vae` raises a `ValueError`.
Enable `--use_1steps_vae` to switch to the single-pass variant, which decodes the whole latent in one forward, applies checkpointing cleanly per decoder block, and lifts the limit — useful when a video reward needs more temporal context. It is numerically equivalent and only affects the `AutoencoderKLWan3_8` (2.2 VAE) path.
### 3.6 Training Validation
You can configure validation parameters to periodically generate test videos during training so you can monitor alignment progress. Under FSDP, `log_validation` runs a pipeline forward that is a cross-rank collective, so **all processes must enter the validation call** (it is intentionally *not* wrapped in `is_main_process`); the trainable expert already carries the LoRA, and the frozen expert is reused rather than re-loaded.
| Parameter | Description | Default Value |
|-----------|-------------|---------------|
| `--validation_steps` | Run validation every N steps | 2000 |
| `--validation_epochs` | Run validation every N epochs | 5 |
| `--validation_prompts` | Prompts for validating video generation | None |
| `--validation_paths` | Reference image paths for I2V validation (inpaint mode only) | None |
**Validation Example**:
```bash
--validation_prompts="A stylish woman strolls down a bustling Tokyo street at night, warm neon glow and animated city signage reflecting on the wet pavement; she wears a black leather jacket over a red dress, confident casual demeanor, medium shot, cinematic." \
--validation_steps=50 \
--validation_epochs=500
```
> Validation videos are saved to the `output_dir` directory. For T2V validation (without a reference image), keep `train_mode="normal"` and omit `--validation_paths`.
---
## 4. Inference Testing
The trained reward LoRA is a normal LoRA and is loaded the same way as [README_TRAIN_LORA.md](./README_TRAIN_LORA.md). For an A14B dual-Transformer model, set `lora_path` / `lora_weight` for the low-noise reward LoRA and `lora_high_path` / `lora_high_weight` for the high-noise reward LoRA in [examples/wan2.2/predict_t2v.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.2/predict_t2v.py) (or `predict_i2v.py`):
```python
# Path to the trained reward LoRA weights, e.g.
# "output_dir_wan2.2_reward_lora/checkpoint-xxx/diffusion_pytorch_model.safetensors"
lora_path = None
lora_high_path = None
lora_weight = 0.55
lora_high_weight = 0.55
```
Run single-GPU inference with:
```bash
python examples/wan2.2/predict_t2v.py
```
For memory-management modes and multi-GPU parallel inference, refer to the Inference Testing section of [README_TRAIN_LORA.md](./README_TRAIN_LORA.md).
---
## 5. Limitations
1. After training to a certain extent, the reward keeps increasing but the quality of the generated videos no longer improves. The model trickily learns shortcuts (adding artifacts in the background, i.e. adversarial patches) to raise the reward.
2. There is still a lack of suitable preference models for video generation. Directly using image preference models cannot evaluate preferences along the temporal dimension (such as dynamism and consistency), and tends to reduce the dynamism of generated videos. This can be partially mitigated by computing the reward using only the first decoded frame, but the impact persists — which is why a video reward (e.g. `VideoAlignReward`) with `--use_1steps_vae` and a larger `--num_decoded_latents`, or a multi-reward mix, is recommended for motion-sensitive alignment.
---
## 6. More Resources
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
- **Standard LoRA training**: [README_TRAIN_LORA.md](./README_TRAIN_LORA.md)
- **Chinese version**: [README_TRAIN_REWARD_zh-CN.md](./README_TRAIN_REWARD_zh-CN.md)
---
## References
<ol>
<li id="ref1">Clark, Kevin, et al. "Directly fine-tuning diffusion models on differentiable rewards.". In ICLR 2024.</li>
<li id="ref2">Prabhudesai, Mihir, et al. "Aligning text-to-image diffusion models with reward backpropagation." arXiv preprint arXiv:2310.03739 (2023).</li>
</ol>
scripts/wan2.2/README_TRAIN_REWARD_zh-CN.md
+374
View File
@@ -0,0 +1,374 @@
# Wan2.2 Reward LoRA 训练指南
本文档提供 Wan2.2 **Reward 反向传播（Reward Backpropagation）** 微调的完整流程，即用 LoRA 优化生成视频，使其更符合人类偏好。内容涵盖环境配置、数据准备、Reward 专用训练参数（单/多 reward、可微 VAE 解码、backprop 策略）、训练验证与推理测试。
> **说明**：Wan2.2 是一个支持文本到视频（T2V）、图像到视频（I2V）、文本-图像到视频（TI2V）的视频生成模型。A14B 系列采用双 Transformer 架构（高噪声/低噪声模型），TI2V-5B 使用单 Transformer 架构。Reward LoRA 训练会跑一段可微的 rollout（一条简短的去噪轨迹），并把 reward 模型的 loss 反传到可训练的专家上，因此比标准 LoRA 微调更耗显存。本指南覆盖为任意 Wan2.2 变体训练自有 reward LoRA 的流程。
我们采用 Reward Backpropagation 技术 <sup>[1](#ref1) [2](#ref2)</sup> 来优化 [Wan2.2](https://github.com/aigc-apps/VideoX-Fun) 生成的视频，使其更好地对齐人类偏好。训练脚本为 [train_reward_lora.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/wan2.2/train_reward_lora.py)；你可以训练自己的 reward LoRA，并作为插件加载到对应的基座模型上。
> [!NOTE]
> Reward 反向传播容易出现 reward-hacking：reward 持续上升，但感知质量可能停滞甚至退化（背景出现对抗伪影）。相比纯图像 reward，更推荐使用**视频 reward**（如 `VideoAlignReward`）或**多 reward 组合**（如 `HPSv3Reward` + `VideoAlignReward`），以缓解时序运动性的下降——见 [3.4](#34-单-reward-与多-reward-训练) 与 [五、局限性](#五局限性)。
---
## 目录
- [一、环境配置](#一环境配置)
- [二、数据准备](#二数据准备)
- [2.1 快速测试数据集](#21-快速测试数据集)
- [2.2 Reward 模型权重](#22-reward-模型权重)
- [2.3 数据集结构与 metadata.json](#23-数据集结构与-metadatajson)
- [三、Reward LoRA 训练](#三reward-lora-训练)
- [3.1 下载预训练基座模型](#31-下载预训练基座模型)
- [3.2 快速开始（FSDP）](#32-快速开始fsdp)
- [3.3 Reward 专用参数解析](#33-reward-专用参数解析)
- [3.4 单 Reward 与多 Reward 训练](#34-单-reward-与多-reward-训练)
- [3.5 可微解码与 `--use_1steps_vae`](#35-可微解码与---use_1steps_vae)
- [3.6 训练验证](#36-训练验证)
- [四、推理测试](#四推理测试)
- [五、局限性](#五局限性)
- [六、更多资源](#六更多资源)
- [参考文献](#参考文献)
---
## 一、环境配置
**方式 1：使用 requirements.txt**
```bash
pip install -r requirements.txt
```
**方式 2：手动安装依赖**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 部分视频 reward 模型（如 `VideoAlignReward`、`HPSv3Reward`）基于 Qwen2-VL 骨干，需要 `transformers==4.51.3`。使用这些 reward 时请安装对应版本。
**方式 3：使用 docker**
使用 docker 的情况下，请保证机器中已经正确安装显卡驱动与 CUDA 环境，然后以此执行以下命令：
```
# pull image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# enter image
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 二、数据准备
### 2.1 快速测试数据集
Reward LoRA 训练从文本 prompt 出发跑可微 rollout，因此标准的 T2V reward 路径（`train_mode="normal"`）只需一个**纯 prompt 数据集**即可，这也是推荐配置。对于 I2V/inpaint 路径，则使用视频数据集来提供参考条件。
```bash
# 下载官方示例数据集
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 Reward 模型权重
不同 reward 的骨干/checkpoint 来源不同：
- **`HPSReward`（HPS v2.1）与 `MPSReward`**：打分头很小，会在**首次运行时自动下载**到 torch hub 缓存（`~/.cache/torch/hub`）；其中 **HPS 还需给 `hpsv2` 内置的 open_clip 补两个特殊配置文件**（见下），MPS 无需额外准备；
- **`VideoAlignReward`（VideoReward）与 `HPSv3Reward`**：基于 Qwen2-VL 骨干，体积较大，**建议提前下载到本地**，再用 `--reward_fn_kwargs` 指向本地路径。
各 reward 的依赖与来源汇总如下（`--reward_fn` 取值见 [3.3](#33-reward-专用参数解析)）：
| Reward | 骨干 / 依赖 | 打分头 / checkpoint | 手动下载 |
|--|--|--|--|
| `HPSReward`（`version="v2.1"`） | open_clip `ViT-H-14`，需 `pip install hpsv2` + `factory.py` 补丁与 BPE 词表 | `HPS_v2.1_compressed.pt`（自动下载） | 否，可用 `model_path` 覆盖 |
| `MPSReward` | `laion/CLIP-ViT-H-14-laion2B-s32B-b79K`（config/tokenizer 自动下载），MPS 代码已内置于 `videox_fun/reward/MPS/` | `MPS_overall.pth`（自动下载） | 否，可用 `model_path` 覆盖 |
| `VideoAlignReward` | `Qwen/Qwen2-VL-2B-Instruct` | `KlingTeam/VideoReward`（含 `model_config.json` 与权重） | **是**，`model_path` 必填 |
| `HPSv3Reward` | `Qwen/Qwen2-VL-7B-Instruct` | `MizzenAI/HPSv3` 的 `HPSv3.safetensors` | 建议，缺省时自动从 HF 下载 |
**HPS v2.1（`HPSReward`）**——除了 `pip install hpsv2`，还要给 `hpsv2` 内置的 open_clip 补两个特殊配置文件：`create_model_and_transforms("ViT-H-14", ...)` 用到了 HPS 分支新增的参数，需要打过补丁的 `factory.py`；`get_tokenizer("ViT-H-14")` 需要 BPE 词表 `bpe_simple_vocab_16e6.txt.gz`。缺任一都会在加载时报错。`ViT-H-14` 骨干与 `HPS_v2.1_compressed.pt` 打分头会在首次运行时自动下载到 torch hub 缓存：
```bash
pip install hpsv2
# 覆盖 hpsv2 内置 open_clip 的 factory.py 补丁，并补上 BPE 词表
site_packages=$(python -c "import site; print(site.getsitepackages()[0])")
wget -O $site_packages/hpsv2/src/open_clip/factory.py https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/package/patches/hpsv2_src_open_clip_factory_patches.py
wget -O $site_packages/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz https://github.com/tgxs002/HPSv2/raw/refs/heads/master/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz
```
```bash
--reward_fn="HPSReward" \
--reward_fn_kwargs='{"version": "v2.1"}'
```
**MPS（`MPSReward`）**——MPS 代码已内置于 `videox_fun/reward/MPS/`，无需打补丁；`MPS_overall.pth` 打分头会自动从 OSS 下载，`laion/CLIP-ViT-H-14-laion2B-s32B-b79K` 的 config 与 tokenizer 会自动从 HuggingFace 下载（若无法访问 HF，可用 `processor_name_or_path` 指向本地镜像）：
```bash
--reward_fn="MPSReward" \
--reward_fn_kwargs='{"max_reward": 1, "loss_scale": 1}'
```
**VideoReward（`VideoAlignReward`）**——需 Qwen2-VL-2B 骨干 + VideoReward 打分头目录（内含 `model_config.json` 与权重），二者都要提前下载，且 `model_path` 必填：
```bash
modelscope download --model Qwen/Qwen2-VL-2B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-2B-Instruct
huggingface-cli download KlingTeam/VideoReward --local-dir models/Personalized_Model/VideoReward
```
```bash
--reward_fn="VideoAlignReward" \
--reward_fn_kwargs='{"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "reward_dim": "TA+MQ", "fps": 16, "use_norm": true, "differentiable": true, "max_reward": 2.0}'
```
**HPSv3（`HPSv3Reward`）**——需 Qwen2-VL-7B 骨干 + `HPSv3.safetensors`；不指定 `checkpoint_path` 时会自动从 HF `MizzenAI/HPSv3` 下载，建议提前落盘：
```bash
modelscope download --model Qwen/Qwen2-VL-7B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-7B-Instruct
modelscope download --model MizzenAI/HPSv3 --local_dir models/Personalized_Model/HPSv3
```
```bash
--reward_fn="HPSv3Reward" \
--reward_fn_kwargs='{"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}'
```
> 使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 必须设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会切断梯度链（见 [3.3](#33-reward-专用参数解析)）。
### 2.3 数据集结构与 metadata.json
```
📦 datasets/
├── 📂 my_dataset/
│ ├── 📂 train/
│ │ ├── 📄 video001.mp4
│ │ ├── 📄 video002.mp4
│ │ └── 📄 ...
│ └── 📄 metadata.json
```
**相对路径格式**（示例）：
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
**关键字段说明**：
- `file_path`：视频路径（相对路径或绝对路径）。
- `text`：视频描述（英文 prompt）。
- `type`：数据类型，固定为 `"video"`。
- `width` / `height`：视频宽高（**强烈建议提供**，用于分桶训练）。可用 `scripts/process_json_add_width_and_height.py` 提取这两个字段。
**相对路径与绝对路径**：
- 相对路径：`export DATASET_NAME="datasets/X-Fun-Videos-Demo/"` 与 `export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"`。
- 绝对路径：`export DATASET_NAME=""` 与 `export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"`。
---
## 三、Reward LoRA 训练
### 3.1 下载预训练基座模型
```bash
# 创建模型目录
mkdir -p models/Diffusion_Transformer
# 下载 Wan2.2 官方权重
# T2V 模型（文本到视频，双 Transformer）
modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-T2V-A14B
# 或 I2V 模型（图像到视频，双 Transformer）
# modelscope download --model Wan-AI/Wan2.2-I2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-I2V-A14B
# 或 TI2V 模型（文本-图像到视频，单 Transformer）
# modelscope download --model Wan-AI/Wan2.2-TI2V-5B --local_dir models/Diffusion_Transformer/Wan2.2-TI2V-5B
```
### 3.2 快速开始（FSDP）
Reward LoRA 训练会把完整的去噪 rollout 与 reward 模型的前向计算图保持在显存中，因此推荐使用 FSDP（`FULL_SHARD`）。当 `--boundary_type="high"` 时，*另一个*（冻结的）专家在训练期已被分片并常驻 GPU，验证阶段会直接复用它，而不会从磁盘再加载一份副本。
**Wan2.2 T2V Reward LoRA 训练示例（单一 HPS v2.1 reward，FSDP）**：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 与 NCCL_P2P_DISABLE=1 用于没有 RDMA 的多机环境。
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
export NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False \
scripts/wan2.2/train_reward_lora.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=1 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir_wan2.2_reward_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--enable_bucket \
--boundary_type="high" \
--rank=128 \
--network_alpha=64 \
--target_name="q,k,v,ffn.0,ffn.2" \
--train_mode="normal" \
--guidance_scale=5.0 \
--num_inference_steps=40 \
--num_decoded_latents=3 \
--reward_fn="HPSReward" \
--reward_fn_kwargs='{"version": "v2.1"}' \
--backprop \
--backprop_strategy="tail" \
--backprop_num_steps=1 \
--low_vram
```
> 本目录下的 `train_reward_lora.sh` 提供了一份可直接运行的模板。请根据你的环境调整 `MODEL_NAME`、数据集路径、`--reward_fn` 与 `--boundary_type`。
### 3.3 Reward 专用参数解析
除标准 LoRA 参数外（见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)），reward 训练脚本还新增以下参数：
| 参数 | 说明 | 示例值 |
|-----------|-------------|---------------|
| `--reward_fn` | Reward 模型类名，可用**逗号分隔的列表**启用多 reward 训练；各 reward 的可微 loss 以加权和合并。可选：`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`、`VideoAlignReward`。 | `"HPSv3Reward,VideoAlignReward"` |
| `--reward_fn_kwargs` | Reward 模型的关键词参数。单 reward 用**扁平 JSON dict**；多 reward 用**以 reward 名为 key 的嵌套 JSON dict**。 | `'{"VideoAlignReward": {"max_reward": 2.0}, "HPSv3Reward": {"max_reward": 12}}'` |
| `--multi_reward_weights` | 多 reward 各 loss 权重的 JSON dict。会归一化到和为 1；多 reward 时若省略则使用等权重；单 reward 时忽略。 | `'{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}'` |
| `--num_inference_steps` | 可微 rollout 与验证中的去噪步数。 | 40 |
| `--num_decoded_latents` | 解码回像素以供 reward 评估的潜帧数。分块因果 VAE 受梯度检查点限制，除非设置 `--use_1steps_vae`，否则要求 `<= 3`。 | 3 |
| `--num_reward_frames` | 随机选取该数量的解码帧用于 reward 评估。`None` = 使用全部。仅对**图像** reward 生效（见下）。 | 2 |
| `--use_1steps_vae` | 使用单段式 Wan2.2 VAE（`AutoencoderKLWan3_8_1Steps`）做可微解码；数值等价，一次前向解码整个 latent，干净地按解码块应用梯度检查点，并解除 `<= 3` 的限制。仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。 | - |
| `--backprop` | 启用 reward 反向传播训练模式。 | - |
| `--backprop_strategy` | rollout 中哪些步保留梯度图：`last`、`tail`、`uniform`、`random`。 | `tail` |
| `--backprop_num_steps` | backprop 步数，仅当 `--backprop_strategy` 为 `tail`/`uniform`/`random` 时生效。 | 1 |
| `--first_frame_mse_weight` | inpaint 模式下首帧 latent MSE 辅助 loss 的权重。 | 1.0 |
**图像 reward 与视频 reward（`is_image_reward`）**：
- **图像 reward**（`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`）逐帧打分，因此帧抽稀（`--num_reward_frames`）是安全的。
- **视频 reward**（`VideoAlignReward`，`is_image_reward=False`）需要完整时序序列来评判运动质量；若压缩帧会把视频退化为近乎静态的输入并破坏运动维度。因此这类 reward 保留全部解码帧（数量受 `--num_decoded_latents` 约束）。
- 在 inpaint 模式（`train_mode="inpaint"`）下，所有 reward 在评估前都会丢弃首帧。
> **重要**：使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 这类视频 reward 必须在 `--reward_fn_kwargs` 中设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会彻底切断梯度链。
### 3.4 单 Reward 与多 Reward 训练
单 reward 使用 `--reward_fn="HPSv3Reward"` 配一个扁平的 `--reward_fn_kwargs`。若要组合多个目标，传入逗号分隔的列表、嵌套的 kwargs dict，以及（可选的）权重。每个 reward 的 loss 与 reward 值会分开记录（如 `train/loss_HPSv3Reward`、`train/reward_VideoAlignReward`），而实际优化的 `loss = Σ wᵢ · lossᵢ`。
**多 reward 示例（HPSv3 + VideoAlign，T2V）**——最强的对齐配置，因为视频 reward 能补上图像 reward 无法覆盖的时序（运动/一致性）信号：
```bash
--num_decoded_latents=13 \
--num_reward_frames=2 \
--use_1steps_vae \
--reward_fn="VideoAlignReward,HPSv3Reward" \
--reward_fn_kwargs='{"VideoAlignReward": {"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "fps": 16, "reward_dim": "TA+MQ", "differentiable": true, "use_norm": true, "max_reward": 2.0, "loss_scale": 1.0}, "HPSv3Reward": {"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}}' \
--multi_reward_weights='{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}' \
--backprop \
--backprop_strategy="tail" \
--backprop_num_steps=1
```
### 3.5 可微解码与 `--use_1steps_vae`
默认的分块因果 VAE（`AutoencoderKLWan3_8`）按块应用梯度检查点，从而把可微解码限制在**前 3 个潜帧**（`--num_decoded_latents <= 3`）。不设置 `--use_1steps_vae` 而使用 `--num_decoded_latents > 3` 会抛出 `ValueError`。
启用 `--use_1steps_vae` 可切换到单段式变体，它在一次前向中解码整个 latent、按解码块干净地应用梯度检查点，并解除该限制——当视频 reward 需要更多时序上下文时很有用。它数值等价，且仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。
### 3.6 训练验证
你可以配置验证参数，在训练过程中周期性生成测试视频，以监控对齐进度。在 FSDP 下，`log_validation` 里的 pipeline 前向是一次跨 rank 的集合通信，因此**所有进程都必须进入验证调用**（它刻意**没有**包在 `is_main_process` 里）；可训练专家已带上 LoRA，冻结专家则被复用而非重新加载。
| 参数 | 说明 | 默认值 |
|-----------|-------------|---------------|
| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
| `--validation_epochs` | 每 N 个 epoch 运行一次验证 | 5 |
| `--validation_prompts` | 用于验证视频生成的 prompt | None |
| `--validation_paths` | I2V 验证的参考图路径（仅 inpaint 模式） | None |
**验证示例**：
```bash
--validation_prompts="A stylish woman strolls down a bustling Tokyo street at night, warm neon glow and animated city signage reflecting on the wet pavement; she wears a black leather jacket over a red dress, confident casual demeanor, medium shot, cinematic." \
--validation_steps=50 \
--validation_epochs=500
```
> 验证视频会保存到 `output_dir` 目录。若做 T2V 验证（无参考图），保持 `train_mode="normal"` 并省略 `--validation_paths`。
---
## 四、推理测试
训练好的 reward LoRA 就是一个普通 LoRA，加载方式与 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 完全一致。对于 A14B 双 Transformer 模型，在 [examples/wan2.2/predict_t2v.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.2/predict_t2v.py)（或 `predict_i2v.py`）中，为 low-noise reward LoRA 设置 `lora_path` / `lora_weight`，为 high-noise reward LoRA 设置 `lora_high_path` / `lora_high_weight`：
```python
# 训练好的 reward LoRA 权重路径，例如
# "output_dir_wan2.2_reward_lora/checkpoint-xxx/diffusion_pytorch_model.safetensors"
lora_path = None
lora_high_path = None
lora_weight = 0.55
lora_high_weight = 0.55
```
单卡推理运行：
```bash
python examples/wan2.2/predict_t2v.py
```
显存管理模式与多卡并行推理的详细说明，请参见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 的推理测试章节。
---
## 五、局限性
1. 训练到一定程度后，reward 持续上升，但生成视频的质量不再进一步提升。模型会取巧地学到一些捷径（在背景里添加伪影，即对抗补丁）来抬高 reward。
2. 目前仍缺乏适合视频生成的偏好模型。直接使用图像偏好模型无法沿时序维度评估偏好（如运动性与一致性），并且会降低生成视频的运动性。可以仅用解码视频的第一帧来计算 reward 来部分缓解，但影响依然存在——这也正是对于运动敏感的对齐，推荐使用视频 reward（如 `VideoAlignReward`）并配合 `--use_1steps_vae` 与更大的 `--num_decoded_latents`、或采用多 reward 组合的原因。
---
## 六、更多资源
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
- **标准 LoRA 训练**：[README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)
- **英文版**：[README_TRAIN_REWARD.md](./README_TRAIN_REWARD.md)
---
## 参考文献
<ol>
<li id="ref1">Clark, Kevin, et al. "Directly fine-tuning diffusion models on differentiable rewards.". In ICLR 2024.</li>
<li id="ref2">Prabhudesai, Mihir, et al. "Aligning text-to-image diffusion models with reward backpropagation." arXiv preprint arXiv:2310.03739 (2023).</li>
</ol>
scripts/wan2.2/train_reward_lora.py
+2417
View File
File diff suppressed because it is too large Load Diff
scripts/wan2.2/train_reward_lora.sh
+50
View File
@@ -0,0 +1,50 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 与 NCCL_P2P_DISABLE=1 用于没有 RDMA 的多机环境。
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
export NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False \
scripts/wan2.2/train_reward_lora.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=1 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir_wan2.2_reward_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--enable_bucket \
--boundary_type="high" \
--rank=128 \
--network_alpha=64 \
--target_name="q,k,v,ffn.0,ffn.2" \
--train_mode="normal" \
--guidance_scale=5.0 \
--num_inference_steps=40 \
--num_decoded_latents=3 \
--reward_fn="HPSReward" \
--reward_fn_kwargs='{"version": "v2.1"}' \
--backprop \
--backprop_strategy="tail" \
--backprop_num_steps=1 \
--low_vram
scripts/wan2.2_fun/README_TRAIN_REWARD.md
+365 -25
View File
@@ -1,40 +1,380 @@
# Wan2.2-Fun-Reward-LoRAs
## Introduction
We explore the Reward Backpropagation technique <sup>[1](#ref1) [2](#ref2)</sup> to optimized the generated videos by [Wan2.2-Fun](https://github.com/aigc-apps/VideoX-Fun) for better alignment with human preferences.
We provide the following pre-trained models (i.e. LoRAs) along with [the training script](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/wan2.2_fun/train_reward_lora.py). You can use these LoRAs to enhance the corresponding base model as a plug-in or train your own reward LoRA.
# Wan2.2 Fun Reward LoRA 训练指南
For more details, please refer to our [GitHub repo](https://github.com/aigc-apps/VideoX-Fun).
本文档提供 Wan2.2 Fun **Reward 反向传播（Reward Backpropagation）** 微调的完整流程，即用 LoRA 优化生成视频，使其更符合人类偏好。内容涵盖环境配置、数据准备、Reward 专用训练参数（单/多 reward、可微 VAE 解码、backprop 策略）、训练验证与推理测试。
| Name | Base Model | Reward Model | Hugging Face | Description |
> **说明**：Wan2.2 Fun 是一个基于 Wan2.2 架构的视频修复模型。Wan2.2 采用双 Transformer 架构（高噪声/低噪声模型），5B 版本使用单 Transformer 架构。Reward LoRA 训练会跑一段可微的 rollout（一条简短的去噪轨迹），并把 reward 模型的 loss 反传到可训练的专家上，因此比标准 LoRA 微调更耗显存。本指南支持 A14B 与 5B 两种变体。
我们采用 Reward Backpropagation 技术 <sup>[1](#ref1) [2](#ref2)</sup> 来优化 [Wan2.2-Fun](https://github.com/aigc-apps/VideoX-Fun) 生成的视频，使其更好地对齐人类偏好。我们随[训练脚本](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/wan2.2_fun/train_reward_lora.py)一并提供以下预训练模型（即 LoRA）。你可以用这些 LoRA 作为插件增强对应的基座模型，也可以训练自己的 reward LoRA。
| 名称 | 基座模型 | Reward 模型 | Hugging Face | 说明 |
|--|--|--|--|--|
| Wan2.2-Fun-A14B-InP-high-noise-HPS2.1.safetensors | [Wan2.2-Fun-A14B-InP (high noise)](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/high_noise_model) | [HPS v2.1](https://github.com/tgxs002/HPSv2) | [🤗Link](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-high-noise-HPS2.1.safetensors) | Official HPS v2.1 reward LoRA (`rank=128` and `network_alpha=64`) for Wan2.2-Fun-A14B-InP (high noise). It is trained with a batch size of 8 for 5,000 steps.|
| Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors | [Wan2.2-Fun-A14B-InP (low noise)](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/low_noise_model) | [MPS](https://github.com/Kwai-Kolors/MPS) | [🤗Link](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors) | Official HPS v2.1 reward LoRA (`rank=128` and `network_alpha=64`) for Wan2.2-Fun-A14B-InP (low noise). It is trained with a batch size of 8 for 2,700 steps.|
| Wan2.2-Fun-A14B-InP-high-noise-MPS.safetensors | [Wan2.2-Fun-A14B-InP (high noise)](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/high_noise_model) | [HPS v2.1](https://github.com/tgxs002/HPSv2) | [🤗Link](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-high-noise-MPS.safetensors) | Official MPS reward LoRA (`rank=128` and `network_alpha=64`) for Wan2.2-Fun-A14B-InP (high noise). It is trained with a batch size of 8 for 5,000 steps.|
| Wan2.2-Fun-A14B-InP-low-noise-MPS.safetensors | [Wan2.2-Fun-A14B-InP (low noise)](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/low_noise_model) | [MPS](https://github.com/Kwai-Kolors/MPS) | [🤗Link](https://huggingface.co/alibaba-pai/Wan2.1-Fun-Reward-LoRAs/resolve/main/Wan2.1-Fun-14B-InP-MPS.safetensors) | Official MPS reward LoRA (`rank=128` and `network_alpha=64`) for Wan2.2-Fun-A14B-InP (low noise). It is trained with a batch size of 8 for xxx steps.|
| Wan2.2-Fun-A14B-InP-high-noise-HPS2.1.safetensors | [Wan2.2-Fun-A14B-InP（high noise）](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/high_noise_model) | [HPS v2.1](https://github.com/tgxs002/HPSv2) | [🤗链接](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-high-noise-HPS2.1.safetensors) | 面向 Wan2.2-Fun-A14B-InP（high noise）的官方 HPS v2.1 reward LoRA（`rank=128`、`network_alpha=64`）。 |
| Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors | [Wan2.2-Fun-A14B-InP（low noise）](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/low_noise_model) | [HPS v2.1](https://github.com/tgxs002/HPSv2) | [🤗链接](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors) | 面向 Wan2.2-Fun-A14B-InP（low noise）的官方 HPS v2.1 reward LoRA（`rank=128`、`network_alpha=64`）。 |
| Wan2.2-Fun-A14B-InP-high-noise-MPS.safetensors | [Wan2.2-Fun-A14B-InP（high noise）](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/high_noise_model) | [MPS](https://github.com/Kwai-Kolors/MPS) | [🤗链接](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-high-noise-MPS.safetensors) | 面向 Wan2.2-Fun-A14B-InP（high noise）的官方 MPS reward LoRA（`rank=128`、`network_alpha=64`）。 |
| Wan2.2-Fun-A14B-InP-low-noise-MPS.safetensors | [Wan2.2-Fun-A14B-InP（low noise）](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/low_noise_model) | [MPS](https://github.com/Kwai-Kolors/MPS) | [🤗链接](https://huggingface.co/alibaba-pai/Wan2.1-Fun-Reward-LoRAs/resolve/main/Wan2.1-Fun-14B-InP-MPS.safetensors) | 面向 Wan2.2-Fun-A14B-InP（low noise）的官方 MPS reward LoRA（`rank=128`、`network_alpha=64`）。 |
> [!NOTE]
> We found that, MPS reward LoRA for the low-noise model converges significantly more slowly than on the other models, and may not deliver satisfactory results. Therefore, for the low-noise model, we recommend using HPSv2.1 reward LoRA.
> 我们发现，low-noise 模型的 MPS reward LoRA 收敛速度明显慢于其他模型，且效果可能不理想。因此对于 low-noise 模型，我们推荐使用 HPS v2.1 reward LoRA。
## Demo
Please refer to [here](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs#demo).
## 演示
请参考[此处](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs#demo)。
## Quick Start
Set `lora_path` along with `lora_weight` for the low noise reward LoRA, while specifying `lora_high_path` and `lora_high_weight` for high noise reward LoRA in [examples/wan2.2_fun/predict_t2v.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.1_fun/predict_t2v.py).
---
## Training
The training code is based on [train_lora.py](./train_lora.py). We provide a shell script to train the HPS v2.1 reward LoRA for the low noise model of Wan2.2-Fun-A14B-InP, which can be trained on a single 8*A100 node with 80GB VRAM. To train reward LoRA for the high noise model, Deepspeed Zero3 with CPU offload is required.
## 目录
- [一、环境配置](#一环境配置)
- [二、数据准备](#二数据准备)
- [2.1 快速测试数据集](#21-快速测试数据集)
- [2.2 Reward 模型权重](#22-reward-模型权重)
- [2.3 数据集结构与 metadata.json](#23-数据集结构与-metadatajson)
- [三、Reward LoRA 训练](#三reward-lora-训练)
- [3.1 下载预训练基座模型](#31-下载预训练基座模型)
- [3.2 快速开始（FSDP）](#32-快速开始fsdp)
- [3.3 Reward 专用参数解析](#33-reward-专用参数解析)
- [3.4 单 Reward 与多 Reward 训练](#34-单-reward-与多-reward-训练)
- [3.5 可微解码与 `--use_1steps_vae`](#35-可微解码与---use_1steps_vae)
- [3.6 训练验证](#36-训练验证)
- [四、推理测试](#四推理测试)
- [五、局限性](#五局限性)
- [六、更多资源](#六更多资源)
- [参考文献](#参考文献)
Please refer to [Setup](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/cogvideox_fun/README_TRAIN_REWARD.md#setup) and [Important Args](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/cogvideox_fun/README_TRAIN_REWARD.md#important-args) before training.
---
## 一、环境配置
## Limitations
1. We observe after training to a certain extent, the reward continues to increase, but the quality of the generated videos does not further improve.
The model trickly learns some shortcuts (by adding artifacts in the background, i.e., adversarial patches) to increase the reward.
2. Currently, there is still a lack of suitable preference models for video generation. Directly using image preference models cannot
evaluate preferences along the temporal dimension (such as dynamism and consistency). Further more, We find using image preference models leads to a decrease
in the dynamism of generated videos. Although this can be mitigated by computing the reward using only the first frame of the decoded video, the impact still persists.
**方式 1：使用 requirements.txt**
## Reference
```bash
pip install -r requirements.txt
```
**方式 2：手动安装依赖**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 部分视频 reward 模型（如 `VideoAlignReward`、`HPSv3Reward`）基于 Qwen2-VL 骨干，需要 `transformers==4.51.3`。使用这些 reward 时请安装对应版本。
**方式 3：使用 docker**
使用 docker 的情况下，请保证机器中已经正确安装显卡驱动与 CUDA 环境，然后以此执行以下命令：
```
# pull image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# enter image
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 二、数据准备
### 2.1 快速测试数据集
Reward LoRA 训练从文本 prompt 出发跑可微 rollout，因此标准的 T2V reward 路径只需一个**纯 prompt 数据集**即可。对于 I2V/inpaint 路径，则使用视频数据集来提供参考条件。
```bash
# 下载官方示例数据集
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 Reward 模型权重
不同 reward 的骨干/checkpoint 来源不同：
- **`HPSReward`（HPS v2.1）与 `MPSReward`**：打分头很小，会在**首次运行时自动下载**到 torch hub 缓存（`~/.cache/torch/hub`）；其中 **HPS 还需给 `hpsv2` 内置的 open_clip 补两个特殊配置文件**（见下），MPS 无需额外准备；
- **`VideoAlignReward`（VideoReward）与 `HPSv3Reward`**：基于 Qwen2-VL 骨干，体积较大，**建议提前下载到本地**，再用 `--reward_fn_kwargs` 指向本地路径。
各 reward 的依赖与来源汇总如下（`--reward_fn` 取值见 [3.3](#33-reward-专用参数解析)）：
| Reward | 骨干 / 依赖 | 打分头 / checkpoint | 手动下载 |
|--|--|--|--|
| `HPSReward`（`version="v2.1"`） | open_clip `ViT-H-14`，需 `pip install hpsv2` + `factory.py` 补丁与 BPE 词表 | `HPS_v2.1_compressed.pt`（自动下载） | 否，可用 `model_path` 覆盖 |
| `MPSReward` | `laion/CLIP-ViT-H-14-laion2B-s32B-b79K`（config/tokenizer 自动下载），MPS 代码已内置于 `videox_fun/reward/MPS/` | `MPS_overall.pth`（自动下载） | 否，可用 `model_path` 覆盖 |
| `VideoAlignReward` | `Qwen/Qwen2-VL-2B-Instruct` | `KlingTeam/VideoReward`（含 `model_config.json` 与权重） | **是**，`model_path` 必填 |
| `HPSv3Reward` | `Qwen/Qwen2-VL-7B-Instruct` | `MizzenAI/HPSv3` 的 `HPSv3.safetensors` | 建议，缺省时自动从 HF 下载 |
**HPS v2.1（`HPSReward`）**——除了 `pip install hpsv2`，还要给 `hpsv2` 内置的 open_clip 补两个特殊配置文件：`create_model_and_transforms("ViT-H-14", ...)` 用到了 HPS 分支新增的参数，需要打过补丁的 `factory.py`；`get_tokenizer("ViT-H-14")` 需要 BPE 词表 `bpe_simple_vocab_16e6.txt.gz`。缺任一都会在加载时报错。`ViT-H-14` 骨干与 `HPS_v2.1_compressed.pt` 打分头会在首次运行时自动下载到 torch hub 缓存：
```bash
pip install hpsv2
# 覆盖 hpsv2 内置 open_clip 的 factory.py 补丁，并补上 BPE 词表
site_packages=$(python -c "import site; print(site.getsitepackages()[0])")
wget -O $site_packages/hpsv2/src/open_clip/factory.py https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/package/patches/hpsv2_src_open_clip_factory_patches.py
wget -O $site_packages/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz https://github.com/tgxs002/HPSv2/raw/refs/heads/master/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz
```
```bash
--reward_fn="HPSReward" \
--reward_fn_kwargs='{"version": "v2.1"}'
```
**MPS（`MPSReward`）**——MPS 代码已内置于 `videox_fun/reward/MPS/`，无需打补丁；`MPS_overall.pth` 打分头会自动从 OSS 下载，`laion/CLIP-ViT-H-14-laion2B-s32B-b79K` 的 config 与 tokenizer 会自动从 HuggingFace 下载（若无法访问 HF，可用 `processor_name_or_path` 指向本地镜像）：
```bash
--reward_fn="MPSReward" \
--reward_fn_kwargs='{"max_reward": 1, "loss_scale": 1}'
```
**VideoReward（`VideoAlignReward`）**——需 Qwen2-VL-2B 骨干 + VideoReward 打分头目录（内含 `model_config.json` 与权重），二者都要提前下载，且 `model_path` 必填：
```bash
modelscope download --model Qwen/Qwen2-VL-2B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-2B-Instruct
huggingface-cli download KlingTeam/VideoReward --local-dir models/Personalized_Model/VideoReward
```
```bash
--reward_fn="VideoAlignReward" \
--reward_fn_kwargs='{"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "reward_dim": "TA+MQ", "fps": 16, "use_norm": true, "differentiable": true, "max_reward": 2.0}'
```
**HPSv3（`HPSv3Reward`）**——需 Qwen2-VL-7B 骨干 + `HPSv3.safetensors`；不指定 `checkpoint_path` 时会自动从 HF `MizzenAI/HPSv3` 下载，建议提前落盘：
```bash
modelscope download --model Qwen/Qwen2-VL-7B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-7B-Instruct
modelscope download --model MizzenAI/HPSv3 --local_dir models/Personalized_Model/HPSv3
```
```bash
--reward_fn="HPSv3Reward" \
--reward_fn_kwargs='{"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}'
```
> 使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 必须设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会切断梯度链（见 [3.3](#33-reward-专用参数解析)）。
### 2.3 数据集结构与 metadata.json
```
📦 datasets/
├── 📂 my_dataset/
│ ├── 📂 train/
│ │ ├── 📄 video001.mp4
│ │ ├── 📄 video002.mp4
│ │ └── 📄 ...
│ └── 📄 metadata.json
```
**相对路径格式**（示例）：
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
**关键字段说明**：
- `file_path`：视频路径（相对路径或绝对路径）。
- `text`：视频描述（英文 prompt）。
- `type`：数据类型，固定为 `"video"`。
- `width` / `height`：视频宽高（**强烈建议提供**，用于分桶训练）。可用 `scripts/process_json_add_width_and_height.py` 提取这两个字段。
**相对路径与绝对路径**：
- 相对路径：`export DATASET_NAME="datasets/X-Fun-Videos-Demo/"` 与 `export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"`。
- 绝对路径：`export DATASET_NAME=""` 与 `export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"`。
---
## 三、Reward LoRA 训练
### 3.1 下载预训练基座模型
```bash
# 创建模型目录
mkdir -p models/Diffusion_Transformer
# 下载 Wan2.2 Fun 官方权重
# A14B 模型（双 Transformer 架构）
modelscope download --model PAI/Wan2.2-Fun-A14B-InP --local_dir models/Diffusion_Transformer/Wan2.2-Fun-A14B-InP
# 或 5B 模型（单 Transformer 架构）
# modelscope download --model PAI/Wan2.2-Fun-5B-InP --local_dir models/Diffusion_Transformer/Wan2.2-Fun-5B-InP
```
### 3.2 快速开始（FSDP）
Reward LoRA 训练会把完整的去噪 rollout 与 reward 模型的前向计算图保持在显存中，因此推荐使用 FSDP（`FULL_SHARD`）。当 `--boundary_type="high"` 时，*另一个*（冻结的）专家在训练期已被分片并常驻 GPU，验证阶段会直接复用它，而不会从磁盘再加载一份副本。
**Wan2.2 Fun Reward LoRA 训练示例（单一 HPS v2.1 reward，FSDP）**：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-Fun-A14B-InP"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 与 NCCL_P2P_DISABLE=1 用于没有 RDMA 的多机环境。
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
export NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False \
scripts/wan2.2_fun/train_reward_lora.py \
--config_path="config/wan2.2/wan_civitai_i2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir_wan2.2_fun_reward_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--enable_bucket \
--boundary_type="low" \
--rank=128 \
--network_alpha=64 \
--target_name="q,k,v,ffn.0,ffn.2" \
--num_inference_steps=25 \
--num_decoded_latents=3 \
--reward_fn="HPSReward" \
--reward_fn_kwargs='{"version": "v2.1"}' \
--backprop \
--backprop_strategy="tail" \
--backprop_num_steps=1 \
--train_mode="inpaint" \
--low_vram
```
> 本目录下的 `train_reward_lora.sh` 提供了一份可直接运行的模板。请根据你的环境调整 `MODEL_NAME`、数据集路径、`--reward_fn` 与 `--boundary_type`。
### 3.3 Reward 专用参数解析
除标准 LoRA 参数外（见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)），reward 训练脚本还新增以下参数：
| 参数 | 说明 | 示例值 |
|-----------|-------------|---------------|
| `--reward_fn` | Reward 模型类名，可用**逗号分隔的列表**启用多 reward 训练；各 reward 的可微 loss 以加权和合并。可选：`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`、`VideoAlignReward`。 | `"HPSv3Reward,VideoAlignReward"` |
| `--reward_fn_kwargs` | Reward 模型的关键词参数。单 reward 用**扁平 JSON dict**；多 reward 用**以 reward 名为 key 的嵌套 JSON dict**。 | `'{"VideoAlignReward": {"max_reward": 2.0}, "HPSv3Reward": {"max_reward": 12}}'` |
| `--multi_reward_weights` | 多 reward 各 loss 权重的 JSON dict。会归一化到和为 1；多 reward 时若省略则使用等权重；单 reward 时忽略。 | `'{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}'` |
| `--num_inference_steps` | 可微 rollout 与验证中的去噪步数。 | 25 |
| `--num_decoded_latents` | 解码回像素以供 reward 评估的潜帧数。分块因果 VAE 受梯度检查点限制，除非设置 `--use_1steps_vae`，否则要求 `<= 3`。 | 3 |
| `--num_reward_frames` | 随机选取该数量的解码帧用于 reward 评估。`None` = 使用全部。仅对**图像** reward 生效（见下）。 | 2 |
| `--use_1steps_vae` | 使用单段式 Wan2.2 VAE（`AutoencoderKLWan3_8_1Steps`）做可微解码；数值等价，一次前向解码整个 latent，干净地按解码块应用梯度检查点，并解除 `<= 3` 的限制。仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。 | - |
| `--backprop` | 启用 reward 反向传播训练模式。 | - |
| `--backprop_strategy` | rollout 中哪些步保留梯度图：`last`、`tail`、`uniform`、`random`。 | `tail` |
| `--backprop_num_steps` | backprop 步数，仅当 `--backprop_strategy` 为 `tail`/`uniform`/`random` 时生效。 | 1 |
| `--first_frame_mse_weight` | inpaint 模式下首帧 latent MSE 辅助 loss 的权重。 | 1.0 |
| `--force_t2v` | 当 `train_mode` 为 `"inpaint"`（I2V）或 `"ti2v"` 时，按纯 T2V 训练（条件置零、使用纯 prompt 数据集）。 | - |
**图像 reward 与视频 reward（`is_image_reward`）**：
- **图像 reward**（`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`）逐帧打分，因此帧抽稀（`--num_reward_frames`）是安全的。
- **视频 reward**（`VideoAlignReward`，`is_image_reward=False`）需要完整时序序列来评判运动质量；若压缩帧会把视频退化为近乎静态的输入并破坏运动维度。因此这类 reward 保留全部解码帧（数量受 `--num_decoded_latents` 约束）。
- 在 inpaint 模式下，所有 reward 在评估前都会丢弃首帧（`--force_t2v` 时除外）。
> **重要**：使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 这类视频 reward 必须在 `--reward_fn_kwargs` 中设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会彻底切断梯度链。
### 3.4 单 Reward 与多 Reward 训练
单 reward 使用 `--reward_fn="HPSReward"` 配一个扁平的 `--reward_fn_kwargs`。若要组合多个目标，传入逗号分隔的列表、嵌套的 kwargs dict，以及（可选的）权重。每个 reward 的 loss 与 reward 值会分开记录（如 `train/loss_HPSv3Reward`、`train/reward_VideoAlignReward`），而实际优化的 `loss = Σ wᵢ · lossᵢ`。
**多 reward 示例（HPSv3 + VideoAlign，T2V）**：
```bash
--num_decoded_latents=13 \
--num_reward_frames=2 \
--use_1steps_vae \
--reward_fn="VideoAlignReward,HPSv3Reward" \
--reward_fn_kwargs='{"VideoAlignReward": {"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "fps": 16, "reward_dim": "TA+MQ", "differentiable": true, "use_norm": true, "max_reward": 2.0, "loss_scale": 1.0}, "HPSv3Reward": {"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}}' \
--multi_reward_weights='{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}' \
--backprop \
--backprop_strategy="tail" \
--backprop_num_steps=1
```
### 3.5 可微解码与 `--use_1steps_vae`
默认的分块因果 VAE（`AutoencoderKLWan3_8`）按块应用梯度检查点，从而把可微解码限制在**前 3 个潜帧**（`--num_decoded_latents <= 3`）。不设置 `--use_1steps_vae` 而使用 `--num_decoded_latents > 3` 会抛出 `ValueError`。
启用 `--use_1steps_vae` 可切换到单段式变体，它在一次前向中解码整个 latent、按解码块干净地应用梯度检查点，并解除该限制——当视频 reward 需要更多时序上下文时很有用。它数值等价，且仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。
### 3.6 训练验证
你可以配置验证参数，在训练过程中周期性生成测试视频，以监控对齐进度。在 FSDP 下，`log_validation` 里的 pipeline 前向是一次跨 rank 的集合通信，因此**所有进程都必须进入验证调用**（它刻意**没有**包在 `is_main_process` 里）；可训练专家已带上 LoRA，冻结专家则被复用而非重新加载。
| 参数 | 说明 | 默认值 |
|-----------|-------------|---------------|
| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
| `--validation_epochs` | 每 N 个 epoch 运行一次验证 | 5 |
| `--validation_prompts` | 用于验证视频生成的 prompt | None |
| `--validation_paths` | I2V 验证的参考图路径（仅 inpaint 模式） | None |
**验证示例**：
```bash
--validation_prompts="A brown dog shaking its head, sitting on a light-colored sofa in a cozy room." \
--validation_steps=50 \
--validation_epochs=500
```
> 验证视频会保存到 `output_dir` 目录。若做 T2V 验证（无参考图），请设置 `train_mode="normal"`（或加上 `--force_t2v`）并省略 `--validation_paths`。
---
## 四、推理测试
训练好的 reward LoRA 就是一个普通 LoRA，加载方式与 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 完全一致。对于 A14B 双 Transformer 模型，在 [examples/wan2.2_fun/predict_t2v.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.2_fun/predict_t2v.py)（或 `predict_i2v.py`）中，为 low-noise reward LoRA 设置 `lora_path` / `lora_weight`，为 high-noise reward LoRA 设置 `lora_high_path` / `lora_high_weight`：
```python
# 训练好的 reward LoRA 权重路径，例如
# "output_dir_wan2.2_fun_reward_lora/checkpoint-xxx/diffusion_pytorch_model.safetensors"
lora_path = None
lora_high_path = None
lora_weight = 0.55
lora_high_weight = 0.55
```
单卡推理运行：
```bash
python examples/wan2.2_fun/predict_t2v.py
```
显存管理模式与多卡并行推理的详细说明，请参见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 的推理测试章节。
---
## 五、局限性
1. 训练到一定程度后，reward 持续上升，但生成视频的质量不再进一步提升。模型会取巧地学到一些捷径（在背景里添加伪影，即对抗补丁）来抬高 reward。
2. 目前仍缺乏适合视频生成的偏好模型。直接使用图像偏好模型无法沿时序维度评估偏好（如运动性与一致性），并且会降低生成视频的运动性。可以仅用解码视频的第一帧来计算 reward 来部分缓解，但影响依然存在——这也正是对于运动敏感的对齐，推荐使用视频 reward（如 `VideoAlignReward`）并配合 `--use_1steps_vae` 与更大的 `--num_decoded_latents` 的原因。
---
## 六、更多资源
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
- **标准 LoRA 训练**：[README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)
- **英文版**：[README_TRAIN_REWARD.md](./README_TRAIN_REWARD.md)
---
## 参考文献
<ol>
<li id="ref1">Clark, Kevin, et al. "Directly fine-tuning diffusion models on differentiable rewards.". In ICLR 2024.</li>
<li id="ref2">Prabhudesai, Mihir, et al. "Aligning text-to-image diffusion models with reward backpropagation." arXiv preprint arXiv:2310.03739 (2023).</li>
scripts/wan2.2_fun/README_TRAIN_REWARD_zh-CN.md
+374
View File
@@ -0,0 +1,374 @@
# Wan2.2 Reward LoRA 训练指南
本文档提供 Wan2.2 **Reward 反向传播（Reward Backpropagation）** 微调的完整流程，即用 LoRA 优化生成视频，使其更符合人类偏好。内容涵盖环境配置、数据准备、Reward 专用训练参数（单/多 reward、可微 VAE 解码、backprop 策略）、训练验证与推理测试。
> **说明**：Wan2.2 是一个支持文本到视频（T2V）、图像到视频（I2V）、文本-图像到视频（TI2V）的视频生成模型。A14B 系列采用双 Transformer 架构（高噪声/低噪声模型），TI2V-5B 使用单 Transformer 架构。Reward LoRA 训练会跑一段可微的 rollout（一条简短的去噪轨迹），并把 reward 模型的 loss 反传到可训练的专家上，因此比标准 LoRA 微调更耗显存。本指南覆盖为任意 Wan2.2 变体训练自有 reward LoRA 的流程。
我们采用 Reward Backpropagation 技术 <sup>[1](#ref1) [2](#ref2)</sup> 来优化 [Wan2.2](https://github.com/aigc-apps/VideoX-Fun) 生成的视频，使其更好地对齐人类偏好。训练脚本为 [train_reward_lora.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/wan2.2/train_reward_lora.py)；你可以训练自己的 reward LoRA，并作为插件加载到对应的基座模型上。
> [!NOTE]
> Reward 反向传播容易出现 reward-hacking：reward 持续上升，但感知质量可能停滞甚至退化（背景出现对抗伪影）。相比纯图像 reward，更推荐使用**视频 reward**（如 `VideoAlignReward`）或**多 reward 组合**（如 `HPSv3Reward` + `VideoAlignReward`），以缓解时序运动性的下降——见 [3.4](#34-单-reward-与多-reward-训练) 与 [五、局限性](#五局限性)。
---
## 目录
- [一、环境配置](#一环境配置)
- [二、数据准备](#二数据准备)
- [2.1 快速测试数据集](#21-快速测试数据集)
- [2.2 Reward 模型权重](#22-reward-模型权重)
- [2.3 数据集结构与 metadata.json](#23-数据集结构与-metadatajson)
- [三、Reward LoRA 训练](#三reward-lora-训练)
- [3.1 下载预训练基座模型](#31-下载预训练基座模型)
- [3.2 快速开始（FSDP）](#32-快速开始fsdp)
- [3.3 Reward 专用参数解析](#33-reward-专用参数解析)
- [3.4 单 Reward 与多 Reward 训练](#34-单-reward-与多-reward-训练)
- [3.5 可微解码与 `--use_1steps_vae`](#35-可微解码与---use_1steps_vae)
- [3.6 训练验证](#36-训练验证)
- [四、推理测试](#四推理测试)
- [五、局限性](#五局限性)
- [六、更多资源](#六更多资源)
- [参考文献](#参考文献)
---
## 一、环境配置
**方式 1：使用 requirements.txt**
```bash
pip install -r requirements.txt
```
**方式 2：手动安装依赖**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 部分视频 reward 模型（如 `VideoAlignReward`、`HPSv3Reward`）基于 Qwen2-VL 骨干，需要 `transformers==4.51.3`。使用这些 reward 时请安装对应版本。
**方式 3：使用 docker**
使用 docker 的情况下，请保证机器中已经正确安装显卡驱动与 CUDA 环境，然后以此执行以下命令：
```
# pull image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# enter image
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 二、数据准备
### 2.1 快速测试数据集
Reward LoRA 训练从文本 prompt 出发跑可微 rollout，因此标准的 T2V reward 路径（`train_mode="normal"`）只需一个**纯 prompt 数据集**即可，这也是推荐配置。对于 I2V/inpaint 路径，则使用视频数据集来提供参考条件。
```bash
# 下载官方示例数据集
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 Reward 模型权重
不同 reward 的骨干/checkpoint 来源不同：
- **`HPSReward`（HPS v2.1）与 `MPSReward`**：打分头很小，会在**首次运行时自动下载**到 torch hub 缓存（`~/.cache/torch/hub`）；其中 **HPS 还需给 `hpsv2` 内置的 open_clip 补两个特殊配置文件**（见下），MPS 无需额外准备；
- **`VideoAlignReward`（VideoReward）与 `HPSv3Reward`**：基于 Qwen2-VL 骨干，体积较大，**建议提前下载到本地**，再用 `--reward_fn_kwargs` 指向本地路径。
各 reward 的依赖与来源汇总如下（`--reward_fn` 取值见 [3.3](#33-reward-专用参数解析)）：
| Reward | 骨干 / 依赖 | 打分头 / checkpoint | 手动下载 |
|--|--|--|--|
| `HPSReward`（`version="v2.1"`） | open_clip `ViT-H-14`，需 `pip install hpsv2` + `factory.py` 补丁与 BPE 词表 | `HPS_v2.1_compressed.pt`（自动下载） | 否，可用 `model_path` 覆盖 |
| `MPSReward` | `laion/CLIP-ViT-H-14-laion2B-s32B-b79K`（config/tokenizer 自动下载），MPS 代码已内置于 `videox_fun/reward/MPS/` | `MPS_overall.pth`（自动下载） | 否，可用 `model_path` 覆盖 |
| `VideoAlignReward` | `Qwen/Qwen2-VL-2B-Instruct` | `KlingTeam/VideoReward`（含 `model_config.json` 与权重） | **是**，`model_path` 必填 |
| `HPSv3Reward` | `Qwen/Qwen2-VL-7B-Instruct` | `MizzenAI/HPSv3` 的 `HPSv3.safetensors` | 建议，缺省时自动从 HF 下载 |
**HPS v2.1（`HPSReward`）**——除了 `pip install hpsv2`，还要给 `hpsv2` 内置的 open_clip 补两个特殊配置文件：`create_model_and_transforms("ViT-H-14", ...)` 用到了 HPS 分支新增的参数，需要打过补丁的 `factory.py`；`get_tokenizer("ViT-H-14")` 需要 BPE 词表 `bpe_simple_vocab_16e6.txt.gz`。缺任一都会在加载时报错。`ViT-H-14` 骨干与 `HPS_v2.1_compressed.pt` 打分头会在首次运行时自动下载到 torch hub 缓存：
```bash
pip install hpsv2
# 覆盖 hpsv2 内置 open_clip 的 factory.py 补丁，并补上 BPE 词表
site_packages=$(python -c "import site; print(site.getsitepackages()[0])")
wget -O $site_packages/hpsv2/src/open_clip/factory.py https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/package/patches/hpsv2_src_open_clip_factory_patches.py
wget -O $site_packages/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz https://github.com/tgxs002/HPSv2/raw/refs/heads/master/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz
```
```bash
--reward_fn="HPSReward" \
--reward_fn_kwargs='{"version": "v2.1"}'
```
**MPS（`MPSReward`）**——MPS 代码已内置于 `videox_fun/reward/MPS/`，无需打补丁；`MPS_overall.pth` 打分头会自动从 OSS 下载，`laion/CLIP-ViT-H-14-laion2B-s32B-b79K` 的 config 与 tokenizer 会自动从 HuggingFace 下载（若无法访问 HF，可用 `processor_name_or_path` 指向本地镜像）：
```bash
--reward_fn="MPSReward" \
--reward_fn_kwargs='{"max_reward": 1, "loss_scale": 1}'
```
**VideoReward（`VideoAlignReward`）**——需 Qwen2-VL-2B 骨干 + VideoReward 打分头目录（内含 `model_config.json` 与权重），二者都要提前下载，且 `model_path` 必填：
```bash
modelscope download --model Qwen/Qwen2-VL-2B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-2B-Instruct
huggingface-cli download KlingTeam/VideoReward --local-dir models/Personalized_Model/VideoReward
```
```bash
--reward_fn="VideoAlignReward" \
--reward_fn_kwargs='{"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "reward_dim": "TA+MQ", "fps": 16, "use_norm": true, "differentiable": true, "max_reward": 2.0}'
```
**HPSv3（`HPSv3Reward`）**——需 Qwen2-VL-7B 骨干 + `HPSv3.safetensors`；不指定 `checkpoint_path` 时会自动从 HF `MizzenAI/HPSv3` 下载，建议提前落盘：
```bash
modelscope download --model Qwen/Qwen2-VL-7B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-7B-Instruct
modelscope download --model MizzenAI/HPSv3 --local_dir models/Personalized_Model/HPSv3
```
```bash
--reward_fn="HPSv3Reward" \
--reward_fn_kwargs='{"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}'
```
> 使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 必须设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会切断梯度链（见 [3.3](#33-reward-专用参数解析)）。
### 2.3 数据集结构与 metadata.json
```
📦 datasets/
├── 📂 my_dataset/
│ ├── 📂 train/
│ │ ├── 📄 video001.mp4
│ │ ├── 📄 video002.mp4
│ │ └── 📄 ...
│ └── 📄 metadata.json
```
**相对路径格式**（示例）：
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
**关键字段说明**：
- `file_path`：视频路径（相对路径或绝对路径）。
- `text`：视频描述（英文 prompt）。
- `type`：数据类型，固定为 `"video"`。
- `width` / `height`：视频宽高（**强烈建议提供**，用于分桶训练）。可用 `scripts/process_json_add_width_and_height.py` 提取这两个字段。
**相对路径与绝对路径**：
- 相对路径：`export DATASET_NAME="datasets/X-Fun-Videos-Demo/"` 与 `export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"`。
- 绝对路径：`export DATASET_NAME=""` 与 `export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"`。
---
## 三、Reward LoRA 训练
### 3.1 下载预训练基座模型
```bash
# 创建模型目录
mkdir -p models/Diffusion_Transformer
# 下载 Wan2.2 官方权重
# T2V 模型（文本到视频，双 Transformer）
modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-T2V-A14B
# 或 I2V 模型（图像到视频，双 Transformer）
# modelscope download --model Wan-AI/Wan2.2-I2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-I2V-A14B
# 或 TI2V 模型（文本-图像到视频，单 Transformer）
# modelscope download --model Wan-AI/Wan2.2-TI2V-5B --local_dir models/Diffusion_Transformer/Wan2.2-TI2V-5B
```
### 3.2 快速开始（FSDP）
Reward LoRA 训练会把完整的去噪 rollout 与 reward 模型的前向计算图保持在显存中，因此推荐使用 FSDP（`FULL_SHARD`）。当 `--boundary_type="high"` 时，*另一个*（冻结的）专家在训练期已被分片并常驻 GPU，验证阶段会直接复用它，而不会从磁盘再加载一份副本。
**Wan2.2 T2V Reward LoRA 训练示例（单一 HPS v2.1 reward，FSDP）**：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 与 NCCL_P2P_DISABLE=1 用于没有 RDMA 的多机环境。
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
export NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False \
scripts/wan2.2/train_reward_lora.py \
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=1 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir_wan2.2_reward_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--enable_bucket \
--boundary_type="high" \
--rank=128 \
--network_alpha=64 \
--target_name="q,k,v,ffn.0,ffn.2" \
--train_mode="normal" \
--guidance_scale=5.0 \
--num_inference_steps=40 \
--num_decoded_latents=3 \
--reward_fn="HPSReward" \
--reward_fn_kwargs='{"version": "v2.1"}' \
--backprop \
--backprop_strategy="tail" \
--backprop_num_steps=1 \
--low_vram
```
> 本目录下的 `train_reward_lora.sh` 提供了一份可直接运行的模板。请根据你的环境调整 `MODEL_NAME`、数据集路径、`--reward_fn` 与 `--boundary_type`。
### 3.3 Reward 专用参数解析
除标准 LoRA 参数外（见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)），reward 训练脚本还新增以下参数：
| 参数 | 说明 | 示例值 |
|-----------|-------------|---------------|
| `--reward_fn` | Reward 模型类名，可用**逗号分隔的列表**启用多 reward 训练；各 reward 的可微 loss 以加权和合并。可选：`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`、`VideoAlignReward`。 | `"HPSv3Reward,VideoAlignReward"` |
| `--reward_fn_kwargs` | Reward 模型的关键词参数。单 reward 用**扁平 JSON dict**；多 reward 用**以 reward 名为 key 的嵌套 JSON dict**。 | `'{"VideoAlignReward": {"max_reward": 2.0}, "HPSv3Reward": {"max_reward": 12}}'` |
| `--multi_reward_weights` | 多 reward 各 loss 权重的 JSON dict。会归一化到和为 1；多 reward 时若省略则使用等权重；单 reward 时忽略。 | `'{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}'` |
| `--num_inference_steps` | 可微 rollout 与验证中的去噪步数。 | 40 |
| `--num_decoded_latents` | 解码回像素以供 reward 评估的潜帧数。分块因果 VAE 受梯度检查点限制，除非设置 `--use_1steps_vae`，否则要求 `<= 3`。 | 3 |
| `--num_reward_frames` | 随机选取该数量的解码帧用于 reward 评估。`None` = 使用全部。仅对**图像** reward 生效（见下）。 | 2 |
| `--use_1steps_vae` | 使用单段式 Wan2.2 VAE（`AutoencoderKLWan3_8_1Steps`）做可微解码；数值等价，一次前向解码整个 latent，干净地按解码块应用梯度检查点，并解除 `<= 3` 的限制。仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。 | - |
| `--backprop` | 启用 reward 反向传播训练模式。 | - |
| `--backprop_strategy` | rollout 中哪些步保留梯度图：`last`、`tail`、`uniform`、`random`。 | `tail` |
| `--backprop_num_steps` | backprop 步数，仅当 `--backprop_strategy` 为 `tail`/`uniform`/`random` 时生效。 | 1 |
| `--first_frame_mse_weight` | inpaint 模式下首帧 latent MSE 辅助 loss 的权重。 | 1.0 |
**图像 reward 与视频 reward（`is_image_reward`）**：
- **图像 reward**（`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`）逐帧打分，因此帧抽稀（`--num_reward_frames`）是安全的。
- **视频 reward**（`VideoAlignReward`，`is_image_reward=False`）需要完整时序序列来评判运动质量；若压缩帧会把视频退化为近乎静态的输入并破坏运动维度。因此这类 reward 保留全部解码帧（数量受 `--num_decoded_latents` 约束）。
- 在 inpaint 模式（`train_mode="inpaint"`）下，所有 reward 在评估前都会丢弃首帧。
> **重要**：使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 这类视频 reward 必须在 `--reward_fn_kwargs` 中设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会彻底切断梯度链。
### 3.4 单 Reward 与多 Reward 训练
单 reward 使用 `--reward_fn="HPSv3Reward"` 配一个扁平的 `--reward_fn_kwargs`。若要组合多个目标，传入逗号分隔的列表、嵌套的 kwargs dict，以及（可选的）权重。每个 reward 的 loss 与 reward 值会分开记录（如 `train/loss_HPSv3Reward`、`train/reward_VideoAlignReward`），而实际优化的 `loss = Σ wᵢ · lossᵢ`。
**多 reward 示例（HPSv3 + VideoAlign，T2V）**——最强的对齐配置，因为视频 reward 能补上图像 reward 无法覆盖的时序（运动/一致性）信号：
```bash
--num_decoded_latents=13 \
--num_reward_frames=2 \
--use_1steps_vae \
--reward_fn="VideoAlignReward,HPSv3Reward" \
--reward_fn_kwargs='{"VideoAlignReward": {"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "fps": 16, "reward_dim": "TA+MQ", "differentiable": true, "use_norm": true, "max_reward": 2.0, "loss_scale": 1.0}, "HPSv3Reward": {"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}}' \
--multi_reward_weights='{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}' \
--backprop \
--backprop_strategy="tail" \
--backprop_num_steps=1
```
### 3.5 可微解码与 `--use_1steps_vae`
默认的分块因果 VAE（`AutoencoderKLWan3_8`）按块应用梯度检查点，从而把可微解码限制在**前 3 个潜帧**（`--num_decoded_latents <= 3`）。不设置 `--use_1steps_vae` 而使用 `--num_decoded_latents > 3` 会抛出 `ValueError`。
启用 `--use_1steps_vae` 可切换到单段式变体，它在一次前向中解码整个 latent、按解码块干净地应用梯度检查点，并解除该限制——当视频 reward 需要更多时序上下文时很有用。它数值等价，且仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。
### 3.6 训练验证
你可以配置验证参数，在训练过程中周期性生成测试视频，以监控对齐进度。在 FSDP 下，`log_validation` 里的 pipeline 前向是一次跨 rank 的集合通信，因此**所有进程都必须进入验证调用**（它刻意**没有**包在 `is_main_process` 里）；可训练专家已带上 LoRA，冻结专家则被复用而非重新加载。
| 参数 | 说明 | 默认值 |
|-----------|-------------|---------------|
| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
| `--validation_epochs` | 每 N 个 epoch 运行一次验证 | 5 |
| `--validation_prompts` | 用于验证视频生成的 prompt | None |
| `--validation_paths` | I2V 验证的参考图路径（仅 inpaint 模式） | None |
**验证示例**：
```bash
--validation_prompts="A stylish woman strolls down a bustling Tokyo street at night, warm neon glow and animated city signage reflecting on the wet pavement; she wears a black leather jacket over a red dress, confident casual demeanor, medium shot, cinematic." \
--validation_steps=50 \
--validation_epochs=500
```
> 验证视频会保存到 `output_dir` 目录。若做 T2V 验证（无参考图），保持 `train_mode="normal"` 并省略 `--validation_paths`。
---
## 四、推理测试
训练好的 reward LoRA 就是一个普通 LoRA，加载方式与 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 完全一致。对于 A14B 双 Transformer 模型，在 [examples/wan2.2/predict_t2v.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.2/predict_t2v.py)（或 `predict_i2v.py`）中，为 low-noise reward LoRA 设置 `lora_path` / `lora_weight`，为 high-noise reward LoRA 设置 `lora_high_path` / `lora_high_weight`：
```python
# 训练好的 reward LoRA 权重路径，例如
# "output_dir_wan2.2_reward_lora/checkpoint-xxx/diffusion_pytorch_model.safetensors"
lora_path = None
lora_high_path = None
lora_weight = 0.55
lora_high_weight = 0.55
```
单卡推理运行：
```bash
python examples/wan2.2/predict_t2v.py
```
显存管理模式与多卡并行推理的详细说明，请参见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 的推理测试章节。
---
## 五、局限性
1. 训练到一定程度后，reward 持续上升，但生成视频的质量不再进一步提升。模型会取巧地学到一些捷径（在背景里添加伪影，即对抗补丁）来抬高 reward。
2. 目前仍缺乏适合视频生成的偏好模型。直接使用图像偏好模型无法沿时序维度评估偏好（如运动性与一致性），并且会降低生成视频的运动性。可以仅用解码视频的第一帧来计算 reward 来部分缓解，但影响依然存在——这也正是对于运动敏感的对齐，推荐使用视频 reward（如 `VideoAlignReward`）并配合 `--use_1steps_vae` 与更大的 `--num_decoded_latents`、或采用多 reward 组合的原因。
---
## 六、更多资源
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
- **标准 LoRA 训练**：[README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)
- **英文版**：[README_TRAIN_REWARD.md](./README_TRAIN_REWARD.md)
---
## 参考文献
<ol>
<li id="ref1">Clark, Kevin, et al. "Directly fine-tuning diffusion models on differentiable rewards.". In ICLR 2024.</li>
<li id="ref2">Prabhudesai, Mihir, et al. "Aligning text-to-image diffusion models with reward backpropagation." arXiv preprint arXiv:2310.03739 (2023).</li>
</ol>
scripts/wan2.2_fun/train_reward_lora.py
+1111 -422
View File
File diff suppressed because it is too large Load Diff
scripts/wan2.2_fun/train_reward_lora.sh
+32 -49
View File
@@ -1,65 +1,48 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-Fun-A14B-InP"
export TRAIN_PROMPT_PATH="MovieGenVideoBench_train.txt"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 与 NCCL_P2P_DISABLE=1 用于没有 RDMA 的多机环境。
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
export NCCL_DEBUG=INFO
# Train HPSv2.1 reward LoRA for the low noise model of Wan2.2-Fun-A14B-InP
accelerate launch --mixed_precision="bf16" --num-processes=8 --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json scripts/wan2.2_fun/train_reward_lora.py \
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False \
scripts/wan2.2_fun/train_reward_lora.py \
--config_path="config/wan2.2/wan_civitai_i2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--max_train_steps=10000 \
--checkpointing_steps=100 \
--learning_rate=1e-05 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=1e-04 \
--seed=42 \
--output_dir="output_dir" \
--output_dir="output_dir_wan2.2_fun_reward_lora" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--max_grad_norm=0.3 \
--vae_mini_batch=1 \
--enable_bucket \
--boundary_type="low" \
--lora_skip_name="ffn" \
--low_vram \
--use_deepspeed \
--prompt_path=$TRAIN_PROMPT_PATH \
--train_sample_height=256 \
--train_sample_width=256 \
--num_inference_steps=40 \
--video_length=81 \
--num_decoded_latents=1 \
--rank=128 \
--network_alpha=64 \
--target_name="q,k,v,ffn.0,ffn.2" \
--num_inference_steps=25 \
--num_decoded_latents=3 \
--reward_fn="HPSReward" \
--reward_fn_kwargs='{"version": "v2.1"}' \
--backprop \
--backprop_strategy="tail" \
--backprop_num_steps=1 \
--backprop
# Train MPS reward LoRA for the high noise model of Wan2.2-Fun-A14B-InP
# accelerate launch --mixed_precision="bf16" --num-processes=8 --use_deepspeed --deepspeed_config_file config/zero_stage3_config_cpu_offload.json scripts/wan2.2_fun/train_reward_lora.py \
# --config_path="config/wan2.2/wan_civitai_i2v.yaml" \
# --pretrained_model_name_or_path=$MODEL_NAME \
# --train_batch_size=1 \
# --gradient_accumulation_steps=1 \
# --max_train_steps=10000 \
# --checkpointing_steps=100 \
# --learning_rate=1e-05 \
# --seed=42 \
# --output_dir="output_dir" \
# --gradient_checkpointing \
# --mixed_precision="bf16" \
# --adam_weight_decay=3e-2 \
# --adam_epsilon=1e-10 \
# --max_grad_norm=0.3 \
# --boundary_type="high" \
# --lora_skip_name="ffn" \
# --low_vram \
# --use_deepspeed \
# --prompt_path=$TRAIN_PROMPT_PATH \
# --train_sample_height=256 \
# --train_sample_width=256 \
# --num_inference_steps=40 \
# --video_length=81 \
# --num_decoded_latents=1 \
# --reward_fn="MPSReward" \
# --backprop_strategy="tail" \
# --backprop_num_steps=1 \
# --backprop
--train_mode="inpaint" \
--low_vram