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

Update multigpus in vsa, sla, spark

Browse Source
This commit is contained in:
bubbliiiing
2026-11-01 10:00:44 +08:00
parent f00a7d480e
commit be4fa825cd
23 changed files with 1232 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/fastvideo/predict_t2v_minimax_h3.py
+11 -4
View File
@@ -40,9 +40,10 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
# GPU at full load, so keep an offloading mode here.
GPU_memory_mode = "model_group_offload"
# Multi GPUs config
# NOTE: the VSA sparse kernel has no sequence-parallel path, so a FastVideo-H3 student runs
# single-GPU only; keep ulysses_degree = ring_degree = 1 (enable_multi_gpus_inference() raises
# otherwise, rather than silently falling back to the dense multi-GPU processor).
# The FastVideo-H3 VSA student supports sequence-parallel inference through ulysses head-parallel
# only (the VSA kernel has no ring-attention path): put every GPU into ulysses_degree, keep
# ring_degree = 1, and ulysses_degree must divide the 56 attention heads (2 / 4 / 7 / 8 / ...).
# The defaults stay single-GPU because the ~19B stack still needs an offloading mode below.
ulysses_degree = 1
ring_degree = 1
# Use FSDP to save more GPU memory in multi gpus. Unused at ulysses_degree = ring_degree = 1.
@@ -247,8 +248,14 @@ if use_qfloat8:
wants_offload = GPU_memory_mode in ("model_cpu_offload", "model_group_offload", "sequential_cpu_offload")
if ulysses_degree > 1 or ring_degree > 1:
# The VSA sparse kernel only supports ulysses head-parallel (no ring-attention path); a VSA
# student under ring_degree>1 would silently fold ring into a larger ulysses degree and give
# none of ring's per-rank memory saving, so reject it instead of mis-attending.
if ring_degree > 1:
raise ValueError(
"FastVideo-H3 VSA multi-GPU requires ring_degree=1; put every GPU into ulysses_degree "
"instead (ulysses_degree must divide the 56 attention heads).")
from functools import partial
# NOTE: raises for the VSA student (no sequence-parallel path); keep ulysses_degree=ring_degree=1.
transformer.enable_multi_gpus_inference()
if fsdp_dit:
fp32_modules = [m for m in transformer.modules()
examples/fastvideo/predict_t2v_wan2.1.py
+6
View File
@@ -239,6 +239,12 @@ pipeline = WanPipeline(
scheduler=scheduler,
)
if ulysses_degree > 1 or ring_degree > 1:
# The VSA kernel has no ring-attention path; a sparse student under ring_degree>1 would
# silently mis-attend. student_topk >= 1.0 is the dense fallback, which ring does support.
if ring_degree > 1 and student_topk < 1.0:
raise ValueError(
"FastVideo VSA multi-GPU requires ring_degree=1; put every GPU into ulysses_degree "
"instead (num_heads must divide ulysses_degree).")
from functools import partial
transformer.enable_multi_gpus_inference()
if fsdp_dit:
examples/sparkdiffusion/predict_t2v_wan2.1.py
+1
View File
@@ -295,6 +295,7 @@ def save_results():
else:
video_path = os.path.join(save_path, prefix + ".mp4")
save_videos_grid(sample, video_path, fps=fps)
print(f"Saved video to: {video_path}")
if ulysses_degree * ring_degree > 1:
import torch.distributed as dist
examples/turbodiffusion/predict_i2v_wan2.2.py
+6 -1
View File
@@ -213,6 +213,12 @@ pipeline = Wan2_2I2VPipeline(
scheduler=scheduler,
)
if ulysses_degree > 1 or ring_degree > 1:
# The SLA sparse kernel only supports ulysses head-parallel (no ring-attention path); a
# sparse expert under ring_degree>1 would silently mis-attend.
if ring_degree > 1:
raise ValueError(
"TurboDiffusion SLA multi-GPU requires ring_degree=1; put every GPU into "
"ulysses_degree instead (num_heads must divide ulysses_degree).")
from functools import partial
transformer.enable_multi_gpus_inference()
transformer_2.enable_multi_gpus_inference()
@@ -289,7 +295,6 @@ def save_results():
image = (image * 255).numpy().astype(np.uint8)
image = Image.fromarray(image)
image.save(video_path)
print(f"Saved image to: {video_path}")
else:
video_path = os.path.join(save_path, prefix + ".mp4")
save_videos_grid(sample, video_path, fps=fps)
examples/turbodiffusion/predict_t2v_wan2.1.py
+7 -2
View File
@@ -67,7 +67,7 @@ sampler_name = "Flow"
# Used when the sampler is in "Flow_Unipc", "Flow_DPM++".
# If you want to generate a 480p video, it is recommended to set the shift value to 3.0.
# If you want to generate a 720p video, it is recommended to set the shift value to 5.0.
shift = 3
shift = 5
# Load pretrained model if need
# Download from https://huggingface.co/TurboDiffusion/TurboWan2.1-T2V-1.3B-480P/
@@ -172,6 +172,12 @@ pipeline = WanPipeline(
scheduler=scheduler,
)
if ulysses_degree > 1 or ring_degree > 1:
# The SLA sparse kernel only supports ulysses head-parallel (no ring-attention path); a
# sparse student under ring_degree>1 would silently mis-attend.
if ring_degree > 1:
raise ValueError(
"TurboDiffusion SLA multi-GPU requires ring_degree=1; put every GPU into "
"ulysses_degree instead (num_heads must divide ulysses_degree).")
from functools import partial
transformer.enable_multi_gpus_inference()
if fsdp_dit:
@@ -234,7 +240,6 @@ def save_results():
image = (image * 255).numpy().astype(np.uint8)
image = Image.fromarray(image)
image.save(video_path)
print(f"Saved image to: {video_path}")
else:
video_path = os.path.join(save_path, prefix + ".mp4")
save_videos_grid(sample, video_path, fps=fps)
scripts/taehv/README_TRAIN.md
+25 -15
View File
@@ -180,25 +180,32 @@ NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
--config_path="config/wan2.2/wan_civitai_5b.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--tae_path=$TAE_PATH \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_size=512 \
--video_sample_stride=1 \
--video_sample_n_frames=33 \
--pretrained_model_name_or_path="$MODEL_NAME" \
--tae_path="$TAE_PATH" \
--train_data_dir="$DATASET_NAME" \
--train_data_meta="$DATASET_META_NAME" \
--fix_sample_size 720 1280 \
--video_sample_strides 1 2 3 \
--video_sample_n_frames=25 \
--validation_n_frames=49 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=4 \
--num_train_epochs=100 \
--checkpointing_steps=500 \
--dataloader_num_workers=8 \
--max_train_steps=10000 \
--checkpointing_steps=1000 \
--validation_steps=1000 \
--validation_paths "asset/inpaint_video.mp4" \
--learning_rate=1e-04 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--lr_warmup_steps=200 \
--latent_loss_weight=1.0 \
--use_ema \
--ema_decay=0.999 \
--seed=42 \
--output_dir="output_dir_taehv_w2.2" \
--mixed_precision="bf16" \
--adam_beta1=0.9 \
--adam_beta2=0.99 \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
@@ -208,6 +215,8 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
--resume_from_checkpoint=latest
```
These settings are also provided in `scripts/taehv/train_taehv.sh`. Edit its training arguments and run `bash scripts/taehv/train_taehv.sh`.
**Wan2.1 / Wan2.2 14B (taew2_1) Example**:
Same as above, with these changes:
@@ -246,6 +255,7 @@ export TAE_PATH="taew2_1.safetensors"
| `--gradient_accumulation_steps` | Gradient accumulation steps | 1 |
| `--dataloader_num_workers` | Number of DataLoader workers | 4 |
| `--num_train_epochs` | Number of training epochs | 100 |
| `--max_train_steps` | Total training steps; overrides `--num_train_epochs` when set | 10000 |
| `--checkpointing_steps` | Save a checkpoint every N steps | 500 |
| `--checkpoints_total_limit` | Max number of checkpoints to store | `None` |
| `--learning_rate` | Initial learning rate | 1e-4 |
@@ -265,7 +275,7 @@ export TAE_PATH="taew2_1.safetensors"
| `--trainable_modules_low_learning_rate` | Trainable modules with lr/2 | `[]` |
| `--resume_from_checkpoint` | Resume training from checkpoint, use `"latest"` to auto-select | `latest` |
| `--validation_steps` / `--validation_epochs` | Run validation every N steps / epochs | 2000 / 5 |
| `--validation_paths` | Video paths for validation reconstruction comparison | `"asset/1.mp4"` |
| `--validation_paths` | Video paths for validation reconstruction comparison | `"asset/inpaint_video.mp4"` |
| `--validation_n_frames` | Validation clip length (4k+1); defaults to `--video_sample_n_frames` | `None` |
**Sample Size Configuration Guide**:
@@ -282,11 +292,11 @@ You can configure validation parameters to periodically reconstruct test videos
|-----------|-------------|-------------------|
| `--validation_steps` | Run validation every N steps | 2000 |
| `--validation_epochs` | Run validation every N epochs | 5 |
| `--validation_paths` | Validation video paths | `"asset/1.mp4"` |
| `--validation_paths` | Validation video paths | `"asset/inpaint_video.mp4"` |
| `--validation_n_frames` | Validation clip length, **4k+1**; omit to reuse `--video_sample_n_frames` | 81 |
```bash
--validation_paths "asset/1.mp4" \
--validation_paths "asset/inpaint_video.mp4" \
--validation_steps=2000 \
--validation_epochs=5 \
--validation_n_frames=81
@@ -304,7 +314,7 @@ You can configure validation parameters to periodically reconstruct test videos
- **Decoder-only distillation**: Add `--freeze_tae_encoder` to only improve decoding quality.
- **Memory**: The teacher full VAE is the main memory consumer. Use `--low_vram` to keep it on CPU between encoding steps, or reduce `--video_sample_n_frames` / `--video_sample_size`.
- **Sequential TAE**: `--use_taehv_sequential` trades speed for O(1) activation memory w.r.t. video length.
- **EMA**: `--use_ema` is recommended for the final deliverable; the final saved `taehv` directory contains the EMA weights. Pass `--ema_decay` (e.g. `0.999`, as `train_taehv_super.sh` does) when short runs need the average to catch up faster than the 0.9999 default.
- **EMA**: `--use_ema` is recommended for the final deliverable; the final saved `taehv` directory contains the EMA weights. Pass `--ema_decay` (e.g. `0.999`, as both `train_taehv.sh` and `train_taehv_super.sh` do) when short runs need the average to catch up faster than the 0.9999 default.
### 3.6 Multi-Node Distributed Training
scripts/taehv/README_TRAIN_zh-CN.md
+25 -15
View File
@@ -180,25 +180,32 @@ NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
--config_path="config/wan2.2/wan_civitai_5b.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--tae_path=$TAE_PATH \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--video_sample_size=512 \
--video_sample_stride=1 \
--video_sample_n_frames=33 \
--pretrained_model_name_or_path="$MODEL_NAME" \
--tae_path="$TAE_PATH" \
--train_data_dir="$DATASET_NAME" \
--train_data_meta="$DATASET_META_NAME" \
--fix_sample_size 720 1280 \
--video_sample_strides 1 2 3 \
--video_sample_n_frames=25 \
--validation_n_frames=49 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=4 \
--num_train_epochs=100 \
--checkpointing_steps=500 \
--dataloader_num_workers=8 \
--max_train_steps=10000 \
--checkpointing_steps=1000 \
--validation_steps=1000 \
--validation_paths "asset/inpaint_video.mp4" \
--learning_rate=1e-04 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--lr_warmup_steps=200 \
--latent_loss_weight=1.0 \
--use_ema \
--ema_decay=0.999 \
--seed=42 \
--output_dir="output_dir_taehv_w2.2" \
--mixed_precision="bf16" \
--adam_beta1=0.9 \
--adam_beta2=0.99 \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
@@ -208,6 +215,8 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
--resume_from_checkpoint=latest
```
相同参数已写入 `scripts/taehv/train_taehv.sh`，也可修改其中的训练参数后，运行 `bash scripts/taehv/train_taehv.sh`。
**Wan2.1 / Wan2.2 14B（taew2_1）训练示例**：
与上面相同，仅做如下修改：
@@ -246,6 +255,7 @@ export TAE_PATH="taew2_1.safetensors"
| `--gradient_accumulation_steps` | 梯度累积步数 | 1 |
| `--dataloader_num_workers` | DataLoader 子进程数 | 4 |
| `--num_train_epochs` | 训练 epoch 数 | 100 |
| `--max_train_steps` | 总训练步数；设置后覆盖 `--num_train_epochs` | 10000 |
| `--checkpointing_steps` | 每 N 步保存 checkpoint | 500 |
| `--checkpoints_total_limit` | 最多保留的 checkpoint 数量 | `None` |
| `--learning_rate` | 初始学习率 | 1e-4 |
@@ -265,7 +275,7 @@ export TAE_PATH="taew2_1.safetensors"
| `--trainable_modules_low_learning_rate` | 以 lr/2 训练的模块 | `[]` |
| `--resume_from_checkpoint` | 恢复训练路径，使用 `"latest"` 自动选择最新 checkpoint | `latest` |
| `--validation_steps` / `--validation_epochs` | 每 N 步 / 每 N 个 epoch 执行一次验证 | 2000 / 5 |
| `--validation_paths` | 用于重建对比验证的视频路径 | `"asset/1.mp4"` |
| `--validation_paths` | 用于重建对比验证的视频路径 | `"asset/inpaint_video.mp4"` |
| `--validation_n_frames` | 验证片段帧数（4k+1），不传则复用 `--video_sample_n_frames` | `None` |
**Sample Size 配置指南**：
@@ -282,11 +292,11 @@ export TAE_PATH="taew2_1.safetensors"
|------|------|--------|
| `--validation_steps` | 每 N 步执行一次验证 | 2000 |
| `--validation_epochs` | 每 N 个 epoch 执行一次验证 | 5 |
| `--validation_paths` | 验证视频路径 | `"asset/1.mp4"` |
| `--validation_paths` | 验证视频路径 | `"asset/inpaint_video.mp4"` |
| `--validation_n_frames` | 验证片段帧数，**必须为 4k+1**；不传则复用 `--video_sample_n_frames` | 81 |
```bash
--validation_paths "asset/1.mp4" \
--validation_paths "asset/inpaint_video.mp4" \
--validation_steps=2000 \
--validation_epochs=5 \
--validation_n_frames=81
@@ -304,7 +314,7 @@ export TAE_PATH="taew2_1.safetensors"
- **Decoder-only 蒸馏**：加上 `--freeze_tae_encoder` 只提升解码质量。
- **显存**：teacher 完整 VAE 是主要的显存消耗者。使用 `--low_vram` 可以将其放在 CPU、仅在编码时搬上 GPU；也可以降低 `--video_sample_n_frames` / `--video_sample_size`。
- **串行 TAE**：`--use_taehv_sequential` 以速度换显存，激活显存与视频长度无关（O(1)）。
- **EMA**：推荐开启 `--use_ema` 作为最终交付物；最终保存的 `taehv` 目录里是 EMA 权重。短训练希望均值更快跟上时，可传 `--ema_decay`（如 `train_taehv_super.sh` 用的 `0.999`）替代默认的 0.9999。
- **EMA**：推荐开启 `--use_ema` 作为最终交付物；最终保存的 `taehv` 目录里是 EMA 权重。短训练希望均值更快跟上时，可传 `--ema_decay`（如 `train_taehv.sh` 与 `train_taehv_super.sh` 用的 `0.999`）替代默认的 0.9999。
### 3.6 多机分布式训练
scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL.md
+77 -18
View File
@@ -34,6 +34,7 @@ This document provides the complete workflow for **Flex-Forcing Stage 2 — Dist
- [3.5 Training with FSDP](#35-training-with-fsdp)
- [3.6 Other Backends](#36-other-backends)
- [3.7 Multi-Node Distributed Training](#37-multi-node-distributed-training)
- [3.8 Optional One-Forcing GAN, EMA, and Precision Controls](#38-optional-one-forcing-gan-ema-and-precision-controls)
- [4. Inference Testing](#4-inference-testing)
- [4.1 Inference Parameters](#41-inference-parameters)
- [4.2 Text-to-Video (T2V) Inference](#42-text-to-video-t2v-inference)
@@ -107,6 +108,8 @@ DMD uses a **TextDataset** (`--train_mode="normal"`) — the generator creates i
`train_distill.sh` defaults to `datasets/internal_datasets/metadata.json`. For a real run, swap in the **VidProM extended** prompts the paper uses.
Enabling any GAN/R1/R2 weight (see **3.8**) switches the loader to **ImageVideoDataset**: captioned real videos are then required, cropped to the critic's fake input frame count.
---
## 3. Distillation Training
@@ -174,15 +177,16 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -204,6 +208,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--use_ema \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
@@ -220,6 +226,7 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
|-----------|-------------|---------------|
| `--config_path` | Component subpaths / `transformer_additional_kwargs` | `config/wan2.1/wan_civitai.yaml` |
| `--pretrained_model_name_or_path` | Base model directory (1.3B) | `$MODEL_NAME` |
| `--real_score_pretrained_model_name_or_path` | 14B real-score teacher directory; defaults to `--pretrained_model_name_or_path` | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
| `--train_data_meta` | Prompt JSON (only the `"text"` field is read) | `$DATASET_META_NAME` |
| `--ode_transformer_path` | Stage 1 ODE checkpoint for generator/critic init | `$STAGE1_CKPT` |
| `--fix_sample_size` | Fixed `H W` of the training clip | `480 832` |
@@ -230,9 +237,9 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
| `--dataloader_num_workers` | DataLoader workers | 8 |
| `--num_train_epochs` | Number of training epochs (overridden by `--max_train_steps`) | 100 |
| `--max_train_steps` | Hard step cap; the paper uses 600 | 600 |
| `--checkpointing_steps` | Save checkpoint every N steps | 50 |
| `--learning_rate` | Generator learning rate | `2e-06` |
| `--learning_rate_critic` | Critic learning rate | `2e-06` |
| `--checkpointing_steps` | Save checkpoint every N steps | 200 |
| `--learning_rate` | Generator learning rate | `2.0e-06` |
| `--learning_rate_critic` | Critic learning rate | `4.0e-07` |
| `--lr_scheduler` | LR scheduler type | `constant_with_warmup` |
| `--lr_warmup_steps` | LR warmup steps | 100 |
| `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 4.5 |
@@ -322,15 +329,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -352,6 +360,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--use_ema \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
@@ -400,15 +410,16 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -430,10 +441,13 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
`--use_ema` is dropped here: `train_distill.py` raises `NotImplementedError` for EMA under DeepSpeed Zero-3.
#### 3.6.2 Training without DeepSpeed and FSDP
**This approach is not recommended because there is no memory-saving backend, which can easily cause out-of-memory errors**. We only provide the training shell for reference.
@@ -464,15 +478,16 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -494,6 +509,8 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--use_ema \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
@@ -538,15 +555,16 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -568,6 +586,8 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--use_ema \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
@@ -602,6 +622,44 @@ NCCL_DEBUG=INFO
- **Data Synchronization**: All machines must be able to access the same data paths (NFS/shared storage)
### 3.8 Optional One-Forcing GAN, EMA, and Precision Controls
These options port the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms already used by the causal-forcing DMD trainer. All GAN/R1/R2 weights default to **0** and stay off in the launcher. `--use_ema`, `--fsdp_reduce_fp32` and `--score_timestep_clamp` default to **off** in argparse, but `train_distill.sh` enables all three; the DeepSpeed Zero-3 command in **3.6.1** drops `--use_ema`, which Zero-3 rejects.
The discriminator shares the **fake-score backbone**, with one learnable register token per selected feature layer. Its head is optimized by the critic optimizer, not the generator optimizer. Generator loss is `DMD + weighted GAN`; critic loss is `denoising + weighted GAN + R1 + R2`.
| Parameter | Description | Default |
|-----------|-------------|---------|
| `--use_ema` | EMA over the generator; checkpoints save a `transformer_ema` subfolder and validation samples from the EMA weights | off (enabled in `train_distill.sh`) |
| `--gan_g_weight` | Non-saturating generator GAN loss weight | `0.0` |
| `--gan_d_weight` | Discriminator GAN loss weight | `0.0` |
| `--gan_feature_layers` | Fake-score block indices used by the head | `21 29` |
| `--gan_block_ffn_dim` | Head cross-attention FFN dimension | `2048` |
| `--gan_block_num_heads` | Head count; must divide the fake-score hidden dimension | `12` |
| `--relativistic_discriminator` | Use paired real/fake logit differences | off |
| `--r1_weight` | APT-style finite-difference penalty on noisy real latents | `0.0` |
| `--r1_sigma` | R1 perturbation standard deviation | `0.01` |
| `--r2_weight` | Normalized finite-difference penalty on noisy generated latents | `0.0` |
| `--r2_sigma` | R2 perturbation standard deviation | `0.01` |
| `--fsdp_reduce_fp32` | FSDP FP32 gradient reduction; parameter gathering stays low precision | off (enabled in `train_distill.sh`) |
| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off (enabled in `train_distill.sh`) |
For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` for the generator and `softplus(-D_real) + softplus(D_fake)` for the discriminator. Relativistic losses instead use `softplus(-(D_fake - D_real))` and `softplus(-(D_real - D_fake))`.
R1 follows the source's APT form: `r1_weight * mean((D_real_perturbed - D_real)^2)`, **without dividing by sigma**. R2 follows `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`. These are finite-difference penalties, not input-gradient double-backward regularizers. Real/fake samples share timestep and noise; perturbations use independent noise. The non-relativistic critic runs fake and real as two sequential 1x-batch forwards with a backward after each (the memory-optimized split), the relativistic critic runs one packed 2x-batch forward, and R1/R2 each re-forward their own 2x batch after the main backward(s).
To enable GAN, pass positive generator/discriminator weights and, optionally, R1/R2 weights to `train_distill.py`. Requirements:
- Captioned real videos, not a prompt-only JSON. The real latents are cropped to the critic's fake input frame count, so the clip must cover at least that many latent frames.
- `--gradient_accumulation_steps=1`; DeepSpeed is not supported by the GAN branch.
- Valid feature-layer indices for the selected fake-score model. The default `21 29` layers and `12` heads fit Wan2.1-T2V-1.3B; other backbones may need different settings.
- Enabling only `gan_g_weight` does not supply a discriminator classification-learning objective. Choose the discriminator weight deliberately; nonzero values are not a guaranteed quality improvement.
Trainable fake-score and head master weights stay **FP32**, while forward uses the configured autocast/FSDP precision. Logits, softplus losses, and finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient communication precision, not the forward dtype or model architecture. The clamp changes score-training timesteps, not the generator rollout or validation schedule.
Resume GAN checkpoints with the same GAN/head configuration. The head is saved with fake-score, including FP32 head tensors, and is recreated before loading. It is not part of the exported generator, so inference needs no discriminator or new GAN flags.
Under `--use_ema`, the shadow stays in FP32 and mirrors the live model's mixed-precision policy; FSDP EMA requires FSDP1 and DeepSpeed Zero-3 is rejected.
---
## 4. Inference Testing
@@ -752,6 +810,7 @@ Any-order editing needs the full KV cache: both entries set `local_attn_size = -
## 5. Additional Resources
- **Flex-Forcing paper**: https://arxiv.org/abs/2607.03509
- **One-Forcing (optional GAN mechanism)**: https://github.com/Aurora-edu/One-Forcing
- **Stage 1**: [README_TRAIN_ODE.md](./README_TRAIN_ODE.md)
- **Shared chunking helpers**: `videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`, `sample_flexible_chunks`, `sample_flex_level0` (the level-0 four-arm mixture, drawn by **both** trainers), `refine_partition`, `build_pyramid_partitions`, `validate_nested_partitions`, `broadcast_chunk_sizes`
- **Model / pipeline**: `videox_fun/models/wan_transformer3d_flex_forcing.py`, `videox_fun/pipeline/pipeline_wan_flex_forcing.py`
scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md
+77 -18
View File
@@ -34,6 +34,7 @@
- [3.5 使用 FSDP 训练](#35-使用-fsdp-训练)
- [3.6 其他后端](#36-其他后端)
- [3.7 多机分布式训练](#37-多机分布式训练)
- [3.8 可选的 One-Forcing GAN、EMA 与精度开关](#38-可选的-one-forcing-ganema-与精度开关)
- [四、推理测试](#四推理测试)
- [4.1 推理参数解析](#41-推理参数解析)
- [4.2 文生视频（T2V）推理](#42-文生视频t2v推理)
@@ -107,6 +108,8 @@ DMD 使用 **TextDataset**（`--train_mode="normal"`）——生成器通过 rol
`train_distill.sh` 默认指向 `datasets/internal_datasets/metadata.json`。真实训练请换成论文使用的 **VidProM extended** prompt 集。
开启任何 GAN/R1/R2 权重（见 **3.8**）会把数据加载切到 **ImageVideoDataset**：此时需要带标注的真实视频，并会裁剪到 critic fake 输入的帧数。
---
## 三、蒸馏训练
@@ -174,15 +177,16 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -204,6 +208,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--use_ema \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
@@ -220,6 +226,7 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
|-----------|-------------|---------------|
| `--config_path` | 组件子路径 / `transformer_additional_kwargs` | `config/wan2.1/wan_civitai.yaml` |
| `--pretrained_model_name_or_path` | 底座模型目录（1.3B） | `$MODEL_NAME` |
| `--real_score_pretrained_model_name_or_path` | 14B 真实分数教师目录；默认回落 `--pretrained_model_name_or_path` | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
| `--train_data_meta` | prompt JSON（只读 `"text"` 字段） | `$DATASET_META_NAME` |
| `--ode_transformer_path` | 用于生成器 / critic 初始化的阶段 1 ODE checkpoint | `$STAGE1_CKPT` |
| `--fix_sample_size` | 固定训练片段的 `H W` | `480 832` |
@@ -230,9 +237,9 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
| `--dataloader_num_workers` | DataLoader 工作进程数 | 8 |
| `--num_train_epochs` | 训练 epoch 数（被 `--max_train_steps` 覆盖） | 100 |
| `--max_train_steps` | 硬性步数上限；论文用 600 | 600 |
| `--checkpointing_steps` | 每 N 步保存一次 checkpoint | 50 |
| `--learning_rate` | 生成器学习率 | `2e-06` |
| `--learning_rate_critic` | critic 学习率 | `2e-06` |
| `--checkpointing_steps` | 每 N 步保存一次 checkpoint | 200 |
| `--learning_rate` | 生成器学习率 | `2.0e-06` |
| `--learning_rate_critic` | critic 学习率 | `4.0e-07` |
| `--lr_scheduler` | 学习率调度器类型 | `constant_with_warmup` |
| `--lr_warmup_steps` | 学习率预热步数 | 100 |
| `--real_guidance_scale` | 真实分数（14B 教师）的 CFG scale | 4.5 |
@@ -333,15 +340,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -363,6 +371,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--use_ema \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
@@ -405,15 +415,16 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -435,10 +446,13 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
此处去掉 `--use_ema`：DeepSpeed Zero-3 下 `train_distill.py` 会对 EMA 抛出 `NotImplementedError`。
DeepSpeed-Zero-3 训练完成之后，得到 `checkpoint-xxx` 状态的保存权重，我们需要将权重文件合并为单个，可以使用提供的`scripts/zero_to_bf16.py`脚本。
```shell
@@ -473,15 +487,16 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -503,6 +518,8 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--use_ema \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
@@ -548,15 +565,16 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -578,6 +596,8 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--use_ema \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
```
@@ -604,6 +624,44 @@ export RANK=1
> - Use `NCCL_DEBUG=INFO` to troubleshoot issues in multi-machine training.
> - If there is no RDMA environment, set `export NCCL_IB_DISABLE=1` and `export NCCL_P2P_DISABLE=1` to avoid communication issues.
### 3.8 可选的 One-Forcing GAN、EMA 与精度开关
这些开关把 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 的机制（因果 forcing 的 DMD trainer 已在用）移植过来。所有 GAN/R1/R2 权重默认是 **0**，launcher 中也保持关闭；`--use_ema`、`--fsdp_reduce_fp32` 与 `--score_timestep_clamp` 在 argparse 中默认**关闭**，但 `train_distill.sh` 把三者都打开了；**3.6.1** 的 DeepSpeed Zero-3 命令则去掉了 Zero-3 不支持的 `--use_ema`。
判别器与 **fake-score 骨干共享**，每个被选中的特征层有一个可学习的 register token。判别器头由 critic 优化器优化，而不是生成器优化器。生成器损失为 `DMD + 加权 GAN`；critic 损失为 `去噪 + 加权 GAN + R1 + R2`。
| 参数 | 说明 | 默认 |
|-----------|-------------|---------|
| `--use_ema` | 对生成器做 EMA；checkpoint 额外保存 `transformer_ema` 子目录，验证样片采样自 EMA 权重 | 关（launcher 已启用） |
| `--gan_g_weight` | 非饱和生成器 GAN 损失权重 | `0.0` |
| `--gan_d_weight` | 判别器 GAN 损失权重 | `0.0` |
| `--gan_feature_layers` | 判别器头使用的 fake-score block 索引 | `21 29` |
| `--gan_block_ffn_dim` | 判别器头 cross-attention FFN 维度 | `2048` |
| `--gan_block_num_heads` | 判别器头 head 数；必须整除 fake-score 隐维度 | `12` |
| `--relativistic_discriminator` | 使用成对的 real/fake logit 差值 | 关 |
| `--r1_weight` | 对含噪真实 latent 的 APT 式有限差分惩罚 | `0.0` |
| `--r1_sigma` | R1 扰动标准差 | `0.01` |
| `--r2_weight` | 对含噪生成 latent 的归一化有限差分惩罚 | `0.0` |
| `--r2_sigma` | R2 扰动标准差 | `0.01` |
| `--fsdp_reduce_fp32` | FSDP 梯度归约用 FP32；参数收集保持低精度 | 关（launcher 已启用） |
| `--score_timestep_clamp` | 把 DMD、生成器 GAN 与 critic 的时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关（launcher 已启用） |
对 logit `D_fake` 与 `D_real`，普通 GAN 损失是生成器 `softplus(-D_fake)`、判别器 `softplus(-D_real) + softplus(D_fake)`；相对论版本改用 `softplus(-(D_fake - D_real))` 与 `softplus(-(D_real - D_fake))`。
R1 沿用来源的 APT 形式：`r1_weight * mean((D_real_perturbed - D_real)^2)`，**不除以 sigma**。R2 是 `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`。两者都是有限差分惩罚，不是输入梯度二重反向的正则。real/fake 样本共享时间步与噪声；扰动使用独立噪声。非相对论 critic 把 fake 与 real 拆成两次 1x batch 前向、各自一次 backward（省显存的拆分路径），相对论 critic 跑一次打包的 2x batch 前向，R1/R2 则在主 backward 之后各自再前向自己的 2x batch。
启用 GAN 需要给 `train_distill.py` 传入正的生成器/判别器权重，可选再加 R1/R2 权重。要求：
- 带标注的真实视频，不能是只有 prompt 的 JSON。真实 latent 会裁剪到 critic fake 输入的帧数，因此片段至少要覆盖那么多 latent 帧。
- `--gradient_accumulation_steps=1`；GAN 分支不支持 DeepSpeed。
- 特征层索引必须对所选 fake-score 模型有效。默认的 `21 29` 层与 `12` head 适配 Wan2.1-T2V-1.3B；换骨干可能需要不同配置。
- 只开 `gan_g_weight` 不会给判别器提供分类学习目标。请有意识地选择判别器权重；非零值不保证质量提升。
可训练的 fake-score 与判别器头 master 权重保持 **FP32**，前向使用配置的 autocast/FSDP 精度。logit、softplus 损失与有限差分归约都在 FP32 下计算。`--fsdp_reduce_fp32` 只改梯度通信精度，不改前向 dtype 或模型结构。时间步 clamp 只改分数训练的时间步，不改生成器 rollout 或验证调度。
恢复 GAN checkpoint 时请使用相同的 GAN/判别器头配置。判别器头随 fake-score 一起保存（含 FP32 的 head 张量），加载前重建。它不属于导出的生成器，因此推理不需要判别器或新的 GAN 旗标。
开启 `--use_ema` 时，shadow 保持 FP32 并与 live 模型的混合精度策略对齐；FSDP EMA 仅支持 FSDP1，DeepSpeed Zero-3 会被拒绝。
---
## 四、推理测试
@@ -741,4 +799,5 @@ num_frame_per_block = 7 # 干净上下文的提交粒度；None = 整
## 五、更多资源
- **Flex-Forcing 论文**：https://arxiv.org/abs/2607.03509
- **One-Forcing（可选 GAN 机制）**：https://github.com/Aurora-edu/One-Forcing
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
scripts/wan2.1_flex_forcing/train_distill.py
+403 -32
View File
@@ -17,6 +17,7 @@
import argparse
import contextlib
import dataclasses
import gc
import json
import logging
@@ -76,7 +77,8 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512,
from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
WanTransformer3DModel,
WanTransformer3DModel_FlexForcing,
WanTransformer3DModel_SelfForcing)
WanTransformer3DModel_SelfForcing,
add_wan_gan_discriminator)
from videox_fun.pipeline import (WanFlexForcingPipeline, WanI2VPipeline,
WanPipeline, WanSelfForcingPipeline)
from videox_fun.utils.discrete_sampler import DiscreteSampling
@@ -84,6 +86,7 @@ from videox_fun.utils.flex_chunking import (build_pyramid_partitions,
chunk_boundaries,
sample_flex_level0,
validate_nested_partitions)
from videox_fun.utils.fsdp_ema import FSDPEMA
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
@@ -171,6 +174,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi
return pred_for_score, score_num_frames, True
def repeat_batch(value, repeats):
"""Repeat a batched tensor or per-sample list along the batch dimension."""
if value is None or repeats == 1:
return value
if torch.is_tensor(value):
return torch.cat([value] * repeats)
return list(value) * repeats
def filter_kwargs(cls, kwargs):
import inspect
sig = inspect.signature(cls.__init__)
@@ -743,6 +755,15 @@ def parse_args():
" https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
),
)
parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true",
help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
)
parser.add_argument(
"--score_timestep_clamp", action="store_true",
help="Clamp DMD, GAN and critic timesteps to [0.02, 0.98] * num_train_timesteps.",
)
parser.add_argument(
"--non_ema_revision",
type=str,
@@ -965,6 +986,16 @@ def parse_args():
default=None,
help=("If you want to load the ode-trained weight into generator transformer3d, input its path."),
)
parser.add_argument(
"--real_score_pretrained_model_name_or_path",
type=str,
default=None,
help=(
"Path to a non-causal pretrained model used as the DMD real_score teacher. "
"For CF Stage 3 this should point at Wan2.1-T2V-14B. If unset, falls back to "
"`--pretrained_model_name_or_path` (1.3B) — useful only for sanity testing."
),
)
parser.add_argument(
"--vae_path",
type=str,
@@ -1170,6 +1201,47 @@ def parse_args():
help="Probability of applying teacher forcing per step (1.0 = always)"
)
parser.add_argument(
"--gan_g_weight", type=float, default=0.0,
help="One-Forcing: non-saturating generator GAN loss weight (0 disables). Positive GAN/R1/R2 weights require real videos.",
)
parser.add_argument(
"--gan_d_weight", type=float, default=0.0,
help="One-Forcing: discriminator GAN loss weight added to critic denoising loss (0 disables).",
)
parser.add_argument(
"--gan_feature_layers", nargs="+", type=int, default=[21, 29],
help="fake_score block indices feeding the discriminator head, one register token per layer.",
)
parser.add_argument(
"--gan_block_ffn_dim", type=int, default=2048,
help="FFN dimension of the discriminator cross-attention blocks.",
)
parser.add_argument(
"--gan_block_num_heads", type=int, default=12,
help="Number of heads of the discriminator cross-attention blocks.",
)
parser.add_argument(
"--relativistic_discriminator", action="store_true",
help="Use relativistic GAN losses based on real/fake logit differences.",
)
parser.add_argument(
"--r1_weight", type=float, default=0.0,
help="Finite-difference R1 penalty weight on noisy real latents (APT form, not divided by sigma).",
)
parser.add_argument(
"--r1_sigma", type=float, default=0.01,
help="Perturbation standard deviation for the R1 penalty.",
)
parser.add_argument(
"--r2_weight", type=float, default=0.0,
help="Finite-difference R2 penalty weight on noisy generated latents (0.5 * squared logit difference / sigma squared)."
)
parser.add_argument(
"--r2_sigma", type=float, default=0.01,
help="Perturbation standard deviation for the R2 penalty.",
)
args = parser.parse_args()
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
if env_local_rank != -1 and env_local_rank != args.local_rank:
@@ -1184,6 +1256,14 @@ def parse_args():
def main():
args = parse_args()
use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight))
if any(not math.isfinite(w) or w < 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight)):
raise ValueError("GAN and regularization weights must be finite and non-negative.")
if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1):
raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.")
if use_gan and (args.r1_weight > 0 and (not math.isfinite(args.r1_sigma) or args.r1_sigma <= 0)
or args.r2_weight > 0 and (not math.isfinite(args.r2_sigma) or args.r2_sigma <= 0)):
raise ValueError("Enabled R1/R2 penalties require finite, positive perturbation standard deviations.")
if args.report_to == "wandb" and args.hub_token is not None:
raise ValueError(
@@ -1247,6 +1327,9 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
@@ -1374,8 +1457,11 @@ def main():
transformer_additional_kwargs=generator_transformer_kwargs,
low_cpu_mem_usage=True,
).to(weight_dtype)
# CF Stage 3 uses a 14B non-causal teacher as real_score; pass an explicit
# `--real_score_pretrained_model_name_or_path` to point at Wan2.1-T2V-14B.
real_score_base = args.real_score_pretrained_model_name_or_path or args.pretrained_model_name_or_path
real_score_transformer3d = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
os.path.join(real_score_base, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
).to(weight_dtype)
@@ -1476,6 +1562,41 @@ def main():
param.requires_grad = True
break
if use_gan:
if (not args.gan_feature_layers or any(i < 0 or i >= len(fake_score_transformer3d.blocks) for i in args.gan_feature_layers)
or args.gan_block_ffn_dim <= 0 or args.gan_block_num_heads <= 0
or fake_score_transformer3d.dim % args.gan_block_num_heads != 0):
raise ValueError("GAN feature layers must exist and the positive head count must divide the fake_score hidden dimension.")
# Keep trainable discriminator master weights in fp32; forward uses autocast.
add_wan_gan_discriminator(
fake_score_transformer3d,
args.gan_feature_layers,
args.gan_block_ffn_dim,
args.gan_block_num_heads,
torch.float32,
seed=args.seed or 0,
)
# Create EMA for the generator_transformer3d.
if args.use_ema:
if zero_stage == 3:
raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.")
if fsdp_plugin is not None:
if getattr(fsdp_plugin, "fsdp_version", 1) != 1:
raise NotImplementedError("FSDPEMA 仅支持 FSDP1。")
ema_module = generator_transformer_cls.from_config(generator_transformer3d.config)
ema_generator_transformer3d = FSDPEMA(
ema_module, source=generator_transformer3d,
accelerator=accelerator, fsdp_plugin=fsdp_plugin,
)
else:
ema_generator_transformer3d = EMAModel(
generator_transformer3d.parameters(),
model_cls=generator_transformer_cls,
model_config=generator_transformer3d.config,
)
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
@@ -1486,7 +1607,10 @@ def main():
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
accelerate_state_dict = {k: v.to(dtype=weight_dtype) for k, v in accelerate_state_dict.items()}
accelerate_state_dict = {
k: v.to(dtype=torch.float32 if k.startswith("gan_head.") else weight_dtype)
for k, v in accelerate_state_dict.items()
}
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
@@ -1503,7 +1627,14 @@ def main():
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
models[0].save_pretrained(os.path.join(output_dir, "transformer"))
model = accelerator.unwrap_model(models[0])
save_dir = os.path.join(output_dir, "transformer")
if hasattr(model, "gan_head"):
from safetensors.torch import save_file
model.save_config(save_dir)
save_file(model.state_dict(), os.path.join(save_dir, "diffusion_pytorch_model.safetensors"), metadata={"format": "pt"})
else:
model.save_pretrained(save_dir)
if not args.use_deepspeed:
weights.pop()
@@ -1529,14 +1660,19 @@ def main():
else WanTransformer3DModel
)
# load diffusers style into model
load_model = model_cls.from_pretrained(
input_dir, subfolder="transformer"
)
unwrapped.register_to_config(**load_model.config)
# GAN heads are recreated from the same CLI configuration before resume.
if hasattr(unwrapped, "gan_head"):
from safetensors.torch import load_file
unwrapped.load_state_dict(load_file(os.path.join(input_dir, "transformer", "diffusion_pytorch_model.safetensors")), strict=True)
else:
# load diffusers style into model
load_model = model_cls.from_pretrained(
input_dir, subfolder="transformer"
)
unwrapped.register_to_config(**load_model.config)
unwrapped.load_state_dict(load_model.state_dict())
del load_model
unwrapped.load_state_dict(load_model.state_dict())
del load_model
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
@@ -1550,6 +1686,32 @@ def main():
accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)
if args.use_ema:
def save_ema_hook(models, weights, output_dir):
if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process:
ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
def load_ema_hook(models, input_dir):
ema_path = os.path.join(input_dir, "transformer_ema")
if isinstance(ema_generator_transformer3d, FSDPEMA):
ema_generator_transformer3d.load_pretrained(ema_path)
else:
_, ema_kwargs = generator_transformer_cls.load_config(ema_path, return_unused_kwargs=True)
load_model = generator_transformer_cls.from_pretrained(
input_dir, subfolder="transformer_ema",
transformer_additional_kwargs=generator_transformer_kwargs,
low_cpu_mem_usage=True,
torch_dtype=torch.float32,
)
load_model = EMAModel(load_model.parameters(), model_cls=generator_transformer_cls, model_config=load_model.config)
load_model.load_state_dict(ema_kwargs)
ema_generator_transformer3d.load_state_dict(load_model.state_dict())
ema_generator_transformer3d.to(accelerator.device)
del load_model
accelerator.register_save_state_pre_hook(save_ema_hook)
accelerator.register_load_state_pre_hook(load_ema_hook)
if args.gradient_checkpointing:
generator_transformer3d.enable_gradient_checkpointing()
fake_score_transformer3d.enable_gradient_checkpointing()
@@ -1646,6 +1808,12 @@ def main():
print(f"Set {name} to lr : {args.learning_rate_critic / 2}")
break
if use_gan:
fake_trainable_params_optim[0]['params'].extend(
param for name, param in fake_score_transformer3d.named_parameters()
if name.startswith("gan_head.") and name not in in_already
)
if args.use_came:
optimizer = optimizer_cls(
trainable_params_optim,
@@ -1687,7 +1855,7 @@ def main():
args.random_hw_adapt = False
# Get the dataset
if args.train_mode != "normal" or args.use_teacher_forcing:
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
train_dataset = ImageVideoDataset(
args.train_data_meta, args.train_data_dir,
video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames,
@@ -1715,7 +1883,7 @@ def main():
return length_to_frame_num
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing:
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan:
aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
batch_sampler_generator = torch.Generator().manual_seed(args.seed)
batch_sampler = AspectRatioBatchImageVideoSampler(
@@ -2058,6 +2226,9 @@ def main():
if args.train_mode != "normal":
clip_image_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if args.use_ema:
ema_generator_transformer3d.to(accelerator.device)
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
if overrode_max_train_steps:
@@ -2148,6 +2319,8 @@ def main():
vae_stream_2 = None
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
min_score_timestep = int(0.02 * noise_scheduler.config.num_train_timesteps)
max_score_timestep = int(0.98 * noise_scheduler.config.num_train_timesteps)
def randomize_denoising_step_indices(
denoising_step_indices_list,
@@ -2208,6 +2381,8 @@ def main():
# must be averaged by contribution count, not by gradient_accumulation_steps.
train_gen_log_count = 0
train_denoising_loss = 0.0
train_gan_g_loss = 0.0
train_gan_d_loss = 0.0
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
for step, batch in enumerate(train_dataloader):
# Data batch sanity check
@@ -2340,7 +2515,7 @@ def main():
pixel_values_bs = pixel_values_bs.sample()
new_pixel_values.append(pixel_values_bs)
return torch.cat(new_pixel_values, dim = 0)
if args.use_teacher_forcing:
if args.use_teacher_forcing or use_gan:
clean_latents = _batch_encode_vae(pixel_values)
else:
clean_latents = None
@@ -2372,7 +2547,7 @@ def main():
clip_context.append(_clip_context)
clip_context = torch.cat(clip_context)
if args.use_teacher_forcing:
if args.use_teacher_forcing or use_gan:
target_shape = clean_latents.size()
else:
target_shape = mask_latents.size()
@@ -3120,6 +3295,8 @@ def main():
# Common code for both KV cache and block mask training
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep)
generator_denoised_input = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
@@ -3225,17 +3402,53 @@ def main():
train_dmd_loss += avg_dmd_loss.item()
train_gen_log_count += 1
# --- One-Forcing Generator GAN Loss ---
generator_loss = dmd_loss
if args.gan_g_weight > 0:
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
gan_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
gan_timestep = gan_timestep.clamp(min_score_timestep, max_score_timestep)
gan_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng)
gan_inputs = [add_noise(generator_pred, gan_noise, gan_timestep)]
if args.relativistic_discriminator:
real_latents = clean_latents[:, :, :generator_pred.size(2)].to(generator_pred.dtype)
gan_inputs.append(add_noise(real_latents, gan_noise, gan_timestep))
num_gan_inputs = len(gan_inputs)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, gan_logits = fake_score_transformer3d(
x=torch.cat(gan_inputs).to(dtype=weight_dtype),
context=repeat_batch(prompt_embeds, num_gan_inputs),
t=gan_timestep.repeat(num_gan_inputs),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, num_gan_inputs) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, num_gan_inputs) if args.train_mode != "normal" else None,
gan_classify=True,
)
gan_logits = gan_logits.float().chunk(num_gan_inputs)
if args.relativistic_discriminator:
gan_g_loss = F.softplus(-(gan_logits[0] - gan_logits[1])).mean()
else:
gan_g_loss = F.softplus(-gan_logits[0]).mean()
gan_g_loss = gan_g_loss * args.gan_g_weight
train_gan_g_loss += accelerator.gather(gan_g_loss.detach().repeat(args.train_batch_size)).mean().item()
generator_loss = generator_loss + gan_g_loss
if args.low_vram:
real_score_transformer3d = real_score_transformer3d.to("cpu")
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
if args.gan_g_weight == 0:
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
torch.cuda.empty_cache()
accelerator.backward(dmd_loss)
accelerator.backward(generator_loss)
if accelerator.sync_gradients:
accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
if args.gan_g_weight > 0:
# Discard discriminator gradients from the generator GAN backward.
critic_optimizer.zero_grad()
if args.low_vram:
fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device)
@@ -3626,6 +3839,8 @@ def main():
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
critic_timestep = critic_timestep.clamp(min_score_timestep, max_score_timestep)
critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng)
fake_score_denoised_input = add_noise(
@@ -3634,15 +3849,25 @@ def main():
critic_timestep
)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
)
# Generated and real latents share critic_noise / critic_timestep. Prepare the
# real input and the R1/R2 perturbations once; they are consumed by the critic
# forward below and by the R1/R2 penalties after the main backward(s).
real_noisy = None
r1_perturbation = None
r2_perturbation = None
if use_gan:
real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype)
real_noisy = add_noise(real_latents, critic_noise, critic_timestep)
del real_latents
if args.r1_weight > 0:
r1_perturbation = args.r1_sigma * torch.randn(
real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng
)
if args.r2_weight > 0:
r2_perturbation = args.r2_sigma * torch.randn(
fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype,
device=fake_score_denoised_input.device, generator=torch_rng
)
def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
noise_pred = noise_pred.float()
@@ -3656,11 +3881,136 @@ def main():
final_loss = masked_loss.mean()
return final_loss
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
accelerator_fake_score_transformer3d.backward(denoising_loss)
# --- One-Forcing Discriminator GAN Loss ---
if use_gan and not args.relativistic_discriminator:
# Memory-optimized critic: the non-relativistic gan_d_loss =
# softplus(-real) + softplus(fake) is separable, so run fake and real as two
# sequential 1x-batch forwards and backward the fake pass first. Backwarding
# fake frees its gradient-checkpointing block inputs before the real forward,
# dropping the packed fake+real+R1+R2 forward peak (up to 4x batch) down to
# 1x batch here; R1/R2 re-forward their own 2x batch below. Both backward(s)
# accumulate into fake_score.grad, giving the identical update as the packed
# single-forward critic.
# -- Fake pass (1x batch): denoising output + fake discriminator logit --
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output, fake_gan_logits = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
gan_classify=True,
)
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
fake_logit = fake_gan_logits.float()
fake_d_loss = F.softplus(fake_logit).mean() * args.gan_d_weight
critic_loss = denoising_loss + fake_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
# Fake graph freed: the transformer blocks' stored inputs are released here.
fake_d_detached = fake_d_loss.detach()
del critic_loss, fake_gan_logits, fake_logit, fake_d_loss, fake_score_denoised_output
# -- Real pass (1x batch): discriminator logit only --
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, real_gan_logits = fake_score_transformer3d(
x=real_noisy,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
gan_classify=True,
)
real_logit = real_gan_logits.float()
real_d_loss = F.softplus(-real_logit).mean() * args.gan_d_weight
accelerator_fake_score_transformer3d.backward(real_d_loss)
# Real graph freed. Combine the two detached halves for logging.
gan_d_loss = fake_d_detached + real_d_loss.detach()
train_gan_d_loss += accelerator.gather(gan_d_loss.repeat(args.train_batch_size)).mean().item()
del fake_d_detached, real_gan_logits, real_logit, real_d_loss
else:
# Combined path: 2x-batch forward for relativistic GAN (both logits coupled by
# softplus(-(real - fake))), or 1x-batch forward when GAN is disabled.
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
if use_gan:
fake_score_denoised_output, gan_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0]
gan_logits = gan_logits.float().chunk(2)
else:
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
)
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
critic_loss = denoising_loss
if use_gan:
fake_logit, real_logit = gan_logits[0], gan_logits[1]
if args.relativistic_discriminator:
gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
else:
gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
gan_d_loss = gan_d_loss * args.gan_d_weight
train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
critic_loss = critic_loss + gan_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
del critic_loss
# R1/R2 gradient penalties: each re-forwards its own 2x batch after the main
# backward(s) and finite-differences the perturbed logit against the clean one,
# matching the packed-formula value while keeping the peak at 2x batch.
if use_gan and args.r1_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, r1_logits = fake_score_transformer3d(
x=torch.cat([real_noisy + r1_perturbation, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
r1_logits = r1_logits.float().chunk(2)
r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean()
accelerator_fake_score_transformer3d.backward(r1_loss)
del r1_loss
if use_gan and args.r2_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, r2_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
r2_logits = r2_logits.float().chunk(2)
r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean()
accelerator_fake_score_transformer3d.backward(r2_loss)
del r2_loss
if accelerator_fake_score_transformer3d.sync_gradients:
accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm)
critic_optimizer.step()
@@ -3674,12 +4024,19 @@ def main():
# Checks if the accelerator has performed an optimization step behind the scenes
if accelerator.sync_gradients:
if args.use_ema:
ema_generator_transformer3d.step(generator_transformer3d.parameters())
progress_bar.update(1)
global_step += 1
accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step)
logs = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}
if use_gan:
logs.update({"train_gan_g_loss": train_gan_g_loss / max(train_gen_log_count, 1), "train_gan_d_loss": train_gan_d_loss})
accelerator.log(logs, step=global_step)
train_dmd_loss = 0.0
train_gen_log_count = 0
train_denoising_loss = 0.0
train_gan_g_loss = 0.0
train_gan_d_loss = 0.0
if global_step % args.checkpointing_steps == 0:
if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
@@ -3722,6 +4079,10 @@ def main():
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
with progress_bar.paused():
if args.use_ema:
# Store the generator parameters temporarily and load the EMA parameters to perform inference.
ema_generator_transformer3d.store(generator_transformer3d.parameters())
ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
log_validation(
vae,
text_encoder,
@@ -3734,6 +4095,9 @@ def main():
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original generator parameters.
ema_generator_transformer3d.restore(generator_transformer3d.parameters())
logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**logs)
@@ -3743,6 +4107,10 @@ def main():
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
with progress_bar.paused():
if args.use_ema:
# Store the generator parameters temporarily and load the EMA parameters to perform inference.
ema_generator_transformer3d.store(generator_transformer3d.parameters())
ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
log_validation(
vae,
text_encoder,
@@ -3755,6 +4123,9 @@ def main():
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original generator parameters.
ema_generator_transformer3d.restore(generator_transformer3d.parameters())
# Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever
# something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto
scripts/wan2.1_flex_forcing/train_distill.sh
+6 -3
View File
@@ -28,15 +28,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--checkpointing_steps=200 \
--learning_rate=2.0e-06 \
--learning_rate_critic=4.0e-07 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_wan2.1_flex_forcing_distill" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -58,5 +59,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--randomize_step_indices \
--flow_euler_rollout \
--use_kv_cache_training \
--use_ema \
--score_timestep_clamp \
--num_frame_per_block=3 \
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
videox_fun/dist/__init__.py
Vendored
+6 -3
View File
@@ -21,10 +21,13 @@ from .longcatvideo_xfuser import (usp_attn_longcatvideo_avatar_forward,
usp_rope_longcatvideo_forward)
from .ltx2_xfuser import (LTX2MultiGPUsAttnProcessor,
LTX2PerturbedMultiGPUsAttnProcessor)
from .minimax_h3_xfuser import MiniMaxH3MultiGPUsAttnProcessor
from .minimax_h3_xfuser import (MiniMaxH3MultiGPUsAttnProcessor,
MiniMaxH3MultiGPUsVSAAttnProcessor)
from .qwen_xfuser import QwenImageMultiGPUsAttnProcessor2_0
from .qwenimage21_xfuser import QwenImage21MultiGPUsAttnProcessor
from .wan_xfuser import (usp_attn_forward, usp_attn_s2v_forward,
from .wan_xfuser import (usp_attn_forward, usp_attn_fastvideo_forward,
usp_attn_s2v_forward,
usp_attn_self_forcing_forward,
usp_attn_sparkwan_forward)
usp_attn_sparkwan_forward,
usp_attn_turbowan_forward)
from .z_image_xfuser import ZMultiGPUsSingleStreamAttnProcessor
videox_fun/dist/fuser.py
Vendored
+9 -1
View File
@@ -41,7 +41,7 @@ def set_multi_gpus_devices(ulysses_degree, ring_degree, classifier_free_guidance
if get_sp_group is None:
raise RuntimeError("xfuser is not installed.")
dist.init_process_group("nccl")
print('parallel inference enabled: ulysses_degree=%d ring_degree=%d classifier_free_guidance_degree=% rank=%d world_size=%d' % (
print('parallel inference enabled: ulysses_degree=%d ring_degree=%d classifier_free_guidance_degree=%d rank=%d world_size=%d' % (
ulysses_degree, ring_degree, classifier_free_guidance_degree, dist.get_rank(),
dist.get_world_size()))
assert dist.get_world_size() == ring_degree * ulysses_degree * classifier_free_guidance_degree, \
@@ -53,6 +53,14 @@ def set_multi_gpus_devices(ulysses_degree, ring_degree, classifier_free_guidance
ulysses_degree=ulysses_degree)
# device = torch.device("cuda:%d" % dist.get_rank())
device = torch.device(f"cuda:{get_world_group().local_rank}")
# Bind this process to its own GPU. ``init_process_group`` above runs before xfuser's
# ``init_distributed_environment``, whose ``set_device`` is guarded by
# ``if not torch.distributed.is_initialized()`` and so never fires; without this call the
# current device stays 0 on every rank while the model lives on ``cuda:{local_rank}``.
# PyTorch ops dispatch on ``tensor.device`` and are unaffected, but a Triton kernel launches
# on ``torch.cuda.current_device()`` and rejects the rank's own pointers as inaccessible,
# so any Triton sparse attention (e.g. VSA) would fail on rank > 0.
torch.cuda.set_device(device)
print('rank=%d device=%s' % (get_world_group().rank, str(device)))
else:
device = "cuda"
videox_fun/dist/minimax_h3_xfuser.py
Vendored
+98
View File
@@ -1,5 +1,7 @@
import torch
import torch.nn.functional as F
from .fuser import get_sp_group, ulysses_all_to_all
from .fuser import xFuserLongContextAttention
@@ -80,3 +82,99 @@ class MiniMaxH3MultiGPUsAttnProcessor:
hidden_states = attn.to_out[0](hidden_states)
hidden_states = attn.to_out[1](hidden_states)
return hidden_states
class MiniMaxH3MultiGPUsVSAAttnProcessor:
"""Sequence-parallel attention processor for the FastVideo VSA student on MiniMax-H3.
Mirrors ``MiniMaxH3MultiGPUsAttnProcessor``'s row-split -> head-split Ulysses scheme, but swaps
the dense ``xFuserLongContextAttention`` core for the VSA block-sparse op (``attn.attn_op``) so
the trained coarse-to-fine selection and the ``to_gate_compress`` compensation branch are kept
instead of being silently dropped by a dense fallback.
Two VSA-specific points under head-parallelism:
* ``to_gate_compress`` is a ``Linear(dim, heads*head_dim)`` that MIXES heads, so it cannot run
in the head-shard layout. It is strictly per-token and SP shards rows while keeping the full
hidden dim, so the gate is evaluated here on this rank's ROW shard (yielding all heads), then
all-to-all'd alongside q/k/v -- element-wise identical to computing it on the full sequence
and slicing heads, and it lands on the same contiguous head block the all-to-all hands back.
* the parent forward pads the packed sequence up to a multiple of the group size, so the
gathered sequence carries ``<= sp_world_size - 1`` zero rows at its tail. VSA builds its tile
geometry from ``vsa_segments``, whose lengths sum to the PADLESS ``valid_length``; the padding
tail is therefore sliced off before the op and the output padded back afterwards. Unlike the
dense processor (which lets the pad rows dilute the softmax by ~1e-5), this keeps VSA
bitwise-equal to the single-GPU run -- the pad rows never enter the kernel.
Requires a padless packed sequence (``attention_mask is None``) and a pure-Ulysses split
(``num_heads`` divisible by the group size, checked by the model's ``enable_multi_gpus_inference``).
A dense student (``topk >= 1.0``) carries no ``attn_op`` and must keep the dense processor.
"""
def __call__(
self,
attn,
hidden_states: torch.Tensor,
rotary_emb=None,
attention_mask=None,
valid_length=None,
vsa_segments=None,
) -> torch.Tensor:
from ..models.minimax_h3_transformer3d import \
apply_minimax_h3_rotary_emb
if attention_mask is not None:
raise ValueError(
"MiniMaxH3MultiGPUsVSAAttnProcessor does not support a masked (padded) packed sequence.")
if not hasattr(attn, "attn_op"):
raise RuntimeError(
"The VSA sequence-parallel processor needs attn.attn_op; a dense student (topk>=1.0) "
"must bind MiniMaxH3MultiGPUsAttnProcessor instead.")
query = attn.to_q(hidden_states).unflatten(-1, (attn.heads, -1))
key = attn.to_k(hidden_states).unflatten(-1, (attn.heads, -1))
value = attn.to_v(hidden_states).unflatten(-1, (attn.heads, -1))
query = attn.norm_q(query)
key = attn.norm_k(key)
# The rotary coordinates are this rank's rows, so apply them before the all-to-all moves rows.
if rotary_emb is not None:
query = apply_minimax_h3_rotary_emb(query, *rotary_emb)
key = apply_minimax_h3_rotary_emb(key, *rotary_emb)
# Cross-head compress gate on this rank's ROW shard (full hidden dim present): [B, chunk, H, D].
gate = attn.attn_op.to_gate_compress(hidden_states).unflatten(-1, (attn.heads, -1))
half_dtypes = (torch.float16, torch.bfloat16)
def half(x):
return x if x.dtype in half_dtypes else x.to(torch.bfloat16)
# row-split -> head-split over the whole (padded) packed sequence, gate alongside q/k/v.
query = ulysses_all_to_all(half(query), scatter_dim=2, gather_dim=1)
key = ulysses_all_to_all(half(key), scatter_dim=2, gather_dim=1)
value = ulysses_all_to_all(half(value), scatter_dim=2, gather_dim=1)
gate = ulysses_all_to_all(half(gate), scatter_dim=2, gather_dim=1)
# Drop the SP padding tail so VSA sees exactly the padless length its segments describe.
padded_length = query.shape[1]
trim = valid_length is not None and valid_length < padded_length
if trim:
query = query[:, :valid_length]
key = key[:, :valid_length]
value = value[:, :valid_length]
gate = gate[:, :valid_length]
out = attn.attn_op(query, key, value, None, segments=vsa_segments, gate=gate)
# Pad the output back to the gathered length before the inverse all-to-all; the zero rows are
# dropped by the model's final all_gather and never touch a real row's attention.
if trim:
out = F.pad(out, (0, 0, 0, 0, 0, padded_length - valid_length))
# head-split -> row-split, back to this rank's rows.
out = ulysses_all_to_all(out.to(hidden_states.dtype), scatter_dim=1, gather_dim=2)
out = out.flatten(2, 3).type_as(query)
out = attn.to_out[0](out)
out = attn.to_out[1](out)
return out
videox_fun/dist/wan_xfuser.py
Vendored
+138 -1
View File
@@ -224,7 +224,10 @@ def usp_attn_sparkwan_forward(self,
q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
x = self.attn_op(q, k, v)
# heads_total/head_start keep the SLA linear branch's fp32 normalizer sum bitwise-equal to
# the single-GPU result despite this rank holding only the H/P head block.
h_per = n // sp_size
x = self.attn_op(q, k, v, heads_total=n, head_start=sp_rank * h_per)
x = ulysses_all_to_all(x, scatter_dim=1, gather_dim=2)
x = x.flatten(2)
x = self.o(x)
@@ -267,6 +270,140 @@ def usp_attn_sparkwan_forward(self,
out = self.o(out)
return out
def usp_attn_turbowan_forward(self,
x,
seq_lens,
grid_sizes,
freqs,
dtype=torch.bfloat16,
t=0):
"""Sequence-parallel self-attention for the TurboDiffusion SLA student.
``TurboWanTransformer3DModel`` swaps ``WanSelfAttention`` for a local variant whose kernel
is ``self.attn_op`` (``FakeLocalAttention`` -> ``SparseLinearAttention``, block-sparse exact
softmax + a ``proj_l`` feature-map linear branch). The stock ``usp_attn_forward`` binds
``xFuserLongContextAttention``'s dense full-softmax and only touches the base-class
``q/k/v/norm_q/norm_k/o`` attributes, so patching it onto this model would silently drop the
SLA kernel -- every needed attribute is inherited, hence no AttributeError, no shape error,
just dense numbers. This variant keeps ``self.attn_op`` live under Ulysses head-parallelism.
Equivalence: ``SparseLinearAttention.forward`` takes and returns ``[B, L, H, D]``, transposes
to BHLD internally and does every reduction over L/D only (``get_block_map`` pooling and
top-k, the sparse kernel, the ``k^T v`` linear branch), and ``proj_l`` is a per-head
``head_dim -> head_dim`` projection -- no cross-head mixing and no head-indexed parameter,
so an all-to-all from the sequence shard ``[B, S/P, H, D]`` to the head shard
``[B, S, H/P, D]`` reproduces the single-GPU result exactly (top-k is a per-query key-block
ratio). RoPE is applied BEFORE the all-to-all because ``rope_apply`` slices the frequency
table by sequence-parallel rank.
Requires a pure-Ulysses split (``ring_degree`` must be 1: the SLA kernel exposes no partial
output / logsumexp interface, so ring's online rescale-concat cannot be wired to it) and
``num_heads`` divisible by the SP degree.
"""
# No sparse kernel on this instance: keep the stock dense SP path.
if not hasattr(self, "attn_op"):
return usp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=dtype, t=t)
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
half_dtypes = (torch.float16, torch.bfloat16)
def half(t_):
return t_ if t_.dtype in half_dtypes else t_.to(dtype)
# query, key, value function
def qkv_fn(x):
q = self.norm_q(self.q(x.to(dtype))).view(b, s, n, d)
k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d)
v = self.v(x.to(dtype)).view(b, s, n, d)
return q, k, v
q, k, v = qkv_fn(x)
q, k = rope_apply_qk(q, k, grid_sizes, freqs)
# seq-shard -> head-shard, SLA over the FULL sequence on a subset of heads, then back.
q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
# heads_total/head_start keep the SLA linear branch's fp32 normalizer sum bitwise-equal to the
# single-GPU result despite this rank holding only the H/P head block.
sp_size = get_sequence_parallel_world_size()
sp_rank = get_sequence_parallel_rank()
h_per = n // sp_size
x = self.attn_op(q, k, v, heads_total=n, head_start=sp_rank * h_per)
x = ulysses_all_to_all(x.to(dtype), scatter_dim=1, gather_dim=2)
# output
x = x.flatten(2)
x = self.o(x)
return x
def usp_attn_fastvideo_forward(self,
x,
seq_lens,
grid_sizes,
freqs,
dtype=torch.bfloat16,
t=0):
"""Sequence-parallel self-attention for the FastVideo VSA student.
Same Ulysses head-parallel scheme as ``usp_attn_sparkwan_forward``, with one VSA-specific
wrinkle: ``VSASparseAttention.to_gate_compress`` is a ``Linear(dim, num_heads*head_dim)``
that MIXES heads, so it cannot run inside the head-shard layout the way the rest of the VSA
kernel can (``sparse`` / ``compressed`` only reduce over L/D and are head-independent).
It is, however, strictly per-token -- and sequence parallelism shards the sequence while
leaving the hidden dim intact, so every rank still holds the full ``C`` for its own tokens.
The gate is therefore evaluated here on the SEQUENCE shard (yielding all ``H`` heads), then
all-to-all'd alongside q/k/v, which is element-wise identical to computing it on the full
sequence and slicing heads -- and lands on exactly the contiguous head block
``[sp_rank*H/P, (sp_rank+1)*H/P)`` the all-to-all hands this rank.
The gathered sequence is ``s * sp_size``, which the parent forward pads to a multiple of the
SP degree, so it can exceed ``grid_sizes.prod()``; ``_vsa_dispatch`` routes that case to its
per-sample branch (attend over the valid grid tokens, pad the output back), keeping the tile
geometry built from the true grid.
Requires a pure-Ulysses split (``ring_degree`` must be 1) and ``num_heads`` divisible by the
SP degree. ``topk >= 1.0`` (the dense fallback) defers to ``usp_attn_forward``.
"""
# Dense student (topk >= 1.0) or a plain Wan self-attention: keep the stock SP path.
if getattr(self, "topk", 1.0) >= 1.0 or not hasattr(self, "attn_op"):
return usp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=dtype, t=t)
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
half_dtypes = (torch.float16, torch.bfloat16)
def half(t_):
return t_ if t_.dtype in half_dtypes else t_.to(dtype)
# query, key, value function
def qkv_fn(x):
q = self.norm_q(self.q(x.to(dtype))).view(b, s, n, d)
k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d)
v = self.v(x.to(dtype)).view(b, s, n, d)
return q, k, v
q, k, v = qkv_fn(x)
q, k = rope_apply_qk(q, k, grid_sizes, freqs)
# Cross-head compress gate, evaluated on this rank's sequence shard where the full hidden
# dim is available: [B, S/P, C] -> [B, S/P, H, D] (all heads).
gate = self.attn_op.to_gate_compress(x.to(dtype)).view(b, s, n, d)
# seq-shard -> head-shard for q/k/v and the gate alike.
q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
gate = ulysses_all_to_all(half(gate), scatter_dim=2, gather_dim=1)
# VSA over the FULL sequence on a head shard; x is unused once gate is supplied.
out = self._vsa_dispatch(q, k, v, None, seq_lens, grid_sizes, gate=gate)
# head-shard -> seq-shard, back to this rank's sequence chunk.
out = ulysses_all_to_all(out.to(dtype), scatter_dim=1, gather_dim=2)
out = out.flatten(2)
out = self.o(out)
return out
@amp.autocast(enabled=False)
@torch.compiler.disable()
def s2v_rope_apply(x, grid_sizes, freqs):
videox_fun/models/attention_kernel.py
+11 -18
View File
@@ -51,17 +51,6 @@ import triton.language as tl
# loads, a fused tl.dot(p, v, acc) accumulate, and a per-query-tile variable kv-block count (vs a
# uniform-MAXK loop over a dummy-padded tile).
#
# This is deliberately NOT the fastest forward available. That same sm90 CUDA kernel (ThunderKittens)
# IS buildable from source on this box -- only its prebuilt wheels are pinned -- and measures
# 4.66 ms/call on the workload above, i.e. 1.35x over this Triton kernel (step-level VSA-vs-dense
# 1.019x -> 1.115x at 480P, 1.233x -> 1.381x at 720P). It stays out of the production path because
# it adds a build dependency and fails SILENTLY three ways: on a non-sm90a device pass it compiles
# to an empty stub returning uninitialised memory; it reinterpret_casts `variable_block_sizes` to
# int32 with no dtype check, so an int64 tensor yields all-NaN; and its lse is [B,H,L,1] in the
# natural-log domain scaled by -1/qk_scale rather than our [B,H,L] base-2. The build recipe, the
# runtime gate it would need, and the measurement set are archived in
# `_fastvideo_sm90_kernel/README.md`.
#
# Deviations from upstream, listed so a diff against the sdist still audits this block:
# 1. its `import torch / triton / triton.language`, and its SPARSE-ADDITION `import math`, are
# hoisted to this module's top import block -- same names, no behavioural change, and the
@@ -495,12 +484,10 @@ class _sparse_linear_attention_varlen(torch.autograd.Function):
@staticmethod
def forward(ctx, q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes, kb_id, qk_scale):
assert q_tile.is_contiguous() and k_tile.is_contiguous() and v_tile.is_contiguous()
B, H, L_BUF, D = q_tile.shape
# Forward runs the vendored official Triton block-sparse kernel, which consumes the
# official (q2k_index, q2k_num) pair map_to_index produced: the first q2k_num[row] entries
# of each row are that query tile's kv-tile ids and the -1 pads are never read. Its M is
# [B,H,L_BUF] base-2 lse (m_i + log2 l_i), byte-compatible with what our backward expects,
# so the save_for_backward contract is unchanged.
# The forward consumes the official (q2k_index, q2k_num) pair map_to_index produced:
# the first q2k_num[row] entries of each row are that query tile's kv-tile ids and the
# -1 pads are never read. It produces M as [B,H,L_BUF] base-2 lse (m_i + log2 l_i),
# which is exactly what the varlen backward kernels below consume.
out, lse = triton_block_sparse_attn_forward(
q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes)
ctx.save_for_backward(q_tile, k_tile, v_tile, q2k_idx, q2k_num,
@@ -625,7 +612,13 @@ def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
K = pooled_score.shape[-1]
topk = min(K, int(topk_ratio * K))
# Clamp to at least one key block. A truncating int() selects ZERO blocks whenever
# topk_ratio * K < 1 (short sequence and/or aggressive ratio -- e.g. RoLa's default 0.03
# below 2176 tokens, SLA's 0.1 below 640), and the sparse kernel then attends to nothing and
# emits all-NaN. Keeps the truncating count everywhere else, so no existing configuration
# changes its selection; matches the lower bound the VSA paths already apply
# (VSASparseAttention.sparse / _h3_compute_topk).
topk = max(1, min(K, int(topk_ratio * K)))
lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices
sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)
videox_fun/models/attention_utils.py
+50 -7
View File
@@ -345,6 +345,32 @@ def attention(
return out
def _head_shard_seq_sum(x, heads_total, head_start):
"""``torch.sum(x, dim=-2, keepdim=True)`` that is bitwise-invariant to Ulysses head-sharding.
``x`` is ``[B, h_local, L, D]``. The pipeline runs under autocast and ``torch.sum`` is in
PyTorch's float32 promotion list, so this normalizer accumulates in fp32 over ``L`` and
returns fp32; the CUDA ReduceConfig picks its accumulation tree from the output count
``B*H*D``. When head-parallelism shrinks ``H`` (e.g. 12 -> 6 -> 3) the tree changes and the
fp32 sum moves in its last bits (~1e-7), which the linear branch's denominator passes through
undamped -- a bf16 output would absorb it, fp32 does not. Padding the head dim back to
``heads_total`` makes every rank reduce over the full ``B*heads_total*D`` outputs, i.e. the
single-GPU tree, then this shard's own head block is sliced back out, so multi-GPU reproduces
the single-GPU result bitwise.
Uninitialized padding is safe: each output element reads only its own head's rows, so the
garbage heads (sliced away) never touch the real ones. ``heads_total=None`` or an already-full
head count falls back to the plain reduction, leaving the single-GPU path unchanged.
"""
h_local = x.shape[1]
if heads_total is None or h_local == heads_total:
return torch.sum(x, dim=-2, keepdim=True)
buf = torch.empty(x.shape[0], heads_total, x.shape[2], x.shape[3],
device=x.device, dtype=x.dtype)
buf[:, head_start:head_start + h_local].copy_(x)
return torch.sum(buf, dim=-2, keepdim=True)[:, head_start:head_start + h_local]
class SparseLinearAttention(nn.Module):
# Modified from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/kernel.py
def __init__(self, head_dim, topk, feature_map='softmax', BLKQ=64, BLKK=64, use_bf16=True, tie_feature_map_qk=True):
@@ -391,13 +417,18 @@ class SparseLinearAttention(nn.Module):
nn.init.zeros_(self.proj_l.weight)
nn.init.zeros_(self.proj_l.bias)
def forward(self, q, k, v, return_sparsity=False):
def forward(self, q, k, v, return_sparsity=False, heads_total=None, head_start=0):
R'''
Args:
q: queries [B, L, H, D]
k: keys [B, L, H, D].
v: values [B, L, H, D].
return_sparsity: whether to return the actual sparsity.
heads_total: full head count before Ulysses head-sharding, or None when unsharded.
Kept so the linear branch's fp32 normalizer sum stays bitwise-equal to the
single-GPU result under sequence parallelism (see _head_shard_seq_sum).
head_start: this shard's offset into the full head range; used only when heads_total
is not None.
'''
dtype = q.dtype
@@ -417,7 +448,7 @@ class SparseLinearAttention(nn.Module):
def calc_linear(q, k, v):
kvsum = k.transpose(-1, -2) @ v
ksum = torch.sum(k, dim=-2, keepdim=True)
ksum = _head_shard_seq_sum(k, heads_total, head_start)
return (q @ kvsum) / (1e-5 + (q * ksum).sum(dim=-1, keepdim=True))
o_l = calc_linear(q, k, v)
@@ -943,7 +974,7 @@ class VSASparseAttention(nn.Module):
oc = oc_blocks.repeat_interleave(self.BLKQ, dim=2)[:, :, :L]
return oc
def forward_h3(self, q, k, v, x, segments, return_sparsity=False):
def forward_h3(self, q, k, v, x, segments, return_sparsity=False, gate=None):
R'''Segment-aware H3 VSA -- the official FastVideo segment/tile mask semantics.
Args:
@@ -951,6 +982,12 @@ class VSASparseAttention(nn.Module):
x: [B, L, C] attention-input hidden states (drives the compress gate).
segments: packed layout -- dense row counts (``int``) and video region
``(t, h, w)`` token grid, in packed order. Built by the model forward.
gate: [B, L, H, D] pre-computed ``to_gate_compress`` output; when given, ``x`` is
ignored. ``to_gate_compress`` is a ``Linear(dim, num_heads*head_dim)`` and thus
mixes heads, so under Ulysses head-parallelism the caller must evaluate it on
its own sequence shard (where the FULL hidden dim is still present) and
all-to-all the result -- passing the head shard in here keeps this module
head-local. ``H`` is then the local head count.
Returns:
[B, L, H, D] attention output (before the output projection).
@@ -985,7 +1022,8 @@ class VSASparseAttention(nn.Module):
qb = _tile(q, self.dtype, 'q')
kb = _tile(k, self.dtype, 'k')
vb = _tile(v, self.dtype, 'v')
gate = self.to_gate_compress(x).view(B, L, H, D)
gate = self.to_gate_compress(x) if gate is None else gate
gate = gate.view(B, L, H, D)
gb = _tile(gate, gate.dtype, 'g')
pq = _h3_pool_tiles(qb, vbs, tile_elems, n_tiles) # [B,H,n_tiles,D] fp32
@@ -1017,12 +1055,15 @@ class VSASparseAttention(nn.Module):
return out, self.topk
return out
def forward(self, q, k, v, x, return_sparsity=False, segments=None):
def forward(self, q, k, v, x, return_sparsity=False, segments=None, gate=None):
R'''
Args:
q, k, v: [B, L, H, D] POST-RoPE queries/keys/values.
x: [B, L, C] attention-input hidden states (drives the compress gate).
return_sparsity: also return the sparse-branch keep ratio.
gate: [B, L, H, D] pre-computed ``to_gate_compress`` output; see ``forward_h3``
(required under Ulysses head-parallelism, where ``x`` is unavailable in the
head-shard layout because the gate Linear mixes heads).
Returns:
[B, L, H, D] attention output (before the output projection).
'''
@@ -1031,7 +1072,8 @@ class VSASparseAttention(nn.Module):
# kernel (forward + backward), so training and inference share one selection
# semantics. Falls back to the flat 1-D block map only when the caller cannot
# supply a clean segment layout (segments is None).
return self.forward_h3(q, k, v, x, segments, return_sparsity=return_sparsity)
return self.forward_h3(q, k, v, x, segments, return_sparsity=return_sparsity,
gate=gate)
B, L, H, D = q.shape
qb = q.transpose(1, 2).contiguous() # [B,H,L,D]
kb = k.transpose(1, 2).contiguous()
@@ -1042,7 +1084,8 @@ class VSASparseAttention(nn.Module):
# Shorter than one block: exact dense attention (sparsity gives no benefit).
o_s = F.scaled_dot_product_attention(
qb.float(), kb.float(), vb.float()).to(qb.dtype)
gate = self.to_gate_compress(x).view(B, L, H, D).transpose(1, 2) # [B,H,L,D]
gate = self.to_gate_compress(x) if gate is None else gate
gate = gate.view(B, L, H, D).transpose(1, 2) # [B,H,L,D]
o_c = self.compressed(qb, kb, vb)
out = o_s + gate.to(o_s.dtype) * o_c.to(o_s.dtype)
out = out.transpose(1, 2) # [B,L,H,D]
videox_fun/models/fastvideo_minimax_h3_transformer3d.py
+15 -8
View File
@@ -240,16 +240,23 @@ class FastVideoMiniMaxH3Transformer3DModel(MiniMaxH3Transformer3DModel):
return tuple(segments)
def enable_multi_gpus_inference(self,):
"""VSA has no sequence-parallel path in this repo.
"""Enable sequence-parallel inference for the VSA sparse student.
The in-repo VSA kernel only implements the single-GPU coarse-to-fine block-sparse forward;
there is no ulysses head-parallel or ring-attention binding for it (unlike the base
``MiniMaxH3MultiGPUsAttnProcessor``). Fail loudly rather than silently falling back to a
DENSE multi-GPU processor (which would bypass VSA and drop the gate params).
Reuses the parent's packed-sequence SP assembly -- it sizes ``sp_world_size`` /
``all_gather``, pads the sequence to a group multiple, hands every rank a contiguous row
slice and checks ``num_heads`` is divisible by the group -- then re-binds the VSA
sequence-parallel processor onto every block that carries an ``attn_op``. The parent
binding (``MiniMaxH3MultiGPUsAttnProcessor``) is a DENSE core; leaving it on a sparse
block would silently bypass VSA and drop the ``to_gate_compress`` params, so the sparse
blocks are swapped to ``MiniMaxH3MultiGPUsVSAAttnProcessor``, which turns the row split
into a head split around the VSA op. A dense student (``topk >= 1.0``) has no ``attn_op``
and keeps the parent's dense processor untouched.
"""
raise NotImplementedError(
"FastVideoMiniMaxH3Transformer3DModel (VSA) supports single-GPU inference only; the "
"VSA kernel has no ulysses/ring sequence-parallel path. Set ulysses_degree=ring_degree=1.")
super().enable_multi_gpus_inference()
from ..dist import MiniMaxH3MultiGPUsVSAAttnProcessor
for block in self.transformer_blocks:
if hasattr(block.attn, "attn_op"):
block.attn.set_processor(MiniMaxH3MultiGPUsVSAAttnProcessor())
def sparse_trainable_params(self):
"""Names of the extra VSA-branch parameters, for logging / freezing."""
videox_fun/models/fastvideo_wan_transformer3d.py
+72 -25
View File
@@ -67,26 +67,45 @@ class FastVideoWanSelfAttention(WanSelfAttention):
q, k = rope_apply_qk(q, k, grid_sizes, freqs)
out = self._vsa_dispatch(q.to(dtype), k.to(dtype), v.to(dtype), x.to(dtype),
seq_lens, grid_sizes)
x = out.to(dtype).flatten(2)
return self.o(x)
def _vsa_dispatch(self, q, k, v, x, seq_lens, grid_sizes, gate=None):
"""Run the VSA kernel over POST-RoPE ``q/k/v`` ([B, L, H, D]) -> [B, L, H, D].
``x`` drives ``attn_op.to_gate_compress`` unless a pre-computed ``gate`` ([B, L, H, D])
is handed in: that Linear mixes heads, so the sequence-parallel path
(``usp_attn_fastvideo_forward``) evaluates it on its own sequence shard and passes the
head-sharded result, with ``x`` left as None.
The segment-aware tile geometry needs the sequence length to equal the true video token
count. Under sequence parallelism the gathered length is padded to a multiple of the SP
degree and can exceed it, so that case takes the per-sample branch below, which attends
over the valid grid tokens only and pads the output back to ``L``.
"""
s = q.shape[1]
g0 = grid_sizes[0]
if (bool((grid_sizes == g0).all()) and s == int(g0.prod())
and bool((seq_lens == s).all())):
segments = ((int(g0[0]), int(g0[1]), int(g0[2])),)
out = self.attn_op(q.to(dtype), k.to(dtype), v.to(dtype), x.to(dtype),
segments=segments)
else:
outputs = []
for i, (grid, length) in enumerate(zip(grid_sizes.tolist(), seq_lens.tolist())):
grid = tuple(int(axis) for axis in grid)
length = int(length)
if min(grid) <= 0 or length != grid[0] * grid[1] * grid[2] or not 0 < length <= s:
raise ValueError("Wan VSA 的 seq_lens 必须等于有效网格 token 数且不超过序列长度。")
oi = self.attn_op(q[i:i + 1, :length].to(dtype), k[i:i + 1, :length].to(dtype),
v[i:i + 1, :length].to(dtype), x[i:i + 1, :length].to(dtype),
segments=(grid,))
outputs.append(torch.nn.functional.pad(oi, (0, 0, 0, 0, 0, s - length)))
out = torch.cat(outputs, dim=0)
x = out.to(dtype).flatten(2)
return self.o(x)
return self.attn_op(q, k, v, x, segments=segments, gate=gate)
outputs = []
for i, (grid, length) in enumerate(zip(grid_sizes.tolist(), seq_lens.tolist())):
grid = tuple(int(axis) for axis in grid)
length = int(length)
if min(grid) <= 0 or length != grid[0] * grid[1] * grid[2] or not 0 < length <= s:
raise ValueError("Wan VSA 的 seq_lens 必须等于有效网格 token 数且不超过序列长度。")
gate_i = None if gate is None else gate[i:i + 1, :length]
# x is only consumed by to_gate_compress, so it is dropped once gate is supplied.
x_i = None if gate is not None else x[i:i + 1, :length]
oi = self.attn_op(q[i:i + 1, :length], k[i:i + 1, :length],
v[i:i + 1, :length], x_i,
segments=(grid,), gate=gate_i)
outputs.append(torch.nn.functional.pad(oi, (0, 0, 0, 0, 0, s - length)))
return torch.cat(outputs, dim=0)
class FastVideoWanTransformer3DModel(WanTransformer3DModel):
@@ -187,17 +206,45 @@ class FastVideoWanTransformer3DModel(WanTransformer3DModel):
self._apply_sparse_self_attention(topk)
def enable_multi_gpus_inference(self,):
"""VSA has no sequence-parallel path in this repo.
"""Bind the VSA-aware USP forward instead of the parent's dense one.
The in-repo VSA kernel only implements the single-GPU coarse-to-fine block-sparse
forward; there is no ulysses head-parallel or ring-attention binding for it (unlike
the RoLa/SLA ``usp_attn_*`` forwards). The 1.3B student fits comfortably on one GPU,
so run inference single-device. Fail loudly rather than silently falling back to the
parent's DENSE ``usp_attn_forward`` (which would bypass VSA and drop the gate params).
``WanTransformer3DModel.enable_multi_gpus_inference`` monkey-patches every
``blocks[i].self_attn.forward`` to ``usp_attn_forward`` (dense full-softmax via
``xFuserLongContextAttention``), which never reaches ``FastVideoWanSelfAttention.attn_op``
and would silently drop the VSA kernel plus the trained ``to_gate_compress`` gate.
Rebind to ``usp_attn_fastvideo_forward``: Ulysses head-parallel, with the cross-head
compress gate evaluated on each rank's sequence shard (where the full hidden dim still
lives) and all-to-all'd next to q/k/v.
Constraint: the VSA kernel has no ring-attention path (no partial-output/logsumexp
interface to rescale across KV chunks), so this is pure Ulysses -- ``ring_degree`` must
be 1 and ``num_heads`` must divide the SP degree.
"""
raise NotImplementedError(
"FastVideoWanTransformer3DModel (VSA) supports single-GPU inference only; the VSA "
"kernel has no ulysses/ring sequence-parallel path. Set ulysses_degree=ring_degree=1.")
import types
from ..dist import (get_sequence_parallel_rank,
get_sequence_parallel_world_size, get_sp_group,
usp_attn_fastvideo_forward)
self.sp_world_size = get_sequence_parallel_world_size()
self.sp_world_rank = get_sequence_parallel_rank()
self.all_gather = get_sp_group().all_gather
# Ulysses head-split needs num_heads divisible by the SP degree.
n_heads = self.blocks[0].self_attn.num_heads
if self.sp_world_size > 1 and n_heads % self.sp_world_size != 0:
raise ValueError(
f"VSA multi-GPU (ulysses) needs num_heads ({n_heads}) divisible by the "
f"sequence-parallel degree ({self.sp_world_size}); set ulysses_degree "
f"accordingly and keep ring_degree=1 (the VSA kernel has no ring path).")
for block in self.blocks:
block.self_attn.forward = types.MethodType(
usp_attn_fastvideo_forward, block.self_attn)
if hasattr(self, 'vace_blocks'):
for block in self.vace_blocks:
block.self_attn.forward = types.MethodType(
usp_attn_fastvideo_forward, block.self_attn)
def sparse_trainable_params(self):
"""Names of the extra VSA-branch parameters, for logging / freezing."""
videox_fun/models/minimax_h3_conversion.py
+5 -2
View File
@@ -111,10 +111,13 @@ class MiniMaxH3MixedPrecisionLoaderMixin:
kwargs["device_map"] = device_map
if max_memory is not None:
kwargs["max_memory"] = max_memory
model = super().from_pretrained(pretrained_model_name_or_path, torch_dtype=torch_dtype, **kwargs)
result = super().from_pretrained(pretrained_model_name_or_path, torch_dtype=torch_dtype, **kwargs)
# ModelMixin returns (model, loading_info) when output_loading_info=True.
# Restore precision on the model while preserving that return contract.
model = result[0] if isinstance(result, tuple) else result
if torch_dtype not in (None, torch.float32) and cls._keep_in_fp32_modules:
restore_fp32_modules(model, pretrained_model_name_or_path, kwargs.get("subfolder", None))
return model
return result
@classmethod
def _from_pretrained_with_missing_keys(cls, pretrained_model_name_or_path, torch_dtype, **kwargs):
videox_fun/models/minimax_h3_transformer3d.py
+106 -17
View File
@@ -13,7 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from contextlib import nullcontext
from contextlib import contextmanager, nullcontext
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, Union
@@ -264,7 +264,11 @@ class MiniMaxH3Attention(nn.Module):
rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
attention_mask: Optional[torch.Tensor] = None,
valid_length: Optional[int] = None,
vsa_segments=None,
) -> torch.Tensor:
if vsa_segments is not None:
return self.processor(self, hidden_states, rotary_emb, attention_mask, valid_length,
vsa_segments=vsa_segments)
return self.processor(self, hidden_states, rotary_emb, attention_mask, valid_length)
@@ -374,6 +378,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
rotary_emb: Tuple[torch.Tensor, torch.Tensor],
attention_mask: Optional[torch.Tensor] = None,
valid_length: Optional[int] = None,
vsa_segments=None,
) -> torch.Tensor:
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(temb)
@@ -382,7 +387,8 @@ class MiniMaxH3TransformerBlock(nn.Module):
norm_hidden_states = norm_hidden_states * (
1.0 + _index_select_rows(scale_msa, adaln_indices)
) + _index_select_rows(shift_msa, adaln_indices)
attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask, valid_length)
attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask, valid_length,
vsa_segments=vsa_segments)
hidden_states = residual + _index_select_rows(gate_msa, adaln_indices) * attn_output
residual = hidden_states
@@ -546,6 +552,9 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin
self.gradient_checkpointing = False
self.gradient_checkpointing_save_on_cpu = False
# Optional non-reentrant checkpoint `context_fn` for the main blocks; the streaming KV cache installs one so
# backward recomputes see the cache view of the original forward.
self.gradient_checkpointing_context_fn = None
self.sp_world_size = 1
self.sp_world_rank = 0
self.all_gather = None
@@ -747,10 +756,14 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin
adaln_indices = adaln_indices[rows]
rotary_emb = (rotary_emb[0][rows], rotary_emb[1][rows])
vsa_segments = attention_kwargs.get("vsa_segments") if attention_kwargs else None
block_kwargs = {"vsa_segments": vsa_segments} if vsa_segments is not None else {}
for block in self.transformer_blocks:
with torch.autograd.graph.save_on_cpu() if self.gradient_checkpointing_save_on_cpu else nullcontext():
if torch.is_grad_enabled() and self.gradient_checkpointing:
ckpt_kwargs = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
if self.gradient_checkpointing_context_fn is not None:
ckpt_kwargs["context_fn"] = self.gradient_checkpointing_context_fn
hidden_states = torch.utils.checkpoint.checkpoint(
block,
hidden_states,
@@ -759,10 +772,12 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin
rotary_emb,
attention_mask,
valid_length,
**block_kwargs,
**ckpt_kwargs,
)
else:
hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb, attention_mask, valid_length)
hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb, attention_mask,
valid_length, **block_kwargs)
if self.sp_world_size > 1:
hidden_states = self.all_gather(hidden_states.contiguous(), dim=1)[:, :sequence_length]
@@ -822,11 +837,20 @@ class MiniMaxH3StreamingKVCache:
Args:
num_layers (`int`): The number of main transformer blocks (50 for the released H3).
dtype (`torch.dtype`, defaults to `torch.bfloat16`): The dtype the clean K/V is stored at.
offload_to_cpu (`bool`, defaults to `False`): Whether the clean K/V is parked on the CPU.
commit_attend_history (`bool`, defaults to `False`): Whether a clean commit pass runs the
live attention pattern (media rows attend `[text K/V; persistent history; current chunk
K/V]`, text rows self-attend) instead of the legacy standalone pass over
`[text | current chunk]`. Enable it when the history must condition the committed K/V,
as in the Self-Forcing clean refresh pass.
"""
def __init__(self, num_layers: int, dtype: torch.dtype = torch.bfloat16):
def __init__(self, num_layers: int, dtype: torch.dtype = torch.bfloat16, offload_to_cpu: bool = False,
commit_attend_history: bool = False):
self.layer_names = tuple(f"transformer_blocks.{index}.attn" for index in range(num_layers))
self.dtype = dtype
self.offload_to_cpu = offload_to_cpu
self.commit_attend_history = commit_attend_history
self._history: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
self._staged: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
self._commit_active = False
@@ -838,6 +862,10 @@ class MiniMaxH3StreamingKVCache:
# Per-forward attention mode: None ("plain"), "live" or "commit".
self.mode: Optional[str] = None
self.token_tags: Optional[torch.Tensor] = None
# Activation-checkpoint replay: while set, `history_pair` exposes only this many leading history rows (the
# history a checkpointed forward saw); `_epoch` changes whenever rows are dropped, invalidating pending replays.
self._history_view_rows: Optional[int] = None
self._epoch = 0
# ------------------------------------------------------------------
# Per-forward setup consumed by the streaming processor.
@@ -902,11 +930,51 @@ class MiniMaxH3StreamingKVCache:
self._staged_row_count = int(key.shape[0])
elif int(key.shape[0]) != self._staged_row_count:
raise RuntimeError(f"{layer_name}: staged row count differs across layers")
self._staged[layer_name] = (key, value)
device = torch.device("cpu") if self.offload_to_cpu else key.device
self._staged[layer_name] = (
key.detach().to(device=device, dtype=self.dtype).contiguous(),
value.detach().to(device=device, dtype=self.dtype).contiguous(),
)
def history_pair(self, layer_name: str) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
r"""Return the persistent `(key, value)` of one layer, or `None` when empty."""
return self._history.get(layer_name)
pair = self._history.get(layer_name)
if pair is None or self._history_view_rows is None:
return pair
if self._history_view_rows == 0:
return None
return pair[0][: self._history_view_rows], pair[1][: self._history_view_rows]
# ------------------------------------------------------------------
# Activation checkpointing.
# ------------------------------------------------------------------
def checkpoint_context(self):
r"""
`context_fn` for non-reentrant `torch.utils.checkpoint` over a main block.
Called when the checkpointed forward runs, it snapshots the cache view that forward sees — the mode, the
token tags and the committed history length — and returns `(forward_context, recompute_context)`; the
recompute context restores that view while backward replays the block, after `end_forward()` cleared the
mode and later commits appended more history. The history is append-only (`commit` concatenates), so the
snapshot is a row count, not a tensor copy: the replay reads a prefix view of the current history.
"""
history_rows = self.history_tokens if self._history_view_rows is None else self._history_view_rows
snapshot = (self.mode, self.token_tags, history_rows, self._epoch)
return nullcontext(), self._replay_view(*snapshot)
@contextmanager
def _replay_view(self, mode, token_tags, history_rows, epoch):
if epoch != self._epoch:
raise RuntimeError(
"the streaming KV history was cleared or trimmed between a checkpointed forward and its backward "
"recompute; clear the cache only after backward"
)
saved = (self.mode, self.token_tags, self._history_view_rows)
self.mode, self.token_tags, self._history_view_rows = mode, token_tags, history_rows
try:
yield
finally:
self.mode, self.token_tags, self._history_view_rows = saved
# ------------------------------------------------------------------
# Transaction control.
@@ -931,6 +999,7 @@ class MiniMaxH3StreamingKVCache:
torch.cat((previous[1], current[1]), dim=0),
)
self._history[layer_name] = combined
del current, previous, combined
self._commit_token_counts.append(len(self._commit_tags))
self._commit_token_tags.append(self._commit_tags)
self._commit_active = False
@@ -1005,6 +1074,7 @@ class MiniMaxH3StreamingKVCache:
def clear(self) -> None:
r"""Clear persistent and staged KV and reset the logical chunk count."""
self._history.clear()
self._epoch += 1
self._commit_token_counts.clear()
self._commit_token_tags.clear()
self._staged.clear()
@@ -1014,6 +1084,7 @@ class MiniMaxH3StreamingKVCache:
self.end_forward()
def _retain_commit_rows(self, selection) -> None:
self._epoch += 1
offsets = [0]
for count in self._commit_token_counts:
offsets.append(offsets[-1] + count)
@@ -1068,9 +1139,11 @@ class MiniMaxH3StreamingKVCacheAttnProcessor:
[`MiniMaxH3AttnProcessor`];
* *live*: text rows self-attend, while media rows attend to
`[text K/V; persistent clean history; current chunk K/V]`;
* *commit*: a plain pass whose output is discarded, staging the media rows' K/V of
every layer into the cache (the K/V of a row depends only on that row's input, so
the commit pass's attention output is irrelevant).
* *commit*: the pass whose output is discarded, staging the media rows' K/V of every
layer into the cache. With `commit_attend_history` off (legacy) it is a plain pass over
`[text | current chunk]`; with it on, the commit pass runs the live attention pattern so
the staged K/V is conditioned on the persistent history, matching the Self-Forcing clean
refresh pass.
Live and commit passes must be padless (`attention_mask is None`); the streaming
packed sequences are single documents by construction.
@@ -1107,19 +1180,19 @@ class MiniMaxH3StreamingKVCacheAttnProcessor:
query = apply_minimax_h3_rotary_emb(query, *rotary_emb)
key = apply_minimax_h3_rotary_emb(key, *rotary_emb)
if cache.mode == "commit":
# The committed rows' K/V depend only on their own inputs; the attention output
# of a commit pass is discarded by the caller.
if cache.mode == "commit" and not cache.commit_attend_history:
# Legacy commit: a plain pass over [text | current chunk]; the attention output is
# discarded by the caller and only the staged media K/V rows matter.
output = attention(query, key, value, causal=False)
rows = cache.commit_row_indices.to(key.device)
staged_key = key.index_select(1, rows).squeeze(0)
staged_value = value.index_select(1, rows).squeeze(0)
cache.stage(
self.layer_name,
staged_key.detach().to(cache.dtype).contiguous(),
staged_value.detach().to(cache.dtype).contiguous(),
staged_key,
staged_value,
)
else: # live
else: # live, or a commit with commit_attend_history: media rows attend the history.
tags = cache.token_tags
text_rows = tags == MINIMAX_H3_STREAMING_TEXT_TAG
media_rows = ~text_rows
@@ -1153,6 +1226,16 @@ class MiniMaxH3StreamingKVCacheAttnProcessor:
values.transpose(1, 2),
).transpose(1, 2)
if cache.mode == "commit" and cache.commit_attend_history:
# The history-conditioned commit pass output is still discarded; only the staged K/V
# rows matter, now computed with the history in context.
rows = cache.commit_row_indices.to(key.device)
cache.stage(
self.layer_name,
key.index_select(1, rows).squeeze(0),
value.index_select(1, rows).squeeze(0),
)
output = output.flatten(2, 3).type_as(query)
output = attn.to_out[0](output)
output = attn.to_out[1](output)
@@ -1201,21 +1284,26 @@ def install_minimax_h3_streaming_kv_cache(
transformer: "MiniMaxH3Transformer3DModel",
*,
dtype: torch.dtype = torch.bfloat16,
offload_to_cpu: bool = False,
commit_attend_history: bool = False,
) -> MiniMaxH3StreamingKVCache:
r"""
Swap every main block's attention processor for the streaming one and return the shared
cache.
The token refiner keeps its standard processor: the streaming contract only concerns the
main block stack over the media sequence.
main block stack over the media sequence. `commit_attend_history` is forwarded to
[`MiniMaxH3StreamingKVCache`].
"""
cache = MiniMaxH3StreamingKVCache(
num_layers=len(transformer.transformer_blocks), dtype=dtype
num_layers=len(transformer.transformer_blocks), dtype=dtype, offload_to_cpu=offload_to_cpu,
commit_attend_history=commit_attend_history,
)
for index, block in enumerate(transformer.transformer_blocks):
block.attn.set_processor(
MiniMaxH3StreamingKVCacheAttnProcessor(cache, cache.layer_names[index])
)
transformer.gradient_checkpointing_context_fn = cache.checkpoint_context
return cache
@@ -1223,3 +1311,4 @@ def uninstall_minimax_h3_streaming_kv_cache(transformer: "MiniMaxH3Transformer3D
r"""Restore the standard attention processor on every main block."""
for block in transformer.transformer_blocks:
block.attn.set_processor(MiniMaxH3AttnProcessor())
transformer.gradient_checkpointing_context_fn = None
videox_fun/models/turbowan_transformer3d.py
+44 -2
View File
@@ -51,8 +51,8 @@ class FakeLocalAttention(nn.Module):
super().__init__()
self.local_attn = SparseLinearAttention(head_dim, topk)
def forward(self, q, k, v):
x = self.local_attn(q, k, v)
def forward(self, q, k, v, heads_total=None, head_start=0):
x = self.local_attn(q, k, v, heads_total=heads_total, head_start=head_start)
return x
@@ -310,3 +310,45 @@ class TurboWanTransformer3DModel(WanTransformer3DModel):
if hasattr(self, "img_emb"):
del self.img_emb
def enable_multi_gpus_inference(self,):
"""Bind the SLA-aware USP forward instead of the parent's dense one.
``WanTransformer3DModel.enable_multi_gpus_inference`` monkey-patches every
``blocks[i].self_attn.forward`` to ``usp_attn_forward`` (dense full-softmax via
``xFuserLongContextAttention``). That path only touches base-class attributes
(``q/k/v/norm_q/norm_k/o/window_size/num_heads/head_dim``), all of which this module's
``WanSelfAttention`` inherits, so the inherited patch used to degrade the SLA student to
dense attention SILENTLY -- no AttributeError, no shape error, just ``self.attn_op``
(``FakeLocalAttention`` -> ``SparseLinearAttention``) never being reached. Rebind to
``usp_attn_turbowan_forward``: Ulysses head-parallel all-to-all around the SLA kernel,
which is head-independent (all reductions over L/D, ``proj_l`` is per-head), so the
multi-GPU result matches single-GPU.
Constraint: the SLA kernel exposes no partial-output/logsumexp interface, so ring
attention cannot be wired to it -- pure Ulysses only, ``ring_degree`` must be 1 and
``num_heads`` must divide the SP degree.
"""
from ..dist import (get_sequence_parallel_rank,
get_sequence_parallel_world_size, get_sp_group,
usp_attn_turbowan_forward)
self.sp_world_size = get_sequence_parallel_world_size()
self.sp_world_rank = get_sequence_parallel_rank()
self.all_gather = get_sp_group().all_gather
# Ulysses head-split needs num_heads divisible by the SP degree.
n_heads = self.blocks[0].self_attn.num_heads
if self.sp_world_size > 1 and n_heads % self.sp_world_size != 0:
raise ValueError(
f"SLA multi-GPU (ulysses) needs num_heads ({n_heads}) divisible by the "
f"sequence-parallel degree ({self.sp_world_size}); set ulysses_degree "
f"accordingly and keep ring_degree=1 (the SLA kernel has no ring path).")
for block in self.blocks:
block.self_attn.forward = types.MethodType(
usp_attn_turbowan_forward, block.self_attn)
if hasattr(self, 'vace_blocks'):
for block in self.vace_blocks:
block.self_attn.forward = types.MethodType(
usp_attn_turbowan_forward, block.self_attn)
videox_fun/pipeline/pipeline_wan.py
+34 -5
View File
@@ -1,5 +1,6 @@
import inspect
import math
import time
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -21,6 +22,17 @@ from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
def stochastic_sampling_timesteps(num_inference_steps, shift, device, num_timesteps=1000):
"""Official FlashHead timestep schedule with shift transform."""
if num_inference_steps == 4:
timesteps = [1000, 750, 500, 250]
else:
timesteps = np.linspace(num_timesteps, 1, num_inference_steps, dtype=np.float32).tolist()
timesteps = torch.tensor(timesteps + [0.0], dtype=torch.float32, device=device)
t = timesteps / num_timesteps
return shift * t / (1 + (shift - 1) * t) * num_timesteps
EXAMPLE_DOC_STRING = """
Examples:
```python
@@ -409,6 +421,7 @@ class WanPipeline(DiffusionPipeline):
max_sequence_length: int = 512,
comfyui_progressbar: bool = False,
shift: int = 5,
stochastic_sampling: bool = False,
) -> Union[WanPipelineOutput, Tuple]:
"""
Function invoked when calling the pipeline for generation.
@@ -471,7 +484,9 @@ class WanPipeline(DiffusionPipeline):
in_prompt_embeds = prompt_embeds
# 4. Prepare timesteps
if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
if stochastic_sampling:
timesteps = stochastic_sampling_timesteps(num_inference_steps, shift, device)
elif isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, mu=1)
elif isinstance(self.scheduler, FlowUniPCMultistepScheduler):
self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
@@ -513,8 +528,9 @@ class WanPipeline(DiffusionPipeline):
# 7. Denoising loop
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
self.transformer.num_inference_steps = num_inference_steps
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps
with self.progress_bar(total=len(denoise_timesteps)) as progress_bar:
for i, t in enumerate(denoise_timesteps):
self.transformer.current_steps = i
if self.interrupt:
@@ -528,6 +544,8 @@ class WanPipeline(DiffusionPipeline):
timestep = t.expand(latent_model_input.shape[0])
# predict noise model_output
torch.cuda.synchronize()
forward_start_time = time.time()
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=device):
noise_pred = self.transformer(
x=latent_model_input,
@@ -535,6 +553,9 @@ class WanPipeline(DiffusionPipeline):
t=timestep,
seq_len=seq_len,
)
torch.cuda.synchronize()
forward_time = time.time() - forward_start_time
print(f"Step {i}/{len(timesteps)} - Transformer forward time: {forward_time:.4f}s")
# perform guidance
if do_classifier_free_guidance:
@@ -542,7 +563,15 @@ class WanPipeline(DiffusionPipeline):
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
if stochastic_sampling:
t_i = (timesteps[i] / 1000).to(weight_dtype)
t_i_1 = (timesteps[i + 1] / 1000).to(weight_dtype)
x_0 = latents - noise_pred * t_i
latents = (1 - t_i_1) * x_0 + t_i_1 * torch.randn(
latents.shape, dtype=latents.dtype, device=device, generator=generator
)
else:
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
if callback_on_step_end is not None:
callback_kwargs = {}
@@ -554,7 +583,7 @@ class WanPipeline(DiffusionPipeline):
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
if i == len(denoise_timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
if comfyui_progressbar:
pbar.update(1)