Update Sparkwan, Reward Model Training, FastVideo (#525)
This commit is contained in:
Bubbliiiing
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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`
|
||||
@@ -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`
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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"
|
||||
@@ -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"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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"
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user