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

Update one-forcing and causal forcing params

Browse Source
This commit is contained in:
bubbliiiing
2026-10-30 16:33:15 +08:00
parent a15b293128
commit f00a7d480e
17 changed files with 4866 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md
+11 -8
View File
@@ -173,7 +173,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
@@ -192,8 +192,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--enable_bucket \
--num_frame_per_block=3 \
--use_kv_cache_training \
--denoising_step_indices_list 1000 667 334 1 \
--real_guidance_scale=6.0 \
--use_ema \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--randomize_step_indices \
--fake_guidance_scale=0.0 \
--gen_update_interval=5 \
@@ -222,7 +224,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--num_train_epochs` | Number of training epochs | 100 |
| `--checkpointing_steps` | Save checkpoint every N steps | 200 |
| `--learning_rate` | Generator learning rate | 2e-06 |
| `--learning_rate_critic` | Critic learning rate | 2e-06 |
| `--learning_rate_critic` | Critic learning rate | 4.0e-07 |
| `--lr_scheduler` | LR scheduler type | `constant_with_warmup` |
| `--lr_warmup_steps` | LR warmup steps | 100 |
| `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_dmd` |
@@ -235,6 +237,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--max_grad_norm` | Gradient clipping threshold | 10.0 |
| `--trainable_modules` | Trainable modules (`"."` = all) | `"."` |
| `--low_vram` | Enable low VRAM mode (offload VAE/text encoder) | - |
| `--use_ema` | Maintain an EMA copy of the generator weights | - |
**Video Sampling Parameters**:
@@ -255,8 +258,8 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| Parameter | Description | Example Value |
|-----------|-------------|---------------|
| `--denoising_step_indices_list` | Denoising step indices (DMD core param). Example runs 4-step; parser default `[1000, 500]` = 2-step | `1000 667 334 1` |
| `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 6.0 |
| `--denoising_step_indices_list` | Denoising step indices (DMD core param). Example runs 4-step; parser default `[1000, 500]` = 2-step | `1000 750 500 250` |
| `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 4.0 |
| `--randomize_step_indices` | Randomize the denoising step indices during training | - |
| `--fake_guidance_scale` | CFG scale for the fake-score (generator). 0.0 = no CFG | 0.0 |
| `--gen_update_interval` | Generator update interval (generator updates every N critic steps) | 5 |
@@ -273,7 +276,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
### 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.
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 remain off, and the DMD launcher additionally enables `--score_timestep_clamp`. 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`.
@@ -290,7 +293,7 @@ The discriminator shares the **fake-score backbone**, with one learnable registe
| `--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 |
| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off (enabled in the DMD launcher) |
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))`.
scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md
+11 -8
View File
@@ -173,7 +173,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
@@ -192,8 +192,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--enable_bucket \
--num_frame_per_block=3 \
--use_kv_cache_training \
--denoising_step_indices_list 1000 667 334 1 \
--real_guidance_scale=6.0 \
--use_ema \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--randomize_step_indices \
--fake_guidance_scale=0.0 \
--gen_update_interval=5 \
@@ -222,7 +224,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--num_train_epochs` | 训练 epoch 数 | 100 |
| `--checkpointing_steps` | 每 N 步保存 checkpoint | 200 |
| `--learning_rate` | 生成器学习率 | 2e-06 |
| `--learning_rate_critic` | 判别器学习率 | 2e-06 |
| `--learning_rate_critic` | 判别器学习率 | 4.0e-07 |
| `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
| `--lr_warmup_steps` | 学习率预热步数 | 100 |
| `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_dmd` |
@@ -235,6 +237,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--max_grad_norm` | 梯度裁剪阈值 | 10.0 |
| `--trainable_modules` | 可训练模块（`"."` = 全部） | `"."` |
| `--low_vram` | 低显存模式（卸载 VAE/文本编码器） | - |
| `--use_ema` | 维护生成器权重的 EMA 副本 | - |
**视频采样参数**：
@@ -255,8 +258,8 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| 参数 | 说明 | 示例值 |
|------|------|--------|
| `--denoising_step_indices_list` | 去噪步骤索引（DMD 核心参数）。示例为 4 步；解析器默认 `[1000, 500]` = 2 步 | `1000 667 334 1` |
| `--real_guidance_scale` | real-score（14B teacher）的 CFG scale | 6.0 |
| `--denoising_step_indices_list` | 去噪步骤索引（DMD 核心参数）。示例为 4 步；解析器默认 `[1000, 500]` = 2 步 | `1000 750 500 250` |
| `--real_guidance_scale` | real-score（14B teacher）的 CFG scale | 4.0 |
| `--randomize_step_indices` | 训练时随机化去噪步骤索引 | - |
| `--fake_guidance_scale` | fake-score（生成器）的 CFG scale。0.0 = 无 CFG | 0.0 |
| `--gen_update_interval` | 生成器更新间隔（每 N 步判别器更新后更新 1 次生成器） | 5 |
@@ -272,7 +275,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
### 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 损失。
这些选项复用 `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 仍关闭，DMD 的 launcher 额外启用 `--score_timestep_clamp`。AR diffusion 和 CCD 不添加 GAN 损失。
判别器共享 **fake-score backbone**，每个选定特征层对应一个可学习 register token。判别器头由 critic optimizer 更新，不加入 generator optimizer。生成器损失为 `DMD + 加权 GAN`，critic 损失为 `denoising + 加权 GAN + R1 + R2`。
@@ -289,7 +292,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--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` | 关闭 |
| `--score_timestep_clamp` | 将 DMD、生成器 GAN 和 critic 时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关闭（DMD launcher 已启用） |
设 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))`。
scripts/wan2.1_causal_forcing/train_ar_diffusion.sh
+1
View File
@@ -37,6 +37,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 \
scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh
+1
View File
@@ -41,6 +41,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 \
scripts/wan2.1_causal_forcing/train_causal_dmd.sh
+6 -3
View File
@@ -38,13 +38,14 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--num_train_epochs=100 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--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 \
@@ -56,8 +57,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--enable_bucket \
--num_frame_per_block=3 \
--use_kv_cache_training \
--denoising_step_indices_list 1000 667 334 1 \
--real_guidance_scale=6.0 \
--use_ema \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--randomize_step_indices \
--fake_guidance_scale=0.0 \
--gen_update_interval=5 \
scripts/wan2.1_one_forcing/README_ONE_FORCING.md
+356
View File
@@ -0,0 +1,356 @@
# Wan2.1 One-Forcing: Stable One-Step Autoregressive Video Generation
This document provides the complete workflow for **One-Forcing** on Wan2.1-T2V-1.3B, a recipe ported from [Aurora-edu/One-Forcing](https://github.com/Aurora-edu/One-Forcing).
> **What is One-Forcing?**
>
> One-Forcing (*Towards Stable One-Step Autoregressive Video Generation*) is **Causal-Forcing Stage 3 — Distribution Matching Distillation (DMD)** *plus* a **shared noised-latent adversarial critic**. The warm-start generator is produced by the first two Causal-Forcing stages, and One-Forcing then polishes it into a **stable one-step, frame-wise autoregressive** video model:
>
> 1. **Stage 1 — AR Diffusion** (`scripts/wan2.1_causal_forcing/train_ar_diffusion.py`): Train the model with **teacher forcing** on clean video latents, learning to denoise chunk-by-chunk (or frame-by-frame) in a causal autoregressive manner. This produces a strong AR backbone. See [README_TRAIN_AR_DIFFUSION.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md).
> 2. **Stage 2 — Causal Consistency Distillation (CCD)** (`scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py`): Distill the multi-step AR model into a **one-step-per-block** consistency model using an EMA teacher with CFG guidance. Its output checkpoint is the **warm-start generator** for One-Forcing. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md).
> 3. **Stage 3 — One-Forcing** (`train_one_forcing.py`, this document), which adds two mechanisms on top of DMD:
> - **DMD core**: the generator matches the real distribution through a 14B real-score teacher and a 1.3B fake-score critic (`train_causal_dmd.py`).
> - **Adversarial critic**: a discriminator head rides on the **fake-score backbone** (one register token per selected feature layer). Both the DMD denoising loss and the adversarial loss operate on the **same noised generated latent**, and the critic uses **one shared noise/timestep** for fake and real samples.
> - **One step**: `denoising_step_list = [1000]` — each frame is generated in a single denoising step (the "One" in One-Forcing).
>
> Instead of training Stage 1/2 yourself, you can also warm-start from the **official ODE-initialized checkpoint** released by the One-Forcing authors (see Section 3).
>
> **The trainer** `scripts/wan2.1_one_forcing/train_one_forcing.py` implements the One-Forcing generator/critic GAN losses, the relativistic option, and finite-difference R1/R2 penalties. For the underlying DMD stage, see [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md).
---
## Table of Contents
- [1. Prerequisites](#1-prerequisites)
- [2. Environment Setup](#2-environment-setup)
- [3. Download Pretrained Models](#3-download-pretrained-models)
- [4. Prepare Training Data](#4-prepare-training-data)
- [4.1 Quick Demo Dataset](#41-quick-demo-dataset)
- [4.2 Dataset Structure](#42-dataset-structure)
- [4.3 metadata.json Format](#43-metadatajson-format)
- [5. Training](#5-training)
- [5.1 Quick Start](#51-quick-start)
- [5.2 Recipe Parameters](#52-recipe-parameters)
- [5.3 GAN Prerequisites](#53-gan-prerequisites)
- [5.4 Optional Regularizers (Relativistic / R1 / R2)](#54-optional-regularizers-relativistic--r1--r2)
- [6. Inference](#6-inference)
- [6.1 Quick Start](#61-quick-start)
- [6.2 1-Step vs 4-Step Selector](#62-1-step-vs-4-step-selector)
- [6.3 First-Frame Schedule](#63-first-frame-schedule)
- [7. Checkpoint Round-Trip](#7-checkpoint-round-trip)
- [8. Additional Resources](#8-additional-resources)
---
## 1. Prerequisites
One-Forcing is the final polish on top of a distilled generator, so it needs:
1. A **warm-start generator checkpoint** — a distilled/ODE generator that the DMD + adversarial critic then polishes (e.g. a **Causal-Forcing Stage-2 CCD** checkpoint, or the official ODE-initialized checkpoint in Section 3).
2. A **Wan2.1-T2V-14B** model as the DMD real-score teacher.
3. **Captioned real videos** — the adversarial critic needs clean real latents, so a prompt-only JSON is **not** sufficient (unlike plain DMD).
---
## 2. 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
```
**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
```
---
## 3. Download Pretrained Models
One-Forcing requires **two** pretrained models:
- **Wan2.1-T2V-1.3B**: base model for the generator/critic.
- **Wan2.1-T2V-14B**: the non-causal real-score teacher used by DMD to compute the real distribution score.
```bash
# Create model directory
mkdir -p models/Diffusion_Transformer
# Download 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
# Download Wan2.1 T2V 14B (DMD real-score teacher)
modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B
```
The official One-Forcing checkpoints are hosted on HuggingFace (`JiaqiFeng/OneForcing`):
```bash
# Warm-start generator for training (the official ODE-initialized checkpoint),
# loaded via --ode_transformer_path.
hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
# Trained One-Forcing generator for inference (see Section 6).
hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
```
If you trained Causal-Forcing Stage 1/2 yourself, you can instead point `--ode_transformer_path` at your own Stage-2 CCD checkpoint.
---
## 4. Prepare Training Data
One-Forcing enables the GAN branch by default (`--gan_g_weight=0.03 --gan_d_weight=0.03`), so the data loader always uses an **ImageVideoDataset**: the adversarial critic needs **captioned real videos** to build its clean latents. A prompt-only JSON is **not** sufficient (unlike plain DMD).
### 4.1 Quick Demo Dataset
```bash
# Download official demo dataset
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 4.2 Dataset Structure
Same as Causal-Forcing Stage 1/2. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md#42-dataset-structure) for details.
### 4.3 metadata.json Format
Each entry needs `file_path`, `type="video"`, and `text`; the launcher uses the `metadata_add_width_height.json` variant, which additionally carries `width`/`height` so bucketing does not need to decode the media files:
```json
[
{
"file_path": "train/00000000.mp4",
"text": "A young woman gently turns her head to the right, tilting it slightly as she offers a soft greeting...",
"type": "video",
"width": 832,
"height": 480
}
]
```
> **Note**: Keep enough frames per video for `--video_sample_stride=1` and `--video_sample_n_frames=81`: the GAN branch aborts when the source yields fewer latent frames than `score_num_frames=21` (see Section 5.3).
---
## 5. Training
### 5.1 Quick Start
The ready-to-use launcher is [train_one_forcing.sh](./train_one_forcing.sh). It calls the One-Forcing trainer (`train_one_forcing.py`) with the GAN recipe enabled:
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
accelerate launch --mixed_precision="bf16" --use_fsdp \
--fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
--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.1_one_forcing/train_one_forcing.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--ode_transformer_path=$STAGE2_CKPT \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--fix_sample_size 480 832 \
--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 \
--max_train_steps=200 \
--checkpointing_steps=100 \
--learning_rate=1.0e-05 \
--learning_rate_critic=1.0e-05 \
--lr_scheduler="constant" \
--seed=42 \
--output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.01 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--num_frame_per_block=1 \
--use_kv_cache_training \
--use_ema \
--ema_decay=0.99 \
--denoising_step_indices_list 1000 \
--real_guidance_scale=6.0 \
--score_timestep_clamp \
--fake_guidance_scale=0.0 \
--gen_update_interval=5 \
--gan_g_weight=0.03 \
--gan_d_weight=0.03 \
--gan_feature_layers 21 29 \
--gan_block_ffn_dim=2048 \
--gan_block_num_heads=12 \
--resume_from_checkpoint="latest" \
--trainable_modules "."
```
Or simply:
```bash
bash scripts/wan2.1_one_forcing/train_one_forcing.sh
```
### 5.2 Recipe Parameters
The One-Forcing recipe sets these launcher flags:
| Parameter | Launcher flag | Value |
|------------------------|---------------|-------|
| `denoising_step_list: [1000]` | `--denoising_step_indices_list` | `1000` (one step) |
| `num_frame_per_block: 1` | `--num_frame_per_block` | `1` (frame-wise) |
| `num_training_frames: 21` / `score_num_frames: 21` | `--video_sample_n_frames=81` @ 480×832 | 21 latent frames |
| `guidance_scale: 5.0` | `--real_guidance_scale` | `6.0` (see CFG note below) |
| `lr: 1e-5` / `lr_critic: 1e-5` | `--learning_rate` / `--learning_rate_critic` | `1.0e-05` |
| `beta1: 0.0`, `beta2: 0.999`, `weight_decay: 0.01` | `--adam_beta1`, `--adam_beta2`, `--adam_weight_decay` | `0.0`, `0.999`, `0.01` |
| `ema_weight: 0.99` | `--use_ema`, `--ema_decay` | on, `0.99` |
| `max_steps: 200` | `--max_train_steps` | `200` |
| `dfake_gen_update_ratio: 5` | `--gen_update_interval` | `5` |
| `gan_feature_layers: [21, 29]` | `--gan_feature_layers` | `21 29` |
| `gan_block_ffn_dim: 2048` | `--gan_block_ffn_dim` | `2048` |
| `gan_block_num_heads: 12` | `--gan_block_num_heads` | `12` |
| `gan_g_weight: 0.03`, `gan_d_weight: 0.03` | `--gan_g_weight`, `--gan_d_weight` | `0.03`, `0.03` |
| `relativistic_discriminator: false` | `--relativistic_discriminator` | omitted (off) |
| `r1_weight: 0.0`, `r2_weight: 0.0` | `--r1_weight`, `--r2_weight` | omitted (off) |
| `context_noise: 0` | `--context_noise` | `0` (default) |
> **CFG convention**: this repo computes guidance as `cond + scale * (cond - uncond)`, and this recipe uses `--real_guidance_scale=6.0`.
For the general training/video-sampling parameters shared with the DMD stage, see [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md#52-key-parameters) Section 5.2.
### 5.3 GAN Prerequisites
The GAN branch is validated at startup in `train_one_forcing.py`. It requires:
- **Captioned real videos** (`type="video"` + `file_path` + `text`), not a prompt-only JSON. Enabling any positive GAN/R1/R2 weight switches the loader from `TextDataset` to `ImageVideoDataset`.
- **`--gradient_accumulation_steps=1`** and **no DeepSpeed** — the trainer raises `ValueError` otherwise.
- **`--video_sample_stride=1`**: with `stride=2`, an 81-frame source yields ~10 latent frames, which is below `score_num_frames=21` and aborts.
- **Valid feature-layer indices** for the fake-score backbone. The defaults `21 29` (layers) and `12` (heads) fit Wan2.1-T2V-1.3B; the head count must divide the fake-score hidden dimension.
Trainable fake-score and head master weights stay **FP32**; the forward pass uses the configured autocast/FSDP precision, and logits/softplus/finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient-communication precision only.
### 5.4 Optional Regularizers (Relativistic / R1 / R2)
These regularizers are **off** by default. To experiment, append to the launcher:
```bash
# Relativistic discriminator: use paired real/fake logit differences
--relativistic_discriminator
# Finite-difference gradient penalties (APT form, not double-backward)
--r1_weight=1.0 --r1_sigma=0.1 # on noisy real latents: r1_weight * mean((D_real' - D_real)^2)
--r2_weight=0.001 --r2_sigma=0.01 # on noisy gen latents: 0.5 * r2_weight * mean(((D_fake' - D_fake)/sigma)^2)
```
For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` (generator) and `softplus(-D_real) + softplus(D_fake)` (discriminator); relativistic losses use the logit differences instead. Real/fake samples share timestep and noise; perturbations use independent noise.
---
## 6. Inference
### 6.1 Quick Start
The inference entry is [examples/wan2.1_one_forcing/predict_t2v.py](../../examples/wan2.1_one_forcing/predict_t2v.py). It loads the **1.3B student** (the 14B model is only the DMD teacher at training time) and runs frame-wise one-step generation. By default it uses the **official One-Forcing checkpoint**:
```bash
hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
```
```python
# examples/wan2.1_one_forcing/predict_t2v.py
model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
transformer_path = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
use_ema = False # only applies to trainer-output checkpoint dirs
stochastic_sampling = True # DMD-distilled student
guidance_scale = 1.0 # CFG is baked into the distilled weights
num_inference_steps = 1 # One-Forcing one-step generation
num_frame_per_block = 1 # frame-wise
sample_size = [480, 832]
video_length = 81
```
Run it:
```bash
python examples/wan2.1_one_forcing/predict_t2v.py
```
`transformer_path` may be the official checkpoint file, a plain `.pt`/`.safetensors` file, or a trainer-output **checkpoint dir** (its live/EMA weights are auto-resolved). The GAN discriminator head is a training-only submodule of fake_score, so inference needs **no discriminator wiring**.
### 6.2 1-Step vs 4-Step Selector
With `stochastic_sampling=True`, `num_inference_steps` selects the warped timestep schedule:
| `num_inference_steps` | Denoise timesteps | Use case |
|-----------------------|-------------------|----------|
| `1` (One-Forcing default) | `[1000]` | Fastest; frame-wise one-step generation |
| `4` | `[1000, 750, 500, 250]` | More stable, slower |
### 6.3 First-Frame Schedule
The shared `WanSelfForcingPipeline` applies **one denoising schedule to all blocks**. A separate first-frame schedule (e.g. `[1000, 750, 500, 250]` for the first frame and `[1000]` for later frames) is therefore not applied; the pipeline is shared by multiple examples and is intentionally left unmodified.
---
## 7. Checkpoint Round-Trip
The FSDP trainer writes each `checkpoint-{step}/` as follows, so training output feeds inference with no key surgery:
| Component | Saved to | Contains `gan_head`? |
|-----------|----------|----------------------|
| **Generator** (live) | `checkpoint-{step}/diffusion_pytorch_model.safetensors` (root, via FSDP hook) | No |
| **Generator** (EMA) | `checkpoint-{step}/transformer_ema/` | No |
| **Fake-score + discriminator head** | `checkpoint-{step}/fake_score/` | Yes (FP32 head tensors) |
- The **generator never contains `gan_head`** keys (the head lives on fake-score), so `predict_t2v.py` loads it cleanly with `strict=False`.
- `_resolve_transformer_path` in the inference script resolves, in order: `transformer_ema/`, `ema_transformer/` (compat alias), `transformer/`, then the root safetensors.
- **Resume** GAN training with the same GAN/head configuration: the head is recreated from CLI flags before loading and restored from `fake_score/`.
---
## 8. Additional Resources
- **One-Forcing (source recipe)**: https://github.com/Aurora-edu/One-Forcing
- **Causal-Forcing (DMD stage)**: https://github.com/thu-ml/Causal-Forcing
- **Stage 1 (AR Diffusion) training guide**: [README_TRAIN_AR_DIFFUSION.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md)
- **Stage 2 (CCD) training guide**: [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md)
- **Underlying DMD trainer guide**: [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md)
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
scripts/wan2.1_one_forcing/README_ONE_FORCING_zh-CN.md
+355
View File
@@ -0,0 +1,355 @@
# Wan2.1 One-Forcing：稳定的单步自回归视频生成
本文档提供在 Wan2.1-T2V-1.3B 上进行 **One-Forcing** 的完整工作流，该配方移植自 [Aurora-edu/One-Forcing](https://github.com/Aurora-edu/One-Forcing)。
> **什么是 One-Forcing？**
>
> One-Forcing（*Towards Stable One-Step Autoregressive Video Generation*）= **Causal-Forcing Stage 3 — 分布匹配蒸馏 (DMD)** *加上*一个**共享噪声 latent 的对抗判别器（critic）**。热启动生成器由 Causal-Forcing 的前两个阶段产出，One-Forcing 再将其打磨为**稳定的单步、逐帧自回归**视频模型：
>
> 1. **Stage 1 — AR Diffusion**（`scripts/wan2.1_causal_forcing/train_ar_diffusion.py`）：在干净视频 latent 上以 **teacher forcing** 训练模型，学习按块（或按帧）因果自回归地去噪，产出强 AR 基座。参见 [README_TRAIN_AR_DIFFUSION_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md)。
> 2. **Stage 2 — 因果一致性蒸馏 (CCD)**（`scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py`）：使用 EMA teacher + CFG 将多步 AR 模型蒸馏为**每块一步**的一致性模型，其产出 checkpoint 即 One-Forcing 的**热启动生成器**。参见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md)。
> 3. **Stage 3 — One-Forcing**（`train_one_forcing.py`，即本文档），在 DMD 之上增加两个机制：
> - **DMD 核心**：生成器通过 14B real-score teacher 与 1.3B fake-score critic 匹配真实分布（`train_causal_dmd.py`）。
> - **对抗 critic**：判别器头挂在 **fake-score backbone** 上（每个选定特征层对应一个 register token）。DMD 去噪损失与对抗损失作用在**同一个噪声生成 latent** 上，critic 对真实/生成样本使用**同一份共享噪声与时间步**。
> - **单步**：`denoising_step_list = [1000]` —— 每帧在单个去噪步内生成（即 One-Forcing 中的 “One”）。
>
> 你也可以不自己训练 Stage 1/2，而是直接使用 One-Forcing 作者发布的**官方 ODE 初始化 checkpoint** 作为热启动权重（见第三节）。
>
> **训练器** `scripts/wan2.1_one_forcing/train_one_forcing.py` 实现了 One-Forcing 的生成器/critic GAN 损失、相对论（relativistic）选项以及有限差分 R1/R2 正则。底层 DMD 阶段请参见 [README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md)。
---
## 目录
- [一、前置条件](#一前置条件)
- [二、环境配置](#二环境配置)
- [三、下载预训练模型](#三下载预训练模型)
- [四、数据准备](#四数据准备)
- [4.1 快速测试数据集](#41-快速测试数据集)
- [4.2 数据集结构](#42-数据集结构)
- [4.3 metadata.json 格式](#43-metadatajson-格式)
- [五、训练](#五训练)
- [5.1 快速开始](#51-快速开始)
- [5.2 配方参数](#52-配方参数)
- [5.3 GAN 前置要求](#53-gan-前置要求)
- [5.4 可选正则（相对论 / R1 / R2）](#54-可选正则相对论--r1--r2)
- [六、推理](#六推理)
- [6.1 快速开始](#61-快速开始)
- [6.2 单步与 4 步选择器](#62-单步与-4-步选择器)
- [6.3 首帧调度](#63-首帧调度)
- [七、Checkpoint 往返（训练→推理）](#七checkpoint-往返训练推理)
- [八、更多资源](#八更多资源)
---
## 一、前置条件
One-Forcing 是在已蒸馏生成器之上的最后打磨，因此需要：
1. **热启动生成器 checkpoint** —— 一个已蒸馏/ODE 生成器，由 DMD + 对抗 critic 进一步打磨（例如 **Causal-Forcing Stage-2 CCD** checkpoint，或第三节的官方 ODE 初始化 checkpoint）。
2. **Wan2.1-T2V-14B** 模型，作为 DMD real-score teacher。
3. **带描述的真实视频** —— 对抗 critic 需要干净的真实 latent，因此仅有提示词的 JSON **不够用**（这一点与普通 DMD 不同）。
---
## 二、环境配置
**方式 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
```
**方式 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
```
---
## 三、下载预训练模型
One-Forcing 需要**两个**预训练模型：
- **Wan2.1-T2V-1.3B**：生成器/判别器的基础模型。
- **Wan2.1-T2V-14B**：非因果的 real-score teacher，用于计算真实分布得分。
```bash
# Create the model directory
mkdir -p models/Diffusion_Transformer
# 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
# 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
```
One-Forcing 官方 checkpoint 托管在 HuggingFace（`JiaqiFeng/OneForcing`）：
```bash
# 训练用热启动生成器（官方 ODE 初始化 checkpoint），通过 --ode_transformer_path 加载。
hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
# 推理用已训练 One-Forcing 生成器（见第六节）。
hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
```
如果你自己训练了 Causal-Forcing Stage 1/2，也可以将 `--ode_transformer_path` 指向你产出的 Stage-2 CCD checkpoint。
---
## 四、数据准备
One-Forcing 默认开启 GAN 分支（`--gan_g_weight=0.03 --gan_d_weight=0.03`），因此数据加载始终使用 **ImageVideoDataset**：对抗 critic 需要**带描述的真实视频**来构建干净 latent，仅有提示词的 JSON **不够用**（这一点与普通 DMD 不同）。
### 4.1 快速测试数据集
```bash
# Download the official demo dataset
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 4.2 数据集结构
与 Causal-Forcing Stage 1/2 相同。详见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md#42-数据集结构)。
### 4.3 metadata.json 格式
每条记录需要 `file_path`、`type="video"` 与 `text`；启动脚本使用 `metadata_add_width_height.json` 变体，额外携带 `width`/`height`，分桶训练无需现场解码媒体文件：
```json
[
{
"file_path": "train/00000000.mp4",
"text": "A young woman gently turns her head to the right, tilting it slightly as she offers a soft greeting...",
"type": "video",
"width": 832,
"height": 480
}
]
```
> **说明**：每条视频须有足够的帧数以满足 `--video_sample_stride=1` 与 `--video_sample_n_frames=81`：当源视频产出的 latent 帧数少于 `score_num_frames=21` 时，GAN 分支会中止（见 5.3）。
---
## 五、训练
### 5.1 快速开始
可直接使用的启动脚本为 [train_one_forcing.sh](./train_one_forcing.sh)。它调用 One-Forcing 训练器（`train_one_forcing.py`）并开启 GAN 配方：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
accelerate launch --mixed_precision="bf16" --use_fsdp \
--fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
--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.1_one_forcing/train_one_forcing.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--ode_transformer_path=$STAGE2_CKPT \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--fix_sample_size 480 832 \
--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 \
--max_train_steps=200 \
--checkpointing_steps=100 \
--learning_rate=1.0e-05 \
--learning_rate_critic=1.0e-05 \
--lr_scheduler="constant" \
--seed=42 \
--output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.01 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--num_frame_per_block=1 \
--use_kv_cache_training \
--use_ema \
--ema_decay=0.99 \
--denoising_step_indices_list 1000 \
--real_guidance_scale=6.0 \
--score_timestep_clamp \
--fake_guidance_scale=0.0 \
--gen_update_interval=5 \
--gan_g_weight=0.03 \
--gan_d_weight=0.03 \
--gan_feature_layers 21 29 \
--gan_block_ffn_dim=2048 \
--gan_block_num_heads=12 \
--resume_from_checkpoint="latest" \
--trainable_modules "."
```
或直接运行：
```bash
bash scripts/wan2.1_one_forcing/train_one_forcing.sh
```
### 5.2 配方参数
One-Forcing 配方设置以下 launcher 参数：
| 参数 | Launcher 参数 | 取值 |
|--------------------|---------------|------|
| `denoising_step_list: [1000]` | `--denoising_step_indices_list` | `1000`（单步） |
| `num_frame_per_block: 1` | `--num_frame_per_block` | `1`（逐帧） |
| `num_training_frames: 21` / `score_num_frames: 21` | `--video_sample_n_frames=81` @ 480×832 | 21 个 latent 帧 |
| `guidance_scale: 5.0` | `--real_guidance_scale` | `6.0`（见下方 CFG 说明） |
| `lr: 1e-5` / `lr_critic: 1e-5` | `--learning_rate` / `--learning_rate_critic` | `1.0e-05` |
| `beta1: 0.0`、`beta2: 0.999`、`weight_decay: 0.01` | `--adam_beta1`、`--adam_beta2`、`--adam_weight_decay` | `0.0`、`0.999`、`0.01` |
| `ema_weight: 0.99` | `--use_ema`、`--ema_decay` | 开启、`0.99` |
| `max_steps: 200` | `--max_train_steps` | `200` |
| `dfake_gen_update_ratio: 5` | `--gen_update_interval` | `5` |
| `gan_feature_layers: [21, 29]` | `--gan_feature_layers` | `21 29` |
| `gan_block_ffn_dim: 2048` | `--gan_block_ffn_dim` | `2048` |
| `gan_block_num_heads: 12` | `--gan_block_num_heads` | `12` |
| `gan_g_weight: 0.03`、`gan_d_weight: 0.03` | `--gan_g_weight`、`--gan_d_weight` | `0.03`、`0.03` |
| `relativistic_discriminator: false` | `--relativistic_discriminator` | 不传（关闭） |
| `r1_weight: 0.0`、`r2_weight: 0.0` | `--r1_weight`、`--r2_weight` | 不传（关闭） |
| `context_noise: 0` | `--context_noise` | `0`（默认） |
> **CFG 约定**：本仓库的引导计算为 `cond + scale * (cond - uncond)`，本配方取 `--real_guidance_scale=6.0`。
与 DMD 阶段共用的通用训练/视频采样参数，参见 [README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md#52-关键参数) 5.2 节。
### 5.3 GAN 前置要求
GAN 分支在 `train_one_forcing.py` 启动时会被校验，要求：
- **带描述的真实视频**（`type="video"` + `file_path` + `text`），不能只提供提示词 JSON。任一 GAN/R1/R2 权重为正时，数据加载会从 `TextDataset` 切换为 `ImageVideoDataset`。
- **`--gradient_accumulation_steps=1`** 且**不使用 DeepSpeed** —— 否则训练器抛出 `ValueError`。
- **`--video_sample_stride=1`**：当 `stride=2` 时，81 帧源仅得到约 10 个 latent 帧，低于 `score_num_frames=21` 而中止。
- fake-score backbone 上**有效的特征层索引**。默认 `21 29`（层）与 `12`（头）适用于 Wan2.1-T2V-1.3B；头数须整除 fake-score 隐藏维度。
可训练的 fake-score 与判别器头 master weights 保持 **FP32**；前向按配置使用 autocast/FSDP 混合精度，logits/softplus/有限差分归约在 FP32 下计算。`--fsdp_reduce_fp32` 只改变梯度通信精度。
### 5.4 可选正则（相对论 / R1 / R2）
这些正则默认**关闭**。若要实验，向 launcher 追加：
```bash
# Relativistic discriminator: use paired real/fake logit differences
--relativistic_discriminator
# Finite-difference gradient penalties (APT form, not double-backward)
--r1_weight=1.0 --r1_sigma=0.1 # on noisy real latents: r1_weight * mean((D_real' - D_real)^2)
--r2_weight=0.001 --r2_sigma=0.01 # on noisy gen latents: 0.5 * r2_weight * mean(((D_fake' - D_fake)/sigma)^2)
```
设 logits 为 `D_fake` 与 `D_real`，普通 GAN 损失为 `softplus(-D_fake)`（生成器）与 `softplus(-D_real) + softplus(D_fake)`（判别器）；相对论模式改用 logit 差值。真实/生成样本共用时间步与噪声，扰动噪声独立抽样。
---
## 六、推理
### 6.1 快速开始
推理入口为 [examples/wan2.1_one_forcing/predict_t2v.py](../../examples/wan2.1_one_forcing/predict_t2v.py)。它加载 **1.3B 学生模型**（14B 仅在训练时作为 DMD teacher），运行逐帧单步生成。默认使用 **One-Forcing 官方 checkpoint**：
```bash
hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
```
```python
# examples/wan2.1_one_forcing/predict_t2v.py
model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
transformer_path = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
use_ema = False # only applies to trainer-output checkpoint dirs
stochastic_sampling = True # DMD-distilled student
guidance_scale = 1.0 # CFG is baked into the distilled weights
num_inference_steps = 1 # One-Forcing one-step generation
num_frame_per_block = 1 # frame-wise
sample_size = [480, 832]
video_length = 81
```
运行：
```bash
python examples/wan2.1_one_forcing/predict_t2v.py
```
`transformer_path` 既可以是官方 checkpoint 文件、纯 `.pt`/`.safetensors` 文件，也可以是训练输出的 **checkpoint 目录**（其 live/EMA 权重会被自动解析）。GAN 判别器头是 fake_score 的训练专用子模块，因此推理**无需任何判别器接线**。
### 6.2 单步与 4 步选择器
当 `stochastic_sampling=True` 时，`num_inference_steps` 决定 warped 时间步调度：
| `num_inference_steps` | 去噪时间步 | 适用场景 |
|-----------------------|------------|----------|
| `1`（One-Forcing 默认） | `[1000]` | 最快；逐帧单步生成 |
| `4` | `[1000, 750, 500, 250]` | 更稳定、更慢 |
### 6.3 首帧调度
共享的 `WanSelfForcingPipeline` 对**所有块使用同一份去噪调度**，因此此处不单独为首帧设置调度（例如首帧 `[1000, 750, 500, 250]`、后续帧 `[1000]`）；该 pipeline 被多个 example 共用，有意保持不变。
---
## 七、Checkpoint 往返（训练→推理）
FSDP 训练器将每个 `checkpoint-{step}/` 按如下布局写出，因此训练产物可直接喂给推理，无需改动 key：
| 组件 | 保存位置 | 是否含 `gan_head`？ |
|------|----------|---------------------|
| **生成器**（live） | `checkpoint-{step}/diffusion_pytorch_model.safetensors`（根目录，经 FSDP hook） | 否 |
| **生成器**（EMA） | `checkpoint-{step}/transformer_ema/` | 否 |
| **Fake-score + 判别器头** | `checkpoint-{step}/fake_score/` | 是（FP32 判别器头张量） |
- **生成器从不包含 `gan_head`** 键（判别器头挂在 fake-score 上），因此 `predict_t2v.py` 可用 `strict=False` 干净加载。
- 推理脚本中的 `_resolve_transformer_path` 依次解析：`transformer_ema/`、`ema_transformer/`（兼容别名）、`transformer/`，最后到根目录 safetensors。
- **恢复** GAN 训练时须使用相同的 GAN/判别器头配置：加载前会根据 CLI 参数重建判别器头，并从 `fake_score/` 还原。
---
## 八、更多资源
- **One-Forcing（源配方）**：https://github.com/Aurora-edu/One-Forcing
- **Causal-Forcing（DMD 阶段）**：https://github.com/thu-ml/Causal-Forcing
- **Stage 1（AR Diffusion）训练指南**：[README_TRAIN_AR_DIFFUSION_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md)
- **Stage 2（CCD）训练指南**：[README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md)
- **底层 DMD 训练器指南**：[README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md)
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
scripts/wan2.1_one_forcing/train_one_forcing.py
+3490
View File
File diff suppressed because it is too large Load Diff
scripts/wan2.1_one_forcing/train_one_forcing.sh
+85
View File
@@ -0,0 +1,85 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
# One-Forcing needs REAL videos to build the adversarial critic's clean_latents, so the dataset
# must be a video meta (type=video + file_path), NOT a prompt-only json. The official demo
# dataset X-Fun-Videos-Demo ships captioned videos and is normalized by --enable_bucket
# --random_hw_adapt --training_with_video_token_length.
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# Warm-start generator: a distilled/ODE checkpoint that the DMD + adversarial critic then polish.
# Default: the official One-Forcing ODE-initialized checkpoint, downloaded via
# hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
# You can also point it at your own Causal-Forcing Stage-2 CCD checkpoint.
export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
# 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
# One-Forcing = Causal-Forcing Stage-3 DMD + a shared noised-latent adversarial critic on fake_score.
# - 1-step frame-wise generation: --denoising_step_indices_list 1000, --num_frame_per_block=1;
# 21 latent frames from --video_sample_n_frames=81 @ 480x832.
# - GAN critic: gan_g_weight/gan_d_weight=0.03, feature layers [21, 29], ffn 2048, heads 12,
# non-relativistic, R1/R2 off. The head rides on fake_score and is training-only.
# - GAN requires real videos, --gradient_accumulation_steps=1, no DeepSpeed, and --video_sample_stride=1
# (stride=2 on 81-frame sources yields ~10 latent frames < score_num_frames=21 and aborts).
# - guidance: local CFG is cond + scale*(cond - uncond), so --real_guidance_scale=6.0.
# - optimizer/EMA/steps: lr=lr_critic=1e-5, betas (0.0, 0.999), weight_decay=0.01,
# --use_ema --ema_decay=0.99, gen_update_interval=5, max_steps=200.
accelerate launch --mixed_precision="bf16" --use_fsdp \
--fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
--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.1_one_forcing/train_one_forcing.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--ode_transformer_path=$STAGE2_CKPT \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--fix_sample_size 480 832 \
--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 \
--max_train_steps=200 \
--checkpointing_steps=100 \
--learning_rate=1.0e-05 \
--learning_rate_critic=1.0e-05 \
--lr_scheduler="constant" \
--seed=42 \
--output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.01 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--num_frame_per_block=1 \
--use_kv_cache_training \
--use_ema \
--ema_decay=0.99 \
--denoising_step_indices_list 1000 \
--real_guidance_scale=6.0 \
--score_timestep_clamp \
--fake_guidance_scale=0.0 \
--gen_update_interval=5 \
--gan_g_weight=0.03 \
--gan_d_weight=0.03 \
--gan_feature_layers 21 29 \
--gan_block_ffn_dim=2048 \
--gan_block_num_heads=12 \
--resume_from_checkpoint="latest" \
--trainable_modules "."
scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md
+62 -24
View File
@@ -179,6 +179,7 @@ The difference between DeepSpeed-Zero-2 and FSDP lies in whether to shard model
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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.
@@ -189,6 +190,7 @@ NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -204,21 +206,26 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \
@@ -247,17 +254,19 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| `--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 |
| `--learning_rate_critic` | Initial learning rate (critic) | 4.0e-07 |
| `--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_wan2.1_self_forcing_distill` |
| `--gradient_checkpointing` | Enable gradient checkpointing | - |
| `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
| `--adam_weight_decay` | AdamW weight decay | 3e-2 |
| `--adam_weight_decay` | AdamW weight decay | 0.0 |
| `--adam_beta1` | AdamW beta1 | 0.0 |
| `--adam_beta2` | AdamW beta2 | 0.999 |
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--vae_mini_batch` | VAE encoding mini-batch size | 1 |
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
| `--max_grad_norm` | Gradient clipping threshold | 10.0 |
| `--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 | - |
@@ -277,7 +286,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
|-----|------|-------|
| `--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 |
| `--real_guidance_scale` | Real guidance scale for scoring | 4.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 |
@@ -365,6 +374,7 @@ You can configure validation parameters to periodically generate test videos dur
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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.
@@ -375,6 +385,7 @@ NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock --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.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -390,21 +401,26 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \
@@ -427,6 +443,7 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/
Training shell command is as follows:
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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.
@@ -437,6 +454,7 @@ 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/wan2.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -452,21 +470,26 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \
@@ -481,6 +504,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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.
@@ -491,6 +515,7 @@ NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -506,21 +531,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \
@@ -540,6 +570,7 @@ Assuming 2 machines, each with 8 GPUs:
**Machine 0 (Master)**:
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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
@@ -555,6 +586,7 @@ 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/wan2.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -570,21 +602,26 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \
@@ -596,6 +633,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
**Machine 1 (Worker)**:
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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" # Same as Master
scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md
+62 -24
View File
@@ -180,6 +180,7 @@ hf download gdhe17/Self-Forcing --local-dir models/Diffusion_Transformer/Self-Fo
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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.
@@ -190,6 +191,7 @@ NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -205,21 +207,26 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \
@@ -248,17 +255,19 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| `--num_train_epochs` | 训练 epoch 数 | 100 |
| `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 |
| `--learning_rate` | 初始学习率（生成器） | 2e-06 |
| `--learning_rate_critic` | 初始学习率（判别器） | 2e-06 |
| `--learning_rate_critic` | 初始学习率（判别器） | 4.0e-07 |
| `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
| `--lr_warmup_steps` | 学习率预热步数 | 100 |
| `--seed` | 随机种子 | 42 |
| `--output_dir` | 输出目录 | `output_dir_wan2.1_self_forcing_distill` |
| `--gradient_checkpointing` | 激活重计算 | - |
| `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
| `--adam_weight_decay` | AdamW 权重衰减 | 0.0 |
| `--adam_beta1` | AdamW beta1 | 0.0 |
| `--adam_beta2` | AdamW beta2 | 0.999 |
| `--adam_epsilon` | AdamW epsilon 值 | 1e-10 |
| `--vae_mini_batch` | VAE 编码时的迷你批次大小 | 1 |
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
| `--max_grad_norm` | 梯度裁剪阈值 | 10.0 |
| `--fsdp_reduce_fp32` | 可选：FSDP 混精下梯度以 fp32 做跨卡 reduce（默认 bf16） | - |
| `--enable_bucket` | 启用分桶训练，不裁剪图片/视频，按分辨率分组训练 | - |
| `--random_hw_adapt` | 自动缩放图片/视频到 `[min_size, max_size]` 范围内的随机尺寸 | - |
@@ -278,7 +287,7 @@ 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 |
| `--real_guidance_scale` | 用于评分的真实 guidance scale | 4.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 |
@@ -366,6 +375,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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.
@@ -376,6 +386,7 @@ NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock --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.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -391,21 +402,26 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \
@@ -428,6 +444,7 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/
训练 shell 命令如下：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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.
@@ -438,6 +455,7 @@ 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/wan2.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -453,21 +471,26 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \
@@ -482,6 +505,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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.
@@ -492,6 +516,7 @@ NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -507,21 +532,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \
@@ -541,6 +571,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
**机器 0（Master）**：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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
@@ -556,6 +587,7 @@ 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/wan2.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -571,21 +603,26 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \
@@ -597,6 +634,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
**机器 1（Worker）**：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
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 相同
scripts/wan2.1_self_forcing/train_distill.sh
+11 -4
View File
@@ -1,4 +1,5 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
export DATASET_NAME="datasets/internal_datasets/"
export DATASET_META_NAME="datasets/internal_datasets/metadata.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
@@ -9,6 +10,7 @@ NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
@@ -24,21 +26,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_self_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--max_grad_norm=10.0 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--denoising_step_indices_list 1000 750 500 250 \
--real_guidance_scale=4.0 \
--score_timestep_clamp \
--use_kv_cache_training \
--num_frame_per_block=3 \
--train_mode="normal" \