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

Update self forcing gan and causal forcing precision

Browse Source
This commit is contained in:
bubbliiiing
2026-10-23 13:10:15 +08:00
parent 42b33f4998
commit e7afc604e1
30 changed files with 7265 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/fastvideo/predict_t2v_minimax_h3.py
+331
View File
@@ -0,0 +1,331 @@
import os
import sys
import torch
current_file_path = os.path.abspath(__file__)
project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.dist import set_multi_gpus_devices, shard_model
from videox_fun.models import (AutoencoderKLMiniMaxH3,
AutoencoderKLMiniMaxH3Audio,
FastVideoMiniMaxH3Transformer3DModel,
Qwen2TokenizerFast,
Qwen3VLForConditionalGeneration,
Qwen3VLProcessor)
from videox_fun.pipeline import MiniMaxH3Pipeline
from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
convert_model_weight_to_float8, merge_lora,
save_videos_with_audio_grid, unmerge_lora)
# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
# model_full_load means that the entire model will be moved to the GPU.
#
# model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
#
# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# model_group_offload transfers internal layer groups between CPU/CUDA,
# balancing memory efficiency and speed between full-module and leaf-level offloading methods.
#
# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use,
# resulting in slower speeds but saving a large amount of GPU memory.
# The MiniMax-H3 stack (~19B transformer + the ~62GB Qwen3-VL conditioner) does not fit a single
# 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).
ulysses_degree = 1
ring_degree = 1
# Use FSDP to save more GPU memory in multi gpus. Unused at ulysses_degree = ring_degree = 1.
fsdp_dit = False
fsdp_text_encoder = False
# Compile will give a speedup in fixed resolution and need a little GPU memory.
# The compile_dit is not compatible with sequential_cpu_offload.
compile_dit = False
# model path
model_name = "models/Diffusion_Transformer/MiniMax-H3"
# ---- FastVideo VSA sparse-student config --------------------------------------------------
# Official FastH3 8-step DMD student (FastVideo/FastVideo-FastH3-8-Step-V2): a MiniMax-H3 backbone
# distilled with data-free DMD2 + VSA-H3 at 80% sparsity, carrying a trained `to_gate_compress` gate
# per block. `model_name` still supplies the pipeline components (text encoder / VAEs / tokenizer /
# schedulers) from the dense base; the transformer comes from `transformer_path` below.
# student_topk is the keep-ratio of key blocks per query (0.20 -> 80% sparsity, the official
# vsa_sparsity=0.8); set it to 1.0 to fall back to dense attention.
student_topk = 0.20
# Transformer weights. Accepts either the official FastH3 diffusers transformer folder (`.../transformer`,
# sharded; its gate keys `attn.to_gate_compress.*` are remapped to this repo's `attn.attn_op.to_gate_compress.*`
# on load), or a single-file native-layout state_dict from this repo's own training (loaded AFTER
# set_student_topk so the trained gate overwrites the zero init). Set to None to run the sparse student on
# the dense base weights (zero-init gate).
transformer_path = "models/Diffusion_Transformer/FastVideo-FastH3-8-Step-V2/transformer"
vae_path = None
lora_path = None
pdd_lora_path = None
# Other params
# MiniMax-H3 generates at a fixed 24 fps, only accepts multiples of 32 as height / width, and snaps video_length up
# to the next 17 * n + 5 the video VAE can decode (the duration has to stay between 5 and 15 seconds).
# Leave sample_size as None to use MiniMax-H3's own 16:9 canvas (768x1344).
sample_size = [704, 1280]
video_length = 124
fps = 24
# Use torch.float16 if GPU does not support torch.bfloat16
# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
weight_dtype = torch.bfloat16
prompt = "A red fox trotting through a snowy pine forest, snow crunching underfoot"
seed = 43
# Number of denoising steps, i.e. of model evaluations. The official FastH3-8-Step-V2 runs eight
# transformer forwards (fastvideo_inference.json: transformer_forwards=8).
num_inference_steps = 8
# The released checkpoint is guidance-distilled: leave guidance_scale at 1 to run one forward pass per step
# with no CFG. A value above 1 enables classifier-free guidance with a negative_prompt, running two passes.
guidance_scale = 1
# The exponential sigma shifts of the two schedules. The official 8-step DMD grid
# (dmd_denoising_steps [999,874,...,125]) is *uniform*; here a uniform sigma grid is produced by
# shift=1.0 (set_shift warps the grid, so the checkpoint's video shift 10 / audio 3 do not map onto
# flow_shift directly). flow_shift=1.0 reproduces the trained video schedule; audio keeps shift 3.
flow_shift = 1.0
audio_flow_shift = 3.0
lora_weight = 0.55
save_path = "samples/fasth3-8step-videos-t2v"
device = set_multi_gpus_devices(ulysses_degree, ring_degree)
# Transformer: build the VSA sparse student, then REBUILD the VSA kernels at the requested keep-ratio.
# set_student_topk recreates attn_op (resetting to_gate_compress to its zero init), so it must run
# BEFORE any trained gate is loaded. `transformer_path` may be the official FastH3 diffusers transformer
# folder (sharded) or a single-file native-layout state_dict; when it is a folder the backbone comes from
# it directly, otherwise from the dense `model_name` base (whose transformer lives in `transformer/`).
is_transformer_dir = transformer_path is not None and os.path.isdir(transformer_path)
transformer = FastVideoMiniMaxH3Transformer3DModel.from_pretrained(
transformer_path if is_transformer_dir else model_name,
subfolder="" if is_transformer_dir else "transformer",
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
)
transformer.set_student_topk(student_topk)
# The official FastH3 checkpoint is sharded and names its gate `attn.to_gate_compress.*`; this repo
# nests it under `attn.attn_op.to_gate_compress.*`. from_pretrained above loaded the (distilled)
# backbone and dropped the gate keys as unexpected, so load just the gate here, remapped, and assert
# every block got one (a silently zero gate would leave the compressed branch untrained -> artifacts).
if is_transformer_dir and student_topk < 1.0:
import json
from safetensors import safe_open
tdir = transformer_path
weight_map = json.load(open(os.path.join(tdir, "diffusion_pytorch_model.safetensors.index.json")))["weight_map"]
shard_to_gate_keys = {}
for key, shard in weight_map.items():
if "to_gate_compress" in key:
shard_to_gate_keys.setdefault(shard, []).append(key)
gate_sd = {}
for shard, keys in shard_to_gate_keys.items():
with safe_open(os.path.join(tdir, shard), framework="pt") as f:
for key in keys:
gate_sd[key.replace("attn.to_gate_compress", "attn.attn_op.to_gate_compress")] = f.get_tensor(key)
m, u = transformer.load_state_dict(gate_sd, strict=False)
gate_loaded = sum(1 for k in gate_sd if k not in m)
print(f"FastH3 gate: {gate_loaded}/{len(gate_sd)} loaded, unexpected {len(u)}", flush=True)
assert len(u) == 0 and gate_loaded == len(gate_sd), (
f"FastH3 gate remap incomplete: loaded {gate_loaded}/{len(gate_sd)}, unexpected {len(u)}.")
# Single-file native-layout state_dict (gate already under attn_op): overwrite the base backbone wholesale.
if transformer_path is not None and not is_transformer_dir:
print(f"From checkpoint: {transformer_path}")
if transformer_path.endswith("safetensors"):
from safetensors.torch import load_file
state_dict = load_file(transformer_path)
else:
state_dict = torch.load(transformer_path, map_location="cpu")
state_dict = state_dict["state_dict"] if isinstance(state_dict, dict) and "state_dict" in state_dict else state_dict
m, u = transformer.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
attn_op_missing = [k for k in m if "attn_op." in k]
if attn_op_missing and student_topk < 1.0:
# A trained VSA student should provide every attn_op param; missing ones mean the compressed
# branch would run on freshly-init (zero) gate values. Fail loudly instead of silently.
raise RuntimeError(
f"{len(attn_op_missing)} VSA attn_op keys are missing from the checkpoint "
f"(student_topk={student_topk}); e.g. {attn_op_missing[:4]}. Check the checkpoint / topk.")
# `strict=False` accepts a file whose keys belong to another model — a LoRA checkpoint, say — by
# loading nothing at all, so an unexpected key is a hard error.
assert len(u) == 0, (
f"{transformer_path} holds {len(u)} key(s) the transformer does not have, e.g. {list(u)[:3]}. A LoRA "
"checkpoint belongs in `lora_path`, not `transformer_path`."
)
pdd_config = None
if pdd_lora_path is not None:
if lora_path is not None:
raise ValueError("`lora_path` and `pdd_lora_path` cannot be used together.")
from videox_fun.models.minimax_h3_pdd import (load_pdd_lora,
pdd_num_inference_steps,
pdd_step_callback)
pdd_config = load_pdd_lora(transformer, pdd_lora_path)
num_inference_steps = pdd_num_inference_steps(pdd_config, num_inference_steps, teacher_default=40)
# Video VAE. The released weights are float32 and the decode runs under float16 autocast, so the VAE is not
# downcast even when the rest of the pipeline is bfloat16 (this is also how the training scripts load it).
vae = AutoencoderKLMiniMaxH3.from_pretrained(
model_name,
subfolder="vae",
low_cpu_mem_usage=True,
)
if vae_path is not None:
print(f"From checkpoint: {vae_path}")
if vae_path.endswith("safetensors"):
from safetensors.torch import load_file
state_dict = load_file(vae_path)
else:
state_dict = torch.load(vae_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = vae.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
# Audio VAE, waveform in / waveform out: MiniMax-H3 has no separate vocoder. Float32 as released, like the video VAE.
audio_vae = AutoencoderKLMiniMaxH3Audio.from_pretrained(
model_name,
subfolder="audio_vae",
low_cpu_mem_usage=True,
)
# Get Tokenizer and Processor
tokenizer = Qwen2TokenizerFast.from_pretrained(os.path.join(model_name, "tokenizer"))
processor = Qwen3VLProcessor.from_pretrained(os.path.join(model_name, "processor"))
# Get Text encoder. MiniMax-H3 reads the unnormalized hidden state after the 50th decoder layer of Qwen3-VL.
text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
os.path.join(model_name, "text_encoder"),
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
)
text_encoder = text_encoder.eval()
# Get Schedulers. MiniMax-H3 steps the video and the audio latents down two schedules inside one transformer call.
scheduler = MiniMaxH3Scheduler.from_pretrained(model_name, subfolder="scheduler")
audio_scheduler = MiniMaxH3Scheduler.from_pretrained(model_name, subfolder="audio_scheduler")
pipeline = MiniMaxH3Pipeline(
vae=vae,
audio_vae=audio_vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
processor=processor,
transformer=transformer,
scheduler=scheduler,
audio_scheduler=audio_scheduler,
)
# The float32 modules of the mixed-precision checkpoint stay untouched by the float8 quantization.
fp8_exclude_module_name = [
"proj_in", "audio_proj_in", "context_embedder", "time_embedder", "time_proj",
"token_refiner", "norm_out", "proj_out", "audio_proj_out",
]
use_qfloat8 = "qfloat8" in GPU_memory_mode
if use_qfloat8:
convert_model_weight_to_float8(transformer, exclude_module_name=fp8_exclude_module_name, device=device)
# The pure *cpu_offload modes (qfloat8 variants excluded, unsupported under FSDP). Under sequence parallelism the
# sharded modules offload through FSDP's own CPUOffload, enabled here at wrap time; the pipeline then routes the
# whole-model offload call to the non-FSDP VAE / audio VAE automatically, so no per-mode branch is needed below.
wants_offload = GPU_memory_mode in ("model_cpu_offload", "model_group_offload", "sequential_cpu_offload")
if ulysses_degree > 1 or ring_degree > 1:
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()
if any(p.dtype == torch.float32 for p in m.parameters(recurse=False))]
shard_fn = partial(shard_model, device_id=device, param_dtype=None, cast_dtype=False,
module_to_wrapper=list(transformer.transformer_blocks),
ignored_modules=fp32_modules, offload_to_cpu=wants_offload)
pipeline.transformer = shard_fn(pipeline.transformer)
print("Add FSDP DIT")
if fsdp_text_encoder:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype,
module_to_wrapper=list(text_encoder.model.language_model.layers),
offload_to_cpu=wants_offload)
pipeline.text_encoder.model = shard_fn(pipeline.text_encoder.model)
print("Add FSDP TEXT ENCODER")
if compile_dit:
for i in range(len(pipeline.transformer.transformer_blocks)):
pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
print("Add Compile")
# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
# The order lives inside the helper: quantization has to happen before the offload hooks
# are installed.
# The FP8 conversion above has already run (before the FSDP sharding, on purpose); only the dequant
# wrapper and the memory placement are left, which is what the preconverted tag installs.
apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype,
quant_tag="qfloat8_preconverted" if GPU_memory_mode.endswith("_and_qfloat8") else None,
exclude_module_name=[])
generator = torch.Generator(device=device).manual_seed(seed)
if lora_path is not None:
pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
pdd_callback = None if pdd_config is None else pdd_step_callback(
transformer, scheduler, audio_scheduler, pdd_config, num_inference_steps
)
with torch.no_grad():
output = pipeline(
prompt=prompt,
height=None if sample_size is None else sample_size[0],
width=None if sample_size is None else sample_size[1],
num_frames=video_length,
num_inference_steps=num_inference_steps,
flow_shift=flow_shift,
audio_flow_shift=audio_flow_shift,
guidance_scale=guidance_scale,
generator=generator,
output_type="pt",
callback_on_step_end=pdd_callback,
)
print(f"[{os.environ.get('RANK', '0')}] generation done, decoding", flush=True)
if lora_path is not None:
pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
sample = output.videos
audio = output.audio
audio_sample_rate = output.sampling_rate
def save_results():
if not os.path.exists(save_path):
os.makedirs(save_path, exist_ok=True)
index = len([path for path in os.listdir(save_path)]) + 1
prefix = str(index).zfill(8)
video_path = os.path.join(save_path, prefix + ".mp4")
save_videos_with_audio_grid(sample, audio, video_path, fps=fps, audio_sample_rate=audio_sample_rate)
if ulysses_degree * ring_degree > 1:
import torch.distributed as dist
if dist.get_rank() == 0:
save_results()
# Keep every rank alive until the saving rank finishes; an early exit of one rank makes the elastic launcher
# terminate the others.
dist.barrier()
else:
save_results()
examples/fastvideo/predict_t2v_wan2.1.py
+313
View File
@@ -0,0 +1,313 @@
import os
import re
import sys
import numpy as np
import torch
from diffusers import FlowMatchEulerDiscreteScheduler
from omegaconf import OmegaConf
from PIL import Image
current_file_path = os.path.abspath(__file__)
project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.dist import set_multi_gpus_devices, shard_model
from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
FastVideoWanTransformer3DModel, WanT5EncoderModel)
from videox_fun.pipeline import WanPipeline
from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
FlowUniPCMultistepScheduler,
apply_gpu_memory_mode, filter_kwargs, merge_lora,
save_videos_grid, unmerge_lora)
# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
# model_full_load means that the entire model will be moved to the GPU.
#
# model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
#
# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# model_group_offload transfers internal layer groups between CPU/CUDA,
# balancing memory efficiency and speed between full-module and leaf-level offloading methods.
#
# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use,
# resulting in slower speeds but saving a large amount of GPU memory.
GPU_memory_mode = "model_cpu_offload"
# Multi GPUs config
# Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used.
# For example, if you are using 8 GPUs, you can set ulysses_degree = 2 and ring_degree = 4.
# If you are using 1 GPU, you can set ulysses_degree = 1 and ring_degree = 1.
# NOTE: the VSA sparse kernel has no sequence-parallel path, so a FastWan student runs single-GPU
# only; keep ulysses_degree = ring_degree = 1 (enable_multi_gpus_inference() raises otherwise).
ulysses_degree = 1
ring_degree = 1
# Use FSDP to save more GPU memory in multi gpus.
fsdp_dit = False
fsdp_text_encoder = True
# Compile will give a speedup in fixed resolution and need a little GPU memory.
# The compile_dit is not compatible with the fsdp_dit and sequential_cpu_offload.
compile_dit = False
# Riflex config
enable_riflex = False
# Index of intrinsic frequency
riflex_k = 6
# Config and model path
config_path = "config/wan2.1/wan_civitai.yaml"
# Base model (Wan-AI/Wan2.1-T2V-1.3B): provides config.json + the dense backbone + VAE + T5.
model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
# ---- FastVideo VSA sparse-student config --------------------------------------------------
# Official released checkpoint: FastVideo/FastWan2.1-T2V-1.3B-Diffusers (VSA sparse attention +
# DMD few-step distillation). student_topk is the keep-ratio of key blocks per query and MUST
# equal the value the checkpoint was trained at: topk=0.20 -> 80% sparsity (--VSA-sparsity 0.8).
student_topk = 0.20
# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++".
# FastWan ships UniPCMultistepScheduler(flow_shift=3.0); "Flow_Unipc" + shift=3.0 mirrors it and
# keeps every denoising step at a high noise level (DMD-style), unlike "Flow" whose last step ~0.
sampler_name = "Flow_Unipc"
# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics. 3.0 is the official
# flow_shift and the recommended value for 480p.
shift = 3.0
# Official released checkpoint (diffusers layout: transformer/diffusion_pytorch_model.safetensors,
# diffusers key naming + the VSA to_gate_compress gate). Download:
# HF_ENDPOINT=https://hf-mirror.com huggingface-cli download \
# FastVideo/FastWan2.1-T2V-1.3B-Diffusers --local-dir <dir>
transformer_path = "models/Diffusion_Transformer/FastWan2.1-T2V-1.3B-Diffusers/transformer/diffusion_pytorch_model.safetensors"
vae_path = None
lora_path = None
# Other params
sample_size = [480, 848]
video_length = 81
fps = 16
# Use torch.float16 if GPU does not support torch.bfloat16
# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
weight_dtype = torch.bfloat16
prompt = "A cinematic aerial shot of a snow-covered mountain village at dawn, warm light spilling from wooden windows, smoke curling from the chimneys while a river of mist drifts through the valley below."
negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
guidance_scale = 1.0 # the DMD student is CFG-free
seed = 1024
num_inference_steps = 3
lora_weight = 0.55
save_path = "samples/fastvideo-wan2.1-official"
def remap_fastwan_diffusers_to_vxf(state_dict):
"""Official FastWan diffusers transformer -> VideoX-Fun civitai WanTransformer3DModel layout.
FastWan ships diffusers naming (attn1/attn2, ffn.net, condition_embedder, proj_out,
scale_shift_table) plus the VSA gate ``to_gate_compress``; VideoX-Fun keeps the native Wan
civitai naming (self_attn/cross_attn, ffn.0/2, time_/text_embedding, head.head, modulation)
and nests the gate under ``self_attn.attn_op``. Only the cross-attn norm is affine, so the
single diffusers ``norm2`` maps to civitai ``norm3``; ``patch_embedding`` is already 5D and
passes through. Turns the 885 released keys into an exact 1:1 match (825 backbone + 30x2 gate).
"""
TOP = {
"condition_embedder.time_embedder.linear_1": "time_embedding.0",
"condition_embedder.time_embedder.linear_2": "time_embedding.2",
"condition_embedder.time_proj": "time_projection.1",
"condition_embedder.text_embedder.linear_1": "text_embedding.0",
"condition_embedder.text_embedder.linear_2": "text_embedding.2",
"proj_out": "head.head",
"scale_shift_table": "head.modulation",
}
out = {}
for k, v in state_dict.items():
m = re.match(r"blocks\.(\d+)\.(.+)$", k)
if m:
blk, rest = m.group(1), m.group(2)
rest = rest.replace("attn1.", "self_attn.").replace("attn2.", "cross_attn.")
rest = (rest.replace(".to_out.0", ".o").replace(".to_q", ".q")
.replace(".to_k", ".k").replace(".to_v", ".v"))
rest = rest.replace("ffn.net.0.proj", "ffn.0").replace("ffn.net.2", "ffn.2")
if rest.startswith("norm2."):
rest = "norm3." + rest[len("norm2."):]
if rest == "scale_shift_table":
rest = "modulation"
if rest.startswith("to_gate_compress"):
rest = "self_attn.attn_op." + rest
nk = "blocks." + blk + "." + rest
else:
nk = k
for src, dst in TOP.items():
if k == src or k.startswith(src + "."):
nk = dst + k[len(src):]
break
out[nk] = v
return out
device = set_multi_gpus_devices(ulysses_degree, ring_degree)
config = OmegaConf.load(config_path)
# Build the sparse student from the base Wan2.1 config, then REBUILD the VSA kernels at the
# released keep-ratio BEFORE loading weights: set_student_topk recreates attn_op (resetting
# to_gate_compress to its zero init), so the trained checkpoint values must be loaded after.
transformer = FastVideoWanTransformer3DModel.from_pretrained(
os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
)
transformer.set_student_topk(student_topk)
if transformer_path is not None:
print(f"From checkpoint: {transformer_path}")
if transformer_path.endswith("safetensors"):
from safetensors.torch import load_file
state_dict = load_file(transformer_path)
else:
state_dict = torch.load(transformer_path, map_location="cpu")
state_dict = state_dict["state_dict"] if isinstance(state_dict, dict) and "state_dict" in state_dict else state_dict
# Official diffusers release -> VideoX-Fun civitai key layout (route the VSA gate under attn_op).
state_dict = remap_fastwan_diffusers_to_vxf(state_dict)
m, u = transformer.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
attn_op_missing = [k for k in m if "attn_op." in k]
if attn_op_missing:
# A VSA student MUST provide every attn_op param; missing ones mean the compressed branch
# would run on freshly-init (zero) gate values. Fail loudly instead of silently.
raise RuntimeError(
f"{len(attn_op_missing)} VSA attn_op keys are missing from the checkpoint "
f"(student_topk={student_topk}); e.g. {attn_op_missing[:4]}. Check the remap / topk.")
if u:
print(f"WARNING: {len(u)} unexpected keys ignored; e.g. {list(u)[:4]}")
# For the official FastWan2.1-T2V-1.3B checkpoint this is expected to be 0 missing / 0 unexpected.
# Get Vae
vae = AutoencoderKLWan.from_pretrained(
os.path.join(model_name, config['vae_kwargs'].get('vae_subpath', 'vae')),
additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
).to(weight_dtype)
if vae_path is not None:
print(f"From checkpoint: {vae_path}")
if vae_path.endswith("safetensors"):
from safetensors.torch import load_file
state_dict = load_file(vae_path)
else:
state_dict = torch.load(vae_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = vae.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
# Get Tokenizer
tokenizer = AutoTokenizer.from_pretrained(
os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
)
# Get Text encoder
text_encoder = WanT5EncoderModel.from_pretrained(
os.path.join(model_name, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
)
# Get Scheduler
Chosen_Scheduler = {
"Flow": FlowMatchEulerDiscreteScheduler,
"Flow_Unipc": FlowUniPCMultistepScheduler,
"Flow_DPM++": FlowDPMSolverMultistepScheduler,
}[sampler_name]
if sampler_name == "Flow_Unipc" or sampler_name == "Flow_DPM++":
config['scheduler_kwargs']['shift'] = 1
scheduler = Chosen_Scheduler(
**filter_kwargs(Chosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
# Get Pipeline
pipeline = WanPipeline(
transformer=transformer,
vae=vae,
tokenizer=tokenizer,
text_encoder=text_encoder,
scheduler=scheduler,
)
if ulysses_degree > 1 or ring_degree > 1:
from functools import partial
transformer.enable_multi_gpus_inference()
if fsdp_dit:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
pipeline.transformer = shard_fn(pipeline.transformer)
print("Add FSDP DIT")
if fsdp_text_encoder:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
pipeline.text_encoder = shard_fn(pipeline.text_encoder)
print("Add FSDP TEXT ENCODER")
if compile_dit:
for i in range(len(pipeline.transformer.blocks)):
pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
print("Add Compile")
# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
generator = torch.Generator(device=device).manual_seed(seed)
if lora_path is not None:
pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
with torch.no_grad():
video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1
if enable_riflex:
pipeline.transformer.enable_riflex(k = riflex_k, L_test = latent_frames)
sample = pipeline(
prompt,
num_frames = video_length,
negative_prompt = negative_prompt,
height = sample_size[0],
width = sample_size[1],
generator = generator,
guidance_scale = guidance_scale,
num_inference_steps = num_inference_steps,
shift = shift,
).videos
if lora_path is not None:
pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
def save_results():
if not os.path.exists(save_path):
os.makedirs(save_path, exist_ok=True)
index = len([path for path in os.listdir(save_path)]) + 1
prefix = str(index).zfill(8)
if video_length == 1:
video_path = os.path.join(save_path, prefix + ".png")
image = sample[0, :, 0]
image = image.transpose(0, 1).transpose(1, 2)
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)
print(f"Saved video to: {video_path}")
if ulysses_degree * ring_degree > 1:
import torch.distributed as dist
if dist.get_rank() == 0:
save_results()
else:
save_results()
scripts/fastvideo/README_TRAIN_DISTILL.md
+648
View File
@@ -0,0 +1,648 @@
# FastVideo VSA Distill Training Guide
This document provides a complete workflow for distilling Wan2.1 into FastWan2.1 with VSA (Video Sparse Attention), including environment configuration, data preparation, distributed training, and inference testing.
> **Note**: FastVideo VSA combines a block-sparse video attention (a learned gate keeps ~20% of key blocks per query) with DMD few-step distillation, so the student runs CFG-free in ~3 steps while attending sparsely. The training recipe, data pipeline, and distributed backends are identical to the TurboDiffusion distill guide; only the sparse student model class differs.
---
## Table of Contents
- [1. Environment Configuration](#1-environment-configuration)
- [2. Data Preparation](#2-data-preparation)
- [2.1 Quick Test Dataset](#21-quick-test-dataset)
- [2.2 Dataset Structure](#22-dataset-structure)
- [2.3 metadata.json Format](#23-metadatajson-format)
- [2.4 Relative vs Absolute Path Usage](#24-relative-vs-absolute-path-usage)
- [3. Distill Training](#3-distill-training)
- [3.1 Download Pretrained Model](#31-download-pretrained-model)
- [3.2 Quick Start (DeepSpeed-Zero-2)](#32-quick-start-deepspeed-zero-2)
- [3.3 Common Training Parameters](#33-common-training-parameters)
- [3.4 Training Validation](#34-training-validation)
- [3.5 Training with FSDP](#35-training-with-fsdp)
- [3.6 Other Backends](#36-other-backends)
- [3.7 Multi-Machine Distributed Training](#37-multi-machine-distributed-training)
- [4. Inference Testing](#4-inference-testing)
- [4.1 Inference Parameters](#41-inference-parameters)
- [4.2 Text-to-Video (T2V) Inference](#42-text-to-video-t2v-inference)
- [4.3 Single-GPU Only](#43-single-gpu-only)
- [5. Additional Resources](#5-additional-resources)
---
## 1. Environment Configuration
**Method 1: Using requirements.txt**
```bash
pip install -r requirements.txt
```
**Method 2: Manual Dependency Installation**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 💡 The VSA block-sparse attention runs on a Triton kernel; Triton ships with PyTorch (>=2.1) on Linux, so no extra install is needed.
**Method 3: Using Docker**
When using Docker, please ensure that the GPU driver and CUDA environment are correctly installed on your machine, then execute the following commands:
```
# pull image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# enter image
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 2. Data Preparation
### 2.1 Quick Test Dataset
We provide a test dataset containing several training samples.
```bash
# Download official demo dataset
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 Dataset Structure
```
📦 datasets/
├── 📂 my_dataset/
│ ├── 📂 train/
│ │ ├── 📄 video001.mp4
│ │ ├── 📄 video002.mp4
│ │ └── 📄 ...
│ └── 📄 metadata.json
```
### 2.3 metadata.json Format
**Relative Path Format** (example):
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
},
{
"file_path": "train/video002.mp4",
"text": "A person walking through a forest, cinematic view",
"type": "video",
"width": 1328,
"height": 1328
}
]
```
**Absolute Path Format**:
```json
[
{
"file_path": "/mnt/data/videos/sunset.mp4",
"text": "A beautiful sunset over the ocean",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
**Key Field Descriptions**:
- `file_path`: Video path (relative or absolute path)
- `text`: Video description (English prompt)
- `type`: Data type, fixed as `"video"`
- `width` / `height`: Video dimensions (**recommended** to provide for bucket training. If not provided, they will be automatically read during training, which may slow down training when data is stored on slow systems like OSS).
- You can use `scripts/process_json_add_width_and_height.py` to extract width and height fields for JSON files without them, supporting both images and videos.
- Usage: `python scripts/process_json_add_width_and_height.py --input_file datasets/X-Fun-Videos-Demo/metadata.json --output_file datasets/X-Fun-Videos-Demo/metadata_add_width_height.json`.
### 2.4 Relative vs Absolute Path Usage
**Relative Path**:
If your data uses relative paths, configure in the training script:
```bash
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
```
**Absolute Path**:
If your data uses absolute paths, configure in the training script:
```bash
export DATASET_NAME=""
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
```
> 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths.
---
## 3. Distill Training
The distillation is a DMD (distribution matching distillation) run with three networks: a sparse **generator** (the FastWan student, `FastVideoWanTransformer3DModel`), a sparse **fake_score** critic (same class), and a frozen **dense** `real_score` teacher (`WanTransformer3DModel`). The student and critic carry the VSA block-sparse self-attention; the teacher keeps full attention.
### 3.1 Download Pretrained Model
```bash
# Create model directory
mkdir -p models/Diffusion_Transformer
# Download Wan2.1 official weights
modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
```
### 3.2 Quick Start (DeepSpeed-Zero-2)
After downloading the data as per **2.1 Quick Test Dataset** and the weights as per **3.1 Download Pretrained Model**, you can directly copy and run the quick start command.
DeepSpeed-Zero-2 and FSDP are recommended for training. Here we use DeepSpeed-Zero-2 as an example.
The difference between DeepSpeed-Zero-2 and FSDP is whether the model weights are sharded. **If you experience insufficient GPU memory when using multiple GPUs with DeepSpeed-Zero-2**, you can switch to FSDP for training.
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.3 Common Training Parameters
**Key Parameter Descriptions**:
| Parameter | Description | Example Value |
|-----|------|-------|
| `--config_path` | Config file path | `config/wan2.1/wan_civitai.yaml` |
| `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` |
| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` |
| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` |
| `--train_batch_size` | Batch size per GPU | 1 |
| `--image_sample_size` | Maximum image training resolution | 640 |
| `--video_sample_size` | Maximum video training resolution | 640 |
| `--token_sample_size` | Token sample size | 640 |
| `--video_sample_stride` | Video sampling stride | 2 |
| `--video_sample_n_frames` | Video sampling frames | 81 |
| `--gradient_accumulation_steps` | Gradient accumulation steps (effectively increases batch) | 1 |
| `--dataloader_num_workers` | DataLoader worker processes | 8 |
| `--num_train_epochs` | Number of training epochs | 100 |
| `--checkpointing_steps` | Save checkpoint every N steps | 50 |
| `--learning_rate` | Initial learning rate (generator) | 2e-06 |
| `--learning_rate_critic` | Initial learning rate (critic) | 2e-06 |
| `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` |
| `--lr_warmup_steps` | Learning rate warmup steps | 100 |
| `--seed` | Random seed | 42 |
| `--output_dir` | Output directory | `output_dir_fastvideo_wan2.1_1.3b` |
| `--gradient_checkpointing` | Enable gradient checkpointing | - |
| `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
| `--adam_weight_decay` | AdamW weight decay | 3e-2 |
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--vae_mini_batch` | VAE encoding mini-batch size | 1 |
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
| `--enable_bucket` | Enable bucket training, no cropping, group by resolution | - |
| `--random_hw_adapt` | Auto-scale images/videos to random size in `[min_size, max_size]` range | - |
| `--training_with_video_token_length` | Train based on token length, supports arbitrary resolutions | - |
| `--uniform_sampling` | Uniform timestep sampling | - |
| `--low_vram` | Low VRAM mode | - |
| `--train_mode` | Training mode: `normal` (standard) or `i2v` (image-to-video) | `normal` |
| `--resume_from_checkpoint` | Resume training path, use `"latest"` to auto-select latest checkpoint | None |
| `--validation_steps` | Run validation every N steps | 2000 |
| `--validation_epochs` | Run validation every N epochs | 5 |
| `--validation_prompts` | Prompts for video generation validation | `"A dog shaking head..."` |
| `--trainable_modules` | Trainable modules (`"."` means all modules) | `"."` |
**Distill-Specific Parameters**:
| Parameter | Description | Example Value |
|-----|------|-------|
| `--denoising_step_indices_list` | Denoising step list (core distill parameter) | `1000 750 500 250` |
| `--real_guidance_scale` | Real guidance scale for scoring | 6.0 |
| `--fake_guidance_scale` | Fake guidance scale for scoring | 0.0 |
| `--gen_update_interval` | Generator update interval | 5 |
| `--negative_prompt` | Negative prompt for distillation | Chinese negative prompt |
| `--validation_paths` | Validation image paths for I2V mode | Image path list |
| `--train_sampling_steps` | Training sampling steps | 1000 |
**VSA Sparsity**:
- The student/critic self-attention is VSA block-sparse. The keep-ratio `student_topk` defaults to `0.20` (i.e. 80% sparsity) inside `FastVideoWanTransformer3DModel`; it is a model-class default, not a CLI flag.
- The learned compression gate `to_gate_compress` is created zero-initialized and becomes trainable automatically under `--trainable_modules "."` (its parameter names contain a dot, like every other module).
- The `real_score` teacher stays dense (`WanTransformer3DModel`) and frozen, so the sparsity is distilled into the student rather than copied from the teacher.
- Keep the inference-time `student_topk` equal to the value the checkpoint was trained at; changing it rebuilds the attention op and resets the learned gate.
**Sample Size Configuration Guide**:
- `video_sample_size` represents the resolution size of videos; when `random_hw_adapt` is True, it represents the minimum value between video and image resolutions.
- `image_sample_size` represents the resolution size of images; when `random_hw_adapt` is True, it represents the maximum value between video and image resolutions.
- `token_sample_size` represents the resolution corresponding to the maximum token length when `training_with_video_token_length` is True.
- Due to potential confusion in configuration, **if you don't require arbitrary resolution for finetuning**, it is recommended to set `video_sample_size`, `image_sample_size`, and `token_sample_size` to the same fixed value, such as **(320, 480, 512, 640, 960)**.
- **All set to 320** represents **240P**.
- **All set to 480** represents **320P**.
- **All set to 640** represents **480P**.
- **All set to 960** represents **720P**.
**Other Parameter Explanations**:
- `enable_bucket` is used to enable bucket training. When enabled, the model does not crop the images and videos at the center, but instead, it trains the entire images and videos after grouping them into buckets based on resolution.
- `random_hw_adapt` is used to enable automatic height and width scaling for images and videos. When `random_hw_adapt` is enabled, the training images will have their height and width set to `image_sample_size` as the maximum and `min(video_sample_size, 512)` as the minimum.
- `training_with_video_token_length` specifies training the model according to token length. For training images and videos, the height and width will be set to `image_sample_size` as the maximum and `video_sample_size` as the minimum.
- `train_mode` is used to specify the training mode, which can be either normal or i2v. If you only wish to achieve text-to-video generation, keep `normal`.
- `resume_from_checkpoint` is used to set the training should be resumed from a previous checkpoint. Use a path or `"latest"` to automatically select the last available checkpoint.
### 3.4 Training Validation
You can configure validation parameters to periodically generate test videos during training, allowing you to monitor training progress and model quality.
**Validation Parameter Descriptions**:
| Parameter | Description | Recommended Value |
|------|------|--------|
| `--validation_steps` | Run validation every N steps | 2000 |
| `--validation_epochs` | Run validation every N epochs | 5 |
| `--validation_prompts` | Prompts for video generation validation | English prompts |
| `--validation_paths` | Validation image paths for I2V mode (i2v mode only) | `"asset/1.png"` |
**Normal Mode Example** (T2V Validation):
```bash
--validation_steps=2000 \
--validation_epochs=5 \
--validation_prompts="A dog shaking head. The video is of high quality, and the view is very clear."
```
**Notes**:
- Validation videos will be saved to the `output_dir` directory.
- Multiple prompts validation format: `--validation_prompts "prompt1" "prompt2" "prompt3"`.
- Distill model validation uses the few steps defined in `denoising_step_indices_list` and runs CFG-free (`guidance_scale=1.0`), matching inference.
### 3.5 Training with FSDP
**If you experience insufficient GPU memory when using multiple GPUs with DeepSpeed-Zero-2**, you can switch to FSDP for training. `scripts/fastvideo/train_distill_wan2.1.sh` is a ready-to-run FSDP launcher (it wraps `WanAttentionBlock` with `FULL_SHARD`).
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.6 Other Backends
#### 3.6.1 Training with DeepSpeed-Zero-3
DeepSpeed Zero-3 is not highly recommended at the moment. In this repository, using FSDP has fewer errors and is more stable.
DeepSpeed Zero-3 is suitable for 14B Wan at high resolutions. After training, you can use the following command to get the final model:
```bash
python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/checkpoint-{your-num-steps}-outputs --max_shard_size 80GB --safe_serialization
```
Training shell command is as follows:
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
#### 3.6.2 Training Without DeepSpeed and FSDP
**This approach is not recommended because without memory-saving backends, it easily causes insufficient GPU memory**. This is only provided as a reference shell for training.
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.7 Multi-Machine Distributed Training
**Suitable for**: Ultra-large datasets, faster training speed
#### 3.7.1 Environment Configuration
Assuming 2 machines, each with 8 GPUs:
**Machine 0 (Master)**:
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export MASTER_ADDR="192.168.1.100" # Master machine IP
export MASTER_PORT=10086
export WORLD_SIZE=2 # Total number of machines
export NUM_PROCESS=16 # Total processes = machines × 8
export RANK=0 # Current machine rank (0 or 1)
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
**Machine 1 (Worker)**:
```bash
export MASTER_ADDR="192.168.1.100" # Same as Master
export MASTER_PORT=10086
export WORLD_SIZE=2
export NUM_PROCESS=16
export RANK=1 # Note this is 1
NCCL_DEBUG=INFO
# Use the same accelerate launch command as Machine 0
```
#### 3.7.2 Multi-Machine Training Notes
- **Network Requirements**:
- RDMA/InfiniBand recommended (high performance)
- Without RDMA, add environment variables:
```bash
export NCCL_IB_DISABLE=1
export NCCL_P2P_DISABLE=1
```
- **Data Synchronization**: All machines must be able to access the same data paths (NFS/shared storage)
---
## 4. Inference Testing
### 4.1 Inference Parameters
**Key Parameter Descriptions** (from `examples/fastvideo/predict_t2v_wan2.1.py`):
| Parameter | Description | Example Value |
|------|------|-------|
| `GPU_memory_mode` | GPU memory mode, see table below for options | `model_cpu_offload` |
| `ulysses_degree` | Ulysses parallelism degree (VSA is single-GPU only, keep 1) | 1 |
| `ring_degree` | Ring parallelism degree (VSA is single-GPU only, keep 1) | 1 |
| `config_path` | Config file path | `config/wan2.1/wan_civitai.yaml` |
| `model_name` | Base model path (config + dense backbone + VAE + T5) | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
| `student_topk` | VSA keep-ratio; must equal the trained value | `0.20` |
| `sampler_name` | Sampler type: `Flow`, `Flow_Unipc`, `Flow_DPM++` | `Flow_Unipc` |
| `shift` | Noise schedule shift (official FastWan flow_shift) | `3.0` |
| `transformer_path` | Path to the FastWan student Transformer weights | trained checkpoint / official FastWan |
| `vae_path` | Path to trained VAE weights | `None` |
| `lora_path` | LoRA weights path | `None` |
| `sample_size` | Generated video resolution `[height, width]` | `[480, 848]` |
| `video_length` | Number of generated frames | `81` |
| `fps` | Frames per second | `16` |
| `weight_dtype` | Model weight precision, use `torch.float16` for GPUs without bf16 support | `torch.bfloat16` |
| `prompt` | Positive prompt describing generated content | `"A cinematic aerial shot..."` |
| `negative_prompt` | Negative prompt to avoid certain content | Chinese negative prompt |
| `guidance_scale` | Guidance strength (the DMD student is CFG-free) | 1.0 |
| `seed` | Random seed for reproducibility | 1024 |
| `num_inference_steps` | Number of inference steps (FastWan uses 3) | 3 |
| `save_path` | Path to save generated videos | `samples/fastvideo-wan2.1-official` |
**GPU Memory Mode Descriptions**:
| Mode | Description | Memory Usage |
|------|------|---------|
| `model_full_load` | Full model loaded to GPU | Highest |
| `model_full_load_and_qfloat8` | Full load + FP8 quantization | High |
| `model_cpu_offload` | Offload model to CPU after use | Medium |
| `model_cpu_offload_and_qfloat8` | CPU offload + FP8 quantization | Medium-Low |
| `model_group_offload` | Layer groups switch between CPU/CUDA | Low |
| `sequential_cpu_offload` | Layer-by-layer offload (slowest) | Lowest |
### 4.2 Text-to-Video (T2V) Inference
Run single-GPU inference:
```bash
python examples/fastvideo/predict_t2v_wan2.1.py
```
Edit `examples/fastvideo/predict_t2v_wan2.1.py` according to your needs. For first-time inference, focus on the following key parameters. For other parameters, refer to the inference parameter descriptions above.
```python
# Choose based on GPU memory
GPU_memory_mode = "model_cpu_offload"
# Base model path (config + dense backbone + VAE + T5)
model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
# VSA keep-ratio; must equal the value the checkpoint was trained at
student_topk = 0.20
# Path to the FastWan student weights (your trained checkpoint, or the official FastWan release)
transformer_path = "models/Personalized_Model/FastWan2.1-T2V-1.3B.pth"
# FastWan sampler + shift (CFG-free few-step student)
sampler_name = "Flow_Unipc"
shift = 3.0
guidance_scale = 1.0
num_inference_steps = 3
# Write based on your generation content
prompt = "A cinematic aerial shot of a snow-covered mountain village at dawn..."
# ...
```
> 💡 The official `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` checkpoint ships in diffusers key naming; `examples/fastvideo/predict_t2v_wan2.1.py` remaps it to the VideoX-Fun civitai layout (including the VSA `to_gate_compress` gate) before loading. A checkpoint produced by this training guide is already in the native layout.
### 4.3 Single-GPU Only
The VSA block-sparse kernel has no sequence-parallel path, so a FastWan student runs on a **single GPU**. Keep `ulysses_degree = ring_degree = 1`; calling `enable_multi_gpus_inference()` raises otherwise. To fit larger resolutions on one card, lower the memory footprint via `GPU_memory_mode` (e.g. `model_cpu_offload` or `sequential_cpu_offload`) rather than splitting across GPUs.
---
## 5. Additional Resources
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
- **FastVideo GitHub**: https://github.com/hao-ai-lab/FastVideo
- **VSA Paper (Video Sparse Attention)**: https://arxiv.org/abs/2505.13389
scripts/fastvideo/README_TRAIN_DISTILL_zh-CN.md
+648
View File
@@ -0,0 +1,648 @@
# FastVideo VSA 蒸馏训练指南
本文档提供了将 Wan2.1 蒸馏为 FastWan2.1（采用 VSA，Video Sparse Attention）的完整工作流，包括环境配置、数据准备、分布式训练和推理测试。
> **注意**: FastVideo VSA 将块稀疏视频注意力（由可学习门控为每个 query 保留约 20% 的 key 块）与 DMD 少步蒸馏结合，使学生模型在约 3 步、无 CFG 的条件下以稀疏注意力推理。训练配方、数据管线与分布式后端与 TurboDiffusion 蒸馏指南完全一致，仅稀疏学生模型类不同。
---
## 目录
- [一、环境配置](#一环境配置)
- [二、数据准备](#二数据准备)
- [2.1 快速测试数据集](#21-快速测试数据集)
- [2.2 数据集结构](#22-数据集结构)
- [2.3 metadata.json 格式](#23-metadatajson-格式)
- [2.4 相对路径与绝对路径使用方案](#24-相对路径与绝对路径使用方案)
- [三、蒸馏训练](#三蒸馏训练)
- [3.1 下载预训练模型](#31-下载预训练模型)
- [3.2 快速开始（DeepSpeed-Zero-2）](#32-快速开始deepspeed-zero-2)
- [3.3 常见训练参数](#33-常见训练参数)
- [3.4 训练验证](#34-训练验证)
- [3.5 使用 FSDP 训练](#35-使用-fsdp-训练)
- [3.6 其他后端](#36-其他后端)
- [3.7 多机分布式训练](#37-多机分布式训练)
- [四、推理测试](#四推理测试)
- [4.1 推理参数解析](#41-推理参数解析)
- [4.2 文本生成视频（T2V）推理](#42-文本生成视频t2v推理)
- [4.3 仅支持单卡](#43-仅支持单卡)
- [五、更多资源](#五更多资源)
---
## 一、环境配置
**方式一：使用 requirements.txt**
```bash
pip install -r requirements.txt
```
**方式二：手动安装依赖**
```bash
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4
```
> 💡 VSA 块稀疏注意力运行在 Triton 内核上；Triton 随 PyTorch（>=2.1）在 Linux 上一并提供，无需额外安装。
**方式三：使用 Docker**
使用 Docker 时，请确保机器上已正确安装 GPU 驱动与 CUDA 环境，然后执行以下命令：
```
# 拉取镜像
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# 进入容器
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
---
## 二、数据准备
### 2.1 快速测试数据集
我们提供了一个包含若干训练样本的测试数据集。
```bash
# 下载官方演示数据集
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
### 2.2 数据集结构
```
📦 datasets/
├── 📂 my_dataset/
│ ├── 📂 train/
│ │ ├── 📄 video001.mp4
│ │ ├── 📄 video002.mp4
│ │ └── 📄 ...
│ └── 📄 metadata.json
```
### 2.3 metadata.json 格式
**相对路径格式**（示例）：
```json
[
{
"file_path": "train/video001.mp4",
"text": "A beautiful sunset over the ocean, golden hour lighting",
"type": "video",
"width": 1024,
"height": 1024
},
{
"file_path": "train/video002.mp4",
"text": "A person walking through a forest, cinematic view",
"type": "video",
"width": 1328,
"height": 1328
}
]
```
**绝对路径格式**：
```json
[
{
"file_path": "/mnt/data/videos/sunset.mp4",
"text": "A beautiful sunset over the ocean",
"type": "video",
"width": 1024,
"height": 1024
}
]
```
**关键字段说明**：
- `file_path`：视频路径（相对或绝对路径）
- `text`：视频描述（英文 prompt）
- `type`：数据类型，固定为 `"video"`
- `width` / `height`：视频尺寸（**建议**提供以便分桶训练。若不提供，训练时会自动读取，当数据存放在 OSS 等较慢系统上时会拖慢训练）。
- 可使用 `scripts/process_json_add_width_and_height.py` 为缺少宽高字段的 JSON 文件提取 width 和 height，同时支持图片与视频。
- 用法：`python scripts/process_json_add_width_and_height.py --input_file datasets/X-Fun-Videos-Demo/metadata.json --output_file datasets/X-Fun-Videos-Demo/metadata_add_width_height.json`。
### 2.4 相对路径与绝对路径使用方案
**相对路径**：
若数据使用相对路径，在训练脚本中配置：
```bash
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
```
**绝对路径**：
若数据使用绝对路径，在训练脚本中配置：
```bash
export DATASET_NAME=""
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
```
> 💡 **建议**：数据集较小且存放在本地时使用相对路径；数据集存放在外部存储（如 NAS、OSS）或跨多台机器共享时使用绝对路径。
---
## 三、蒸馏训练
蒸馏是一次 DMD（distribution matching distillation）训练，包含三个网络：稀疏的 **generator**（FastWan 学生，`FastVideoWanTransformer3DModel`）、稀疏的 **fake_score** critic（同一个类），以及冻结的 **dense** `real_score` 教师（`WanTransformer3DModel`）。学生与 critic 携带 VSA 块稀疏自注意力，教师保持全注意力。
### 3.1 下载预训练模型
```bash
# 创建模型目录
mkdir -p models/Diffusion_Transformer
# 下载 Wan2.1 官方权重
modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
```
### 3.2 快速开始（DeepSpeed-Zero-2）
按照 **2.1 快速测试数据集** 下载数据、按照 **3.1 下载预训练模型** 下载权重后，即可直接复制运行快速开始命令。
推荐使用 DeepSpeed-Zero-2 或 FSDP 训练，这里以 DeepSpeed-Zero-2 为例。
DeepSpeed-Zero-2 与 FSDP 的区别在于模型权重是否分片。**若使用 DeepSpeed-Zero-2 多卡训练时显存不足**，可切换到 FSDP 训练。
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.3 常见训练参数
**关键参数说明**：
| 参数 | 描述 | 示例值 |
|-----|------|-------|
| `--config_path` | 配置文件路径 | `config/wan2.1/wan_civitai.yaml` |
| `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` |
| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` |
| `--train_data_meta` | 训练数据元信息文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` |
| `--train_batch_size` | 每张卡的 batch size | 1 |
| `--image_sample_size` | 图片训练最大分辨率 | 640 |
| `--video_sample_size` | 视频训练最大分辨率 | 640 |
| `--token_sample_size` | Token 采样尺寸 | 640 |
| `--video_sample_stride` | 视频采样步长 | 2 |
| `--video_sample_n_frames` | 视频采样帧数 | 81 |
| `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 |
| `--dataloader_num_workers` | DataLoader 工作进程数 | 8 |
| `--num_train_epochs` | 训练轮数 | 100 |
| `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 |
| `--learning_rate` | 初始学习率（generator） | 2e-06 |
| `--learning_rate_critic` | 初始学习率（critic） | 2e-06 |
| `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
| `--lr_warmup_steps` | 学习率 warmup 步数 | 100 |
| `--seed` | 随机种子 | 42 |
| `--output_dir` | 输出目录 | `output_dir_fastvideo_wan2.1_1.3b` |
| `--gradient_checkpointing` | 启用梯度检查点 | - |
| `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--vae_mini_batch` | VAE 编码 mini-batch 大小 | 1 |
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
| `--enable_bucket` | 启用分桶训练，不裁剪，按分辨率分组 | - |
| `--random_hw_adapt` | 将图片/视频自动缩放到 `[min_size, max_size]` 范围内的随机尺寸 | - |
| `--training_with_video_token_length` | 基于 token 长度训练，支持任意分辨率 | - |
| `--uniform_sampling` | 均匀时间步采样 | - |
| `--low_vram` | 低显存模式 | - |
| `--train_mode` | 训练模式：`normal`（标准）或 `i2v`（图生视频） | `normal` |
| `--resume_from_checkpoint` | 续训路径，使用 `"latest"` 自动选择最新 checkpoint | None |
| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
| `--validation_epochs` | 每 N 轮运行一次验证 | 5 |
| `--validation_prompts` | 用于视频生成验证的 prompt | `"A dog shaking head..."` |
| `--trainable_modules` | 可训练模块（`"."` 表示全部模块） | `"."` |
**蒸馏专用参数**：
| 参数 | 描述 | 示例值 |
|-----|------|-------|
| `--denoising_step_indices_list` | 去噪步列表（蒸馏核心参数） | `1000 750 500 250` |
| `--real_guidance_scale` | 打分用的 real guidance scale | 6.0 |
| `--fake_guidance_scale` | 打分用的 fake guidance scale | 0.0 |
| `--gen_update_interval` | generator 更新间隔 | 5 |
| `--negative_prompt` | 蒸馏用的负向 prompt | 中文负向 prompt |
| `--validation_paths` | I2V 模式的验证图片路径 | 图片路径列表 |
| `--train_sampling_steps` | 训练采样步数 | 1000 |
**VSA 稀疏度**：
- 学生/critic 的自注意力为 VSA 块稀疏。保留比 `student_topk` 在 `FastVideoWanTransformer3DModel` 内默认为 `0.20`（即 80% 稀疏度）；它是模型类默认值，而非命令行旗标。
- 可学习的压缩门控 `to_gate_compress` 以零初始化创建，并在 `--trainable_modules "."` 下自动变为可训练（其参数名与其它模块一样包含点号）。
- `real_score` 教师保持 dense（`WanTransformer3DModel`）且冻结，因此稀疏性是被蒸馏进学生，而非从教师拷贝。
- 推理时的 `student_topk` 须与 checkpoint 训练时的取值一致；修改它会重建注意力算子并重置已训练好的门控。
**采样尺寸配置指南**：
- `video_sample_size` 表示视频的分辨率尺寸；当 `random_hw_adapt` 为 True 时，表示视频与图片分辨率之间的最小值。
- `image_sample_size` 表示图片的分辨率尺寸；当 `random_hw_adapt` 为 True 时，表示视频与图片分辨率之间的最大值。
- `token_sample_size` 表示 `training_with_video_token_length` 为 True 时最大 token 长度对应的分辨率。
- 由于配置容易混淆，**若微调不需要任意分辨率**，建议将 `video_sample_size`、`image_sample_size`、`token_sample_size` 设为同一固定值，例如 **(320, 480, 512, 640, 960)**。
- **全部设为 320** 代表 **240P**。
- **全部设为 480** 代表 **320P**。
- **全部设为 640** 代表 **480P**。
- **全部设为 960** 代表 **720P**。
**其它参数说明**：
- `enable_bucket` 用于启用分桶训练。启用后模型不在中心裁剪图片与视频，而是按分辨率分桶后对完整图片与视频进行训练。
- `random_hw_adapt` 用于对图片与视频自动高宽缩放。启用后，训练图片的高宽以 `image_sample_size` 为最大、`min(video_sample_size, 512)` 为最小。
- `training_with_video_token_length` 指定按 token 长度训练。训练图片与视频的高宽以 `image_sample_size` 为最大、`video_sample_size` 为最小。
- `train_mode` 用于指定训练模式，可为 normal 或 i2v。若只需文本生成视频，保持 `normal`。
- `resume_from_checkpoint` 用于设置从之前的 checkpoint 续训。使用路径或 `"latest"` 自动选择最近可用的 checkpoint。
### 3.4 训练验证
你可以配置验证参数，在训练期间周期性生成测试视频，以便监控训练进度与模型质量。
**验证参数说明**：
| 参数 | 描述 | 推荐值 |
|------|------|--------|
| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
| `--validation_epochs` | 每 N 轮运行一次验证 | 5 |
| `--validation_prompts` | 用于视频生成验证的 prompt | 英文 prompt |
| `--validation_paths` | I2V 模式的验证图片路径（仅 i2v 模式） | `"asset/1.png"` |
**Normal 模式示例**（T2V 验证）：
```bash
--validation_steps=2000 \
--validation_epochs=5 \
--validation_prompts="A dog shaking head. The video is of high quality, and the view is very clear."
```
**说明**：
- 验证视频会保存到 `output_dir` 目录。
- 多 prompt 验证格式：`--validation_prompts "prompt1" "prompt2" "prompt3"`。
- 蒸馏模型验证使用 `denoising_step_indices_list` 定义的少步、并以无 CFG（`guidance_scale=1.0`）运行，与推理保持一致。
### 3.5 使用 FSDP 训练
**若使用 DeepSpeed-Zero-2 多卡训练时显存不足**，可切换到 FSDP 训练。`scripts/fastvideo/train_distill_wan2.1.sh` 是一个开箱即用的 FSDP launcher（它以 `FULL_SHARD` 包装 `WanAttentionBlock`）。
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.6 其他后端
#### 3.6.1 使用 DeepSpeed-Zero-3 进行训练
目前不强烈推荐 DeepSpeed Zero-3。在本仓库中，使用 FSDP 报错更少、更稳定。
DeepSpeed Zero-3 适合高分辨率下的 14B Wan。训练后可用以下命令得到最终模型：
```bash
python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/checkpoint-{your-num-steps}-outputs --max_shard_size 80GB --safe_serialization
```
训练 shell 命令如下：
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
#### 3.6.2 不使用 DeepSpeed 与 FSDP 训练
**不推荐这种方式，因为没有省显存后端时极易导致显存不足**。这里仅作为训练参考 shell 提供。
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
### 3.7 多机分布式训练
**适用于**：超大数据集、更快的训练速度
#### 3.7.1 环境配置
假设有 2 台机器，每台 8 张 GPU：
**机器 0（Master）**：
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export MASTER_ADDR="192.168.1.100" # Master 机器 IP
export MASTER_PORT=10086
export WORLD_SIZE=2 # 机器总数
export NUM_PROCESS=16 # 总进程数 = 机器数 × 8
export RANK=0 # 当前机器 rank（0 或 1）
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
```
**机器 1（Worker）**：
```bash
export MASTER_ADDR="192.168.1.100" # 与 Master 相同
export MASTER_PORT=10086
export WORLD_SIZE=2
export NUM_PROCESS=16
export RANK=1 # 注意这里是 1
NCCL_DEBUG=INFO
# 使用与机器 0 相同的 accelerate launch 命令
```
#### 3.7.2 多机训练注意事项
- **网络要求**：
- 推荐 RDMA/InfiniBand（高性能）
- 没有 RDMA 时，添加环境变量：
```bash
export NCCL_IB_DISABLE=1
export NCCL_P2P_DISABLE=1
```
- **数据同步**：所有机器必须能访问相同的数据路径（NFS/共享存储）
---
## 四、推理测试
### 4.1 推理参数解析
**关键参数说明**（取自 `examples/fastvideo/predict_t2v_wan2.1.py`）：
| 参数 | 描述 | 示例值 |
|------|------|-------|
| `GPU_memory_mode` | GPU 显存模式，选项见下表 | `model_cpu_offload` |
| `ulysses_degree` | Ulysses 并行度（VSA 仅单卡，保持 1） | 1 |
| `ring_degree` | Ring 并行度（VSA 仅单卡，保持 1） | 1 |
| `config_path` | 配置文件路径 | `config/wan2.1/wan_civitai.yaml` |
| `model_name` | 基座模型路径（config + dense 主干 + VAE + T5） | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
| `student_topk` | VSA 保留比；须与训练时取值一致 | `0.20` |
| `sampler_name` | 采样器类型：`Flow`、`Flow_Unipc`、`Flow_DPM++` | `Flow_Unipc` |
| `shift` | 噪声调度 shift（官方 FastWan flow_shift） | `3.0` |
| `transformer_path` | FastWan 学生 Transformer 权重路径 | 训练 checkpoint / 官方 FastWan |
| `vae_path` | 训练好的 VAE 权重路径 | `None` |
| `lora_path` | LoRA 权重路径 | `None` |
| `sample_size` | 生成视频分辨率 `[height, width]` | `[480, 848]` |
| `video_length` | 生成帧数 | `81` |
| `fps` | 每秒帧数 | `16` |
| `weight_dtype` | 模型权重精度，不支持 bf16 的 GPU 用 `torch.float16` | `torch.bfloat16` |
| `prompt` | 描述生成内容的正向 prompt | `"A cinematic aerial shot..."` |
| `negative_prompt` | 用于规避特定内容的负向 prompt | 中文负向 prompt |
| `guidance_scale` | 引导强度（DMD 学生无 CFG） | 1.0 |
| `seed` | 复现用随机种子 | 1024 |
| `num_inference_steps` | 推理步数（FastWan 用 3 步） | 3 |
| `save_path` | 生成视频保存路径 | `samples/fastvideo-wan2.1-official` |
**GPU 显存模式说明**：
| 模式 | 描述 | 显存占用 |
|------|------|---------|
| `model_full_load` | 整个模型加载到 GPU | 最高 |
| `model_full_load_and_qfloat8` | 全量加载 + FP8 量化 | 高 |
| `model_cpu_offload` | 用后将模型卸载到 CPU | 中 |
| `model_cpu_offload_and_qfloat8` | CPU 卸载 + FP8 量化 | 中低 |
| `model_group_offload` | 层组在 CPU/CUDA 间切换 | 低 |
| `sequential_cpu_offload` | 逐层卸载（最慢） | 最低 |
### 4.2 文本生成视频（T2V）推理
运行单卡推理：
```bash
python examples/fastvideo/predict_t2v_wan2.1.py
```
根据需求编辑 `examples/fastvideo/predict_t2v_wan2.1.py`。首次推理重点关注以下关键参数，其它参数参见上文推理参数说明。
```python
# 根据 GPU 显存选择
GPU_memory_mode = "model_cpu_offload"
# 基座模型路径（config + dense 主干 + VAE + T5）
model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
# VSA 保留比；须与 checkpoint 训练时取值一致
student_topk = 0.20
# FastWan 学生权重路径（你训练的 checkpoint，或官方 FastWan 发布权重）
transformer_path = "models/Personalized_Model/FastWan2.1-T2V-1.3B.pth"
# FastWan 采样器 + shift（无 CFG 少步学生）
sampler_name = "Flow_Unipc"
shift = 3.0
guidance_scale = 1.0
num_inference_steps = 3
# 根据你的生成内容编写
prompt = "A cinematic aerial shot of a snow-covered mountain village at dawn..."
# ...
```
> 💡 官方 `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` checkpoint 采用 diffusers 键名；`examples/fastvideo/predict_t2v_wan2.1.py` 会在加载前将其 remap 到 VideoX-Fun civitai 布局（包含 VSA `to_gate_compress` 门控）。本训练指南产出的 checkpoint 已是原生布局。
### 4.3 仅支持单卡
VSA 块稀疏内核没有序列并行路径，因此 FastWan 学生只能在**单卡**上运行。请保持 `ulysses_degree = ring_degree = 1`；否则调用 `enable_multi_gpus_inference()` 会 raise。要在单卡上容纳更大分辨率，请通过 `GPU_memory_mode`（如 `model_cpu_offload` 或 `sequential_cpu_offload`）降低显存占用，而不是跨卡拆分。
---
## 五、更多资源
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
- **FastVideo GitHub**：https://github.com/hao-ai-lab/FastVideo
- **VSA 论文（Video Sparse Attention）**：https://arxiv.org/abs/2505.13389
scripts/fastvideo/train_distill_wan2.1.py
+2411
View File
File diff suppressed because it is too large Load Diff
scripts/fastvideo/train_distill_wan2.1.sh
+48
View File
@@ -0,0 +1,48 @@
# This shell script is used to distill a regular Wan2.1-1.3B into FastWan2.1-1.3B (VSA sparse student)
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
--fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \
--config_path="config/wan2.1/wan_civitai.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=640 \
--video_sample_size=640 \
--token_sample_size=640 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-06 \
--learning_rate_critic=2e-06 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir_fastvideo_wan2.1_1.3b" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--train_mode="normal" \
--trainable_modules "." \
--low_vram
scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md
+2
View File
@@ -172,6 +172,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -211,6 +212,7 @@ bash scripts/wan2.1_causal_forcing/train_ar_diffusion.sh
| `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_ar_diffusion` |
| `--gradient_checkpointing` | Enable activation checkpointing | - |
| `--mixed_precision` | `fp16` / `bf16` | `bf16` |
| `--fsdp_reduce_fp32` | FP32 gradient reduction under FSDP mixed precision; forward and parameter gathering stay low precision. CLI default is off; enabled in the launcher | enabled |
| `--adam_weight_decay` | AdamW weight decay | 3e-2 |
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md
+2
View File
@@ -172,6 +172,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
@@ -211,6 +212,7 @@ bash scripts/wan2.1_causal_forcing/train_ar_diffusion.sh
| `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_ar_diffusion` |
| `--gradient_checkpointing` | 启用激活重计算 | - |
| `--mixed_precision` | `fp16` / `bf16` | `bf16` |
| `--fsdp_reduce_fp32` | FSDP 混合精度下使用 FP32 梯度归约，前向与参数聚合仍保持低精度。CLI 默认关闭，launcher 已启用 | 启用 |
| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md
+2
View File
@@ -183,6 +183,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_ccd" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
@@ -228,6 +229,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh
| `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_ccd` |
| `--gradient_checkpointing` | Enable activation checkpointing | - |
| `--mixed_precision` | `fp16` / `bf16` | `bf16` |
| `--fsdp_reduce_fp32` | FP32 gradient reduction under FSDP mixed precision; forward and parameter gathering stay low precision. CLI default is off; enabled in the launcher | enabled |
| `--adam_weight_decay` | AdamW weight decay | 0.0 |
| `--adam_beta1` | AdamW beta1 (CCD uses 0.0) | 0.0 |
| `--adam_beta2` | AdamW beta2 | 0.999 |
scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md
+2
View File
@@ -183,6 +183,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_ccd" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
@@ -228,6 +229,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh
| `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_ccd` |
| `--gradient_checkpointing` | 启用激活重计算 | - |
| `--mixed_precision` | `fp16` / `bf16` | `bf16` |
| `--fsdp_reduce_fp32` | FSDP 混合精度下使用 FP32 梯度归约，前向与参数聚合仍保持低精度。CLI 默认关闭，launcher 已启用 | 启用 |
| `--adam_weight_decay` | AdamW 权重衰减 | 0.0 |
| `--adam_beta1` | AdamW beta1（CCD 使用 0.0） | 0.0 |
| `--adam_beta2` | AdamW beta2 | 0.999 |
scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md
+45 -5
View File
@@ -26,6 +26,7 @@ This document provides the complete workflow for **Causal-Forcing Stage 3 — Di
- [5.1 Quick Start](#51-quick-start)
- [5.2 Key Parameters](#52-key-parameters)
- [5.3 DMD-Specific Parameters](#53-dmd-specific-parameters)
- [5.4 Optional One-Forcing GAN and Precision Controls](#54-optional-one-forcing-gan-and-precision-controls)
- [6. Use the Trained Checkpoint](#6-use-the-trained-checkpoint)
- [7. Additional Resources](#7-additional-resources)
@@ -103,7 +104,7 @@ The Stage 2 CCD checkpoint (generator/critic init) is loaded via `--ode_transfor
## 4. Prepare Training Data
DMD uses a **TextDataset** (`train_mode="normal"`) — it only needs prompts, not video data, because the generator creates its own training samples via autoregressive rollout. However, a `metadata.json` with prompts is still required.
With `train_mode="normal"`, teacher forcing disabled, and all GAN/R1/R2 weights set to zero, DMD uses a **TextDataset**: only prompts are required because the generator creates training samples via autoregressive rollout. Enabling any GAN/R1/R2 weight switches to **ImageVideoDataset** and requires captioned real videos; see Section 5.4.
### 4.1 Quick Demo Dataset
@@ -118,7 +119,7 @@ Same as Stage 1/2. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](./README_TRA
### 4.3 metadata.json Format
DMD only uses the `"text"` field from each entry. The `file_path` and other fields are ignored when `--train_mode="normal"` (TextDataset mode).
In prompt-only TextDataset mode, DMD only uses the `"text"` field from each entry. GAN/R1/R2 or teacher forcing requires video metadata with `file_path`, `type="video"`, and `text`, as in Stage 1/2.
```json
[
@@ -131,7 +132,7 @@ DMD only uses the `"text"` field from each entry. The `file_path` and other fiel
]
```
> **Note**: You can reuse any video metadata.json — only the `text` field matters for DMD.
> **Note**: Existing video metadata can be reused. Only `text` is read in prompt-only mode; GAN training also reads the video files.
---
@@ -179,6 +180,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_dmd" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
@@ -212,7 +214,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--real_score_pretrained_model_name_or_path` | 14B real-score teacher for DMD | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
| `--config_path` | Model config YAML | `config/wan2.1/wan_civitai.yaml` |
| `--train_data_dir` | Data root directory | `""` |
| `--train_data_meta` | Path to `metadata.json` (prompts only) | `datasets/X-Fun-Videos-Demo/metadata.json` |
| `--train_data_meta` | Metadata path: prompts for plain DMD, captioned videos for GAN/R1/R2 | `datasets/X-Fun-Videos-Demo/metadata.json` |
| `--ode_transformer_path` | Stage 2 CCD checkpoint for generator/critic init | `$STAGE2_CKPT` |
| `--train_batch_size` | Per-GPU batch size | 1 |
| `--gradient_accumulation_steps` | Gradient accumulation steps | 1 |
@@ -264,11 +266,48 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--context_noise` | Context noise level for KV cache update | 0 |
| `--use_teacher_forcing` | Enable teacher forcing (pass clean_x to transformer) | - |
| `--teacher_forcing_prob` | Probability of applying teacher forcing per step (1.0 = always) | 1.0 |
| `--train_mode` | Training mode: `normal` (TextDataset, prompt-only) or `i2v` | `normal` |
| `--train_mode` | `normal` (prompt-only unless teacher forcing or GAN/R1/R2 is enabled) or `i2v` | `normal` |
| `--resume_from_checkpoint` | Resume from checkpoint. Use `"latest"` to auto-select | `"latest"` |
---
### 5.4 Optional One-Forcing GAN and Precision Controls
These options reuse the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms in `scripts/wan2.1_self_forcing/train_distill.py`. All GAN/R1/R2 weights default to **0**, and both precision flags default to **off** in argparse. The AR diffusion, CCD, and DMD launchers enable `--fsdp_reduce_fp32`; GAN/R1/R2 and `--score_timestep_clamp` remain off. AR diffusion and CCD do not gain GAN losses.
The discriminator shares the **fake-score backbone**, with one learnable register token per selected feature layer. Its head is optimized by the critic optimizer, not the generator optimizer. Generator loss is `DMD + weighted GAN`; critic loss is `denoising + weighted GAN + R1 + R2`.
| Parameter | Description | Default |
|-----------|-------------|---------|
| `--gan_g_weight` | Non-saturating generator GAN loss weight | `0.0` |
| `--gan_d_weight` | Discriminator GAN loss weight | `0.0` |
| `--gan_feature_layers` | Fake-score block indices used by the head | `21 29` |
| `--gan_block_ffn_dim` | Head cross-attention FFN dimension | `2048` |
| `--gan_block_num_heads` | Head count; must divide fake-score hidden dimension | `12` |
| `--relativistic_discriminator` | Use paired real/fake logit differences | off |
| `--r1_weight` | APT-style finite-difference penalty on noisy real latents | `0.0` |
| `--r1_sigma` | R1 perturbation standard deviation | `0.01` |
| `--r2_weight` | Normalized finite-difference penalty on noisy generated latents | `0.0` |
| `--r2_sigma` | R2 perturbation standard deviation | `0.01` |
| `--fsdp_reduce_fp32` | FSDP FP32 gradient reduction; parameter gathering stays low precision | off (enabled in the launcher) |
| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off |
For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` for the generator and `softplus(-D_real) + softplus(D_fake)` for the discriminator. Relativistic losses instead use `softplus(-(D_fake - D_real))` and `softplus(-(D_real - D_fake))`.
R1 follows the source's APT form: `r1_weight * mean((D_real_perturbed - D_real)^2)`, **without dividing by sigma**. R2 follows `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`. These are finite-difference penalties, not input-gradient double-backward regularizers. Real/fake samples share timestep and noise; perturbations use independent noise.
To enable GAN, pass positive generator/discriminator weights and, optionally, R1/R2 weights to `train_causal_dmd.py`. Requirements:
- Captioned real videos, not a prompt-only JSON. Use the Stage 1/2 video metadata schema and enough frames for `score_num_frames` and the selected causal block geometry.
- `--gradient_accumulation_steps=1`; DeepSpeed is not supported by the GAN branch.
- Valid feature-layer indices for the selected fake-score model. The default `21 29` layers and `12` heads fit Wan2.1-T2V-1.3B; other backbones may need different settings.
- Enabling only `gan_g_weight` does not supply a discriminator classification-learning objective. Choose the discriminator weight deliberately; nonzero values are not a guaranteed quality improvement.
Trainable fake-score and head master weights stay **FP32**, while forward uses the configured autocast/FSDP precision. Logits, softplus losses, and finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient communication precision, not the forward dtype or model architecture. The clamp changes score-training timesteps, not the generator rollout or validation schedule.
Resume GAN checkpoints with the same GAN/head configuration. The head is saved with fake-score, including FP32 head tensors, and is recreated before loading. It is not part of the exported generator, so inference needs no discriminator or new GAN flags.
---
## 6. Use the Trained Checkpoint
The Stage 3 DMD checkpoint is the **final model** in the Causal-Forcing pipeline. Use it for inference:
@@ -300,4 +339,5 @@ python examples/wan2.1_causal_forcing/predict_t2v.py
## 7. Additional Resources
- **Causal-Forcing Paper**: https://github.com/thu-ml/Causal-Forcing
- **One-Forcing (optional GAN mechanism)**: https://github.com/Aurora-edu/One-Forcing
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md
+57 -17
View File
@@ -26,6 +26,7 @@
- [5.1 快速开始](#51-快速开始)
- [5.2 关键参数](#52-关键参数)
- [5.3 DMD 特有参数](#53-dmd-特有参数)
- [5.4 可选 One-Forcing GAN 与精度控制](#54-可选-one-forcing-gan-与精度控制)
- [六、使用训练好的 Checkpoint](#六使用训练好的-checkpoint)
- [七、更多资源](#七更多资源)
@@ -39,7 +40,7 @@ Stage 3 需要：
2. **Wan2.1-T2V-14B** 模型：作为 DMD real-score teacher。
```bash
# 示例：Stage 2 CCD 训练的 checkpoint
# Example: Stage 2 CCD checkpoint
export STAGE2_CKPT="output_dir_wan2.1_causal_forcing_ccd/checkpoint-5000/transformer/diffusion_pytorch_model.safetensors"
```
@@ -70,10 +71,10 @@ pip install deepspeed==0.17.0 numpy==1.26.4
**方式 3：使用 Docker**
```bash
# 拉取镜像
# Pull the image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
# 进入容器
# Enter the container
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
```
@@ -87,13 +88,13 @@ DMD 需要**两个**预训练模型：
- **Wan2.1-T2V-14B**：非因果的 real-score teacher，用于计算真实分布得分。
```bash
# 创建模型目录
# Create the model directory
mkdir -p models/Diffusion_Transformer
# 下载 Wan2.1 T2V 1.3B（student 基础模型）
# Download the Wan2.1 T2V 1.3B student base model
modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
# 下载 Wan2.1 T2V 14B（DMD real-score teacher）
# Download the Wan2.1 T2V 14B real-score teacher
modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B
```
@@ -103,12 +104,12 @@ Stage 2 CCD checkpoint（生成器/判别器初始化）通过 `--ode_transforme
## 四、数据准备
DMD 使用 **TextDataset**（`train_mode="normal"`）— 只需要提示词，不需要视频数据，因为生成器通过自回归 rollout 自行创建训练样本。但仍需提供包含提示词的 `metadata.json`。
当 `train_mode="normal"`、未启用 teacher forcing 且 GAN/R1/R2 权重均为零时，DMD 使用 **TextDataset**：只需要提示词，生成器通过自回归 rollout 自行创建训练样本。任一 GAN/R1/R2 权重为正时会切换为 **ImageVideoDataset**，需要带描述的真实视频，详见 5.4。
### 4.1 快速测试数据集
```bash
# 下载官方示例数据集
# Download the official demo dataset
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
```
@@ -118,7 +119,7 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun
### 4.3 metadata.json 格式
DMD 仅使用每条记录的 `"text"` 字段。当 `--train_mode="normal"`（TextDataset 模式）时，`file_path` 和其他字段会被忽略。
仅提示词的 TextDataset 模式只读取每条记录的 `"text"` 字段。GAN/R1/R2 或 teacher forcing 需要与 Stage 1/2 相同的视频元数据，包括 `file_path`、`type="video"` 和 `text`。
```json
[
@@ -131,7 +132,7 @@ DMD 仅使用每条记录的 `"text"` 字段。当 `--train_mode="normal"`（Tex
]
```
> **说明**：你可以复用任何视频 metadata.json — DMD 只关心 `text` 字段。
> **说明**：可以复用现有视频元数据。仅提示词模式只读取 `text`，GAN 训练还会读取视频文件。
---
@@ -179,6 +180,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
--output_dir="output_dir_wan2.1_causal_forcing_dmd" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--fsdp_reduce_fp32 \
--adam_weight_decay=0.0 \
--adam_beta1=0.0 \
--adam_beta2=0.999 \
@@ -212,7 +214,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--real_score_pretrained_model_name_or_path` | 14B real-score teacher | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
| `--config_path` | 模型配置 YAML | `config/wan2.1/wan_civitai.yaml` |
| `--train_data_dir` | 数据根目录 | `""` |
| `--train_data_meta` | `metadata.json` 路径（仅含提示词） | `datasets/X-Fun-Videos-Demo/metadata.json` |
| `--train_data_meta` | 元数据路径：普通 DMD 使用提示词，GAN/R1/R2 使用带描述的视频 | `datasets/X-Fun-Videos-Demo/metadata.json` |
| `--ode_transformer_path` | Stage 2 CCD checkpoint（生成器/判别器初始化） | `$STAGE2_CKPT` |
| `--train_batch_size` | 每 GPU batch 大小 | 1 |
| `--gradient_accumulation_steps` | 梯度累积步数 | 1 |
@@ -264,23 +266,60 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
| `--context_noise` | KV 缓存更新的上下文噪声级别 | 0 |
| `--use_teacher_forcing` | 启用 teacher forcing（将 clean_x 传给 transformer） | - |
| `--teacher_forcing_prob` | 每步应用 teacher forcing 的概率（1.0 = 始终） | 1.0 |
| `--train_mode` | 训练模式：`normal`（TextDataset，仅提示词）或 `i2v` | `normal` |
| `--train_mode` | `normal`（未启用 teacher forcing 或 GAN/R1/R2 时仅提示词）或 `i2v` | `normal` |
| `--resume_from_checkpoint` | 恢复训练。使用 `"latest"` 自动选择最新 checkpoint | `"latest"` |
---
### 5.4 可选 One-Forcing GAN 与精度控制
这些选项复用 `scripts/wan2.1_self_forcing/train_distill.py` 中的 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 机制。GAN/R1/R2 权重默认均为 **0**，两个精度开关在 argparse 中默认均**关闭**。AR diffusion、CCD 和 DMD 的 launcher 均启用 `--fsdp_reduce_fp32`；GAN/R1/R2 和 `--score_timestep_clamp` 仍关闭。AR diffusion 和 CCD 不添加 GAN 损失。
判别器共享 **fake-score backbone**，每个选定特征层对应一个可学习 register token。判别器头由 critic optimizer 更新，不加入 generator optimizer。生成器损失为 `DMD + 加权 GAN`，critic 损失为 `denoising + 加权 GAN + R1 + R2`。
| 参数 | 说明 | 默认值 |
|------|------|--------|
| `--gan_g_weight` | 非饱和生成器 GAN 损失权重 | `0.0` |
| `--gan_d_weight` | 判别器 GAN 损失权重 | `0.0` |
| `--gan_feature_layers` | 判别器头使用的 fake-score block 索引 | `21 29` |
| `--gan_block_ffn_dim` | 判别器交叉注意力 FFN 维度 | `2048` |
| `--gan_block_num_heads` | 判别器头数，须整除 fake-score 隐藏维度 | `12` |
| `--relativistic_discriminator` | 使用配对真实/生成样本 logit 差值 | 关闭 |
| `--r1_weight` | 噪声真实 latent 上的 APT 形式有限差分正则权重 | `0.0` |
| `--r1_sigma` | R1 扰动标准差 | `0.01` |
| `--r2_weight` | 噪声生成 latent 上的归一化有限差分正则权重 | `0.0` |
| `--r2_sigma` | R2 扰动标准差 | `0.01` |
| `--fsdp_reduce_fp32` | FSDP 使用 FP32 梯度归约，参数仍以低精度聚合 | 关闭（launcher 已启用） |
| `--score_timestep_clamp` | 将 DMD、生成器 GAN 和 critic 时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关闭 |
设 logits 为 `D_fake` 和 `D_real`，普通 GAN 的生成器损失为 `softplus(-D_fake)`，判别器损失为 `softplus(-D_real) + softplus(D_fake)`。相对论模式分别使用 `softplus(-(D_fake - D_real))` 和 `softplus(-(D_real - D_fake))`。
R1 保持源实现的 APT 形式：`r1_weight * mean((D_real_perturbed - D_real)^2)`，**不除以 sigma**。R2 为 `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`。二者都是有限差分正则，不是需要 double backward 的输入梯度惩罚。真实/生成样本共用时间步和噪声，扰动噪声独立抽样。
开启 GAN 时，向 `train_causal_dmd.py` 传入正的生成器/判别器权重，并按需配置 R1/R2。要求：
- 使用带描述的真实视频，不能只提供提示词 JSON。元数据采用 Stage 1/2 视频格式，帧数须满足 `score_num_frames` 与因果块几何。
- `--gradient_accumulation_steps=1`；GAN 分支不支持 DeepSpeed。
- 特征层索引必须存在。默认 `21 29` 与 `12` 个头适用于 Wan2.1-T2V-1.3B，其他 backbone 可能需要调整。
- 仅开启 `gan_g_weight` 不会提供判别器的分类学习目标。须明确选择判别器权重，非零权重不保证质量提升。
可训练 fake-score 与判别器头的 master weights 保持 **FP32**，前向按配置使用 autocast/FSDP 混合精度。Logits、softplus 损失及有限差分归约使用 FP32。`--fsdp_reduce_fp32` 只改变梯度通信精度，不改变前向 dtype 或模型结构。时间步 clamp 只作用于 score 训练，不改变生成器 rollout 或 validation schedule。
恢复 GAN checkpoint 时应使用相同的 GAN/判别器头配置。判别器头随 fake-score 保存，包括 FP32 判别器头张量，并在加载前重建。导出的 generator 不包含该头，推理无需判别器或新增 GAN 参数。
---
## 六、使用训练好的 Checkpoint
Stage 3 DMD checkpoint 是 Causal-Forcing 流水线的**最终模型**，直接用于推理：
```python
# 在 examples/wan2.1_causal_forcing/predict_t2v.py 中
# In examples/wan2.1_causal_forcing/predict_t2v.py
transformer_path = "output_dir_wan2.1_causal_forcing_dmd/checkpoint-{N}/diffusion_pytorch_model.safetensors"
# DMD Stage 3 推理配置
guidance_scale = 1.0 # CFG 已烘焦到蒸馏权重中
num_inference_steps = 4 # 4 步 DMD
# DMD Stage 3 inference config
guidance_scale = 1.0 # CFG is baked into distilled weights
num_inference_steps = 4 # 4-step DMD
stochastic_sampling = True
num_frame_per_block = 3 # 分块生成
num_frame_per_block = 3 # Chunk-wise generation
```
或运行：
@@ -299,4 +338,5 @@ python examples/wan2.1_causal_forcing/predict_t2v.py
## 七、更多资源
- **Causal-Forcing 论文**：https://github.com/thu-ml/Causal-Forcing
- **One-Forcing（可选 GAN 机制）**：https://github.com/Aurora-edu/One-Forcing
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
scripts/wan2.1_causal_forcing/train_ar_diffusion.py
+13
View File
@@ -18,6 +18,7 @@ scripts/wan2.1_self_forcing/train_distill.py and Causal-Forcing
# See the License for the specific language governing permissions and
import argparse
import dataclasses
import gc
import logging
import math
@@ -314,6 +315,10 @@ def parse_args():
parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true",
help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
)
parser.add_argument(
"--logging_dir",
type=str,
@@ -596,6 +601,9 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
@@ -1254,6 +1262,8 @@ def main():
training_weights_table = training_weights_table.to(accelerator.device)
for epoch in range(first_epoch, args.num_train_epochs):
if global_step >= args.max_train_steps:
break
train_loss = 0.0
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
for step, batch in enumerate(train_dataloader):
@@ -1487,6 +1497,9 @@ def main():
if global_step >= args.max_train_steps:
break
if global_step >= args.max_train_steps:
break
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
with progress_bar.paused():
log_validation(
scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py
+14 -1
View File
@@ -18,6 +18,7 @@ scripts/wan2.1_self_forcing/train_distill.py and Causal-Forcing
# See the License for the specific language governing permissions and
import argparse
import dataclasses
import gc
import logging
import math
@@ -317,6 +318,10 @@ def parse_args():
parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true",
help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
)
parser.add_argument(
"--logging_dir",
type=str,
@@ -610,6 +615,9 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
@@ -769,7 +777,7 @@ def main():
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
torch_dtype=torch.float32,
)
ema_transformer3d.num_frame_per_block = args.num_frame_per_block
ema_transformer3d.independent_first_frame = args.independent_first_frame
@@ -1437,6 +1445,8 @@ def main():
torch.cuda.empty_cache()
for epoch in range(first_epoch, args.num_train_epochs):
if global_step >= args.max_train_steps:
break
train_loss = 0.0
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
for step, batch in enumerate(train_dataloader):
@@ -1737,6 +1747,9 @@ def main():
if global_step >= args.max_train_steps:
break
if global_step >= args.max_train_steps:
break
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
with progress_bar.paused():
log_validation(
scripts/wan2.1_causal_forcing/train_causal_dmd.py
+392 -91
View File
@@ -17,6 +17,7 @@
import argparse
import contextlib
import dataclasses
import gc
import json
import logging
@@ -74,10 +75,12 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512,
get_random_mask)
from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
WanTransformer3DModel,
WanTransformer3DModel_SelfForcing)
WanTransformer3DModel_SelfForcing,
add_wan_gan_discriminator)
from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
WanSelfForcingPipeline)
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.fsdp_ema import FSDPEMA
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
@@ -158,6 +161,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi
return pred_for_score, score_num_frames, True
def repeat_batch(value, repeats):
"""Repeat a batched tensor or per-sample list along the batch dimension."""
if value is None or repeats == 1:
return value
if torch.is_tensor(value):
return torch.cat([value] * repeats)
return list(value) * repeats
def filter_kwargs(cls, kwargs):
import inspect
sig = inspect.signature(cls.__init__)
@@ -494,6 +506,14 @@ def parse_args():
),
)
parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true",
help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
)
parser.add_argument(
"--score_timestep_clamp", action="store_true",
help="Clamp DMD, GAN and critic timesteps to [0.02, 0.98] * num_train_timesteps.",
)
parser.add_argument(
"--dataloader_num_workers",
type=int,
@@ -838,6 +858,47 @@ def parse_args():
help="Probability of applying teacher forcing per step (1.0 = always)"
)
parser.add_argument(
"--gan_g_weight", type=float, default=0.0,
help="One-Forcing: non-saturating generator GAN loss weight (0 disables). Positive GAN/R1/R2 weights require real videos.",
)
parser.add_argument(
"--gan_d_weight", type=float, default=0.0,
help="One-Forcing: discriminator GAN loss weight added to critic denoising loss (0 disables).",
)
parser.add_argument(
"--gan_feature_layers", nargs="+", type=int, default=[21, 29],
help="fake_score block indices feeding the discriminator head, one register token per layer.",
)
parser.add_argument(
"--gan_block_ffn_dim", type=int, default=2048,
help="FFN dimension of the discriminator cross-attention blocks.",
)
parser.add_argument(
"--gan_block_num_heads", type=int, default=12,
help="Number of heads of the discriminator cross-attention blocks.",
)
parser.add_argument(
"--relativistic_discriminator", action="store_true",
help="Use relativistic GAN losses based on real/fake logit differences.",
)
parser.add_argument(
"--r1_weight", type=float, default=0.0,
help="Finite-difference R1 penalty weight on noisy real latents (APT form, not divided by sigma).",
)
parser.add_argument(
"--r1_sigma", type=float, default=0.01,
help="Perturbation standard deviation for the R1 penalty.",
)
parser.add_argument(
"--r2_weight", type=float, default=0.0,
help="Finite-difference R2 penalty weight on noisy generated latents (0.5 * squared logit difference / sigma squared)."
)
parser.add_argument(
"--r2_sigma", type=float, default=0.01,
help="Perturbation standard deviation for the R2 penalty.",
)
args = parser.parse_args()
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
if env_local_rank != -1 and env_local_rank != args.local_rank:
@@ -848,6 +909,14 @@ def parse_args():
def main():
args = parse_args()
use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight))
if any(not math.isfinite(w) or w < 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight)):
raise ValueError("GAN and regularization weights must be finite and non-negative.")
if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1):
raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.")
if use_gan and (args.r1_weight > 0 and (not math.isfinite(args.r1_sigma) or args.r1_sigma <= 0)
or args.r2_weight > 0 and (not math.isfinite(args.r2_sigma) or args.r2_sigma <= 0)):
raise ValueError("Enabled R1/R2 penalties require finite, positive perturbation standard deviations.")
if args.report_to == "wandb" and args.hub_token is not None:
raise ValueError(
@@ -875,6 +944,8 @@ def main():
deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None
fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None
if use_gan and deepspeed_plugin is not None:
raise ValueError("The GAN branch does not support DeepSpeed.")
if deepspeed_plugin is not None:
zero_stage = int(deepspeed_plugin.zero_stage)
fsdp_stage = 0
@@ -896,6 +967,9 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
@@ -1111,28 +1185,55 @@ def main():
param.requires_grad = True
break
if use_gan:
if (not args.gan_feature_layers or any(i < 0 or i >= len(fake_score_transformer3d.blocks) for i in args.gan_feature_layers)
or args.gan_block_ffn_dim <= 0 or args.gan_block_num_heads <= 0
or fake_score_transformer3d.dim % args.gan_block_num_heads != 0):
raise ValueError("GAN feature layers must exist and the positive head count must divide the fake_score hidden dimension.")
# Keep trainable discriminator master weights in fp32; forward uses autocast.
add_wan_gan_discriminator(
fake_score_transformer3d,
args.gan_feature_layers,
args.gan_block_ffn_dim,
args.gan_block_num_heads,
torch.float32,
seed=args.seed or 0,
)
# Create EMA for the generator_transformer3d.
if args.use_ema:
if zero_stage == 3:
raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.")
ema_generator_transformer3d = EMAModel(
generator_transformer3d.parameters(),
model_cls=WanTransformer3DModel_SelfForcing,
model_config=generator_transformer3d.config,
)
if fsdp_plugin is not None:
if getattr(fsdp_plugin, "fsdp_version", 1) != 1:
raise NotImplementedError("FSDPEMA 仅支持 FSDP1。")
ema_module = WanTransformer3DModel_SelfForcing.from_config(generator_transformer3d.config)
ema_generator_transformer3d = FSDPEMA(
ema_module, source=generator_transformer3d,
accelerator=accelerator, fsdp_plugin=fsdp_plugin,
)
else:
ema_generator_transformer3d = EMAModel(
generator_transformer3d.parameters(),
model_cls=WanTransformer3DModel_SelfForcing,
model_config=generator_transformer3d.config,
)
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
if fsdp_stage != 0 or zero_stage == 3:
if fsdp_plugin is not None or zero_stage == 3:
def save_model_hook(models, weights, output_dir):
accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
if accelerator.is_main_process:
from safetensors.torch import save_file
safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
accelerate_state_dict = {k: v.to(dtype=weight_dtype) for k, v in accelerate_state_dict.items()}
accelerate_state_dict = {
k: v.to(dtype=torch.float32 if k.startswith("gan_head.") else weight_dtype)
for k, v in accelerate_state_dict.items()
}
save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
@@ -1149,10 +1250,14 @@ def main():
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
def save_model_hook(models, weights, output_dir):
if accelerator.is_main_process:
if args.use_ema and models[0] is generator_transformer3d:
ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
models[0].save_pretrained(os.path.join(output_dir, "transformer"))
model = accelerator.unwrap_model(models[0])
save_dir = os.path.join(output_dir, "transformer")
if hasattr(model, "gan_head"):
from safetensors.torch import save_file
model.save_config(save_dir)
save_file(model.state_dict(), os.path.join(save_dir, "diffusion_pytorch_model.safetensors"), metadata={"format": "pt"})
else:
model.save_pretrained(save_dir)
if not args.use_deepspeed:
weights.pop()
@@ -1160,34 +1265,23 @@ def main():
pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)
def load_model_hook(models, input_dir):
if args.use_ema and models[0] is generator_transformer3d:
ema_path = os.path.join(input_dir, "transformer_ema")
_, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True)
load_model = WanTransformer3DModel_SelfForcing.from_pretrained(
input_dir, subfolder="transformer_ema",
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
)
load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config)
load_model.load_state_dict(ema_kwargs)
ema_generator_transformer3d.load_state_dict(load_model.state_dict())
ema_generator_transformer3d.to(accelerator.device)
del load_model
for i in range(len(models)):
# pop models so that they are not loaded again
model = models.pop()
# load diffusers style into model
load_model = WanTransformer3DModel.from_pretrained(
input_dir, subfolder="transformer",
low_cpu_mem_usage=True,
)
model.register_to_config(**load_model.config)
model.load_state_dict(load_model.state_dict())
del load_model
# GAN heads are recreated from the same CLI configuration before resume.
if hasattr(model, "gan_head"):
from safetensors.torch import load_file
model.load_state_dict(load_file(os.path.join(input_dir, "transformer", "diffusion_pytorch_model.safetensors")), strict=True)
else:
load_model = WanTransformer3DModel.from_pretrained(
input_dir,
subfolder="transformer",
low_cpu_mem_usage=True,
)
model.register_to_config(**load_model.config)
model.load_state_dict(load_model.state_dict())
del load_model
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
@@ -1201,6 +1295,32 @@ def main():
accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)
if args.use_ema:
def save_ema_hook(models, weights, output_dir):
if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process:
ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
def load_ema_hook(models, input_dir):
ema_path = os.path.join(input_dir, "transformer_ema")
if isinstance(ema_generator_transformer3d, FSDPEMA):
ema_generator_transformer3d.load_pretrained(ema_path)
else:
_, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True)
load_model = WanTransformer3DModel_SelfForcing.from_pretrained(
input_dir, subfolder="transformer_ema",
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=torch.float32,
)
load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config)
load_model.load_state_dict(ema_kwargs)
ema_generator_transformer3d.load_state_dict(load_model.state_dict())
ema_generator_transformer3d.to(accelerator.device)
del load_model
accelerator.register_save_state_pre_hook(save_ema_hook)
accelerator.register_load_state_pre_hook(load_ema_hook)
if args.gradient_checkpointing:
generator_transformer3d.enable_gradient_checkpointing()
fake_score_transformer3d.enable_gradient_checkpointing()
@@ -1267,8 +1387,8 @@ def main():
fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters()))
fake_trainable_params_optim = [
{'params': [], 'lr': args.learning_rate},
{'params': [], 'lr': args.learning_rate / 2},
{'params': [], 'lr': args.learning_rate_critic},
{'params': [], 'lr': args.learning_rate_critic / 2},
]
in_already = []
for name, param in fake_score_transformer3d.named_parameters():
@@ -1281,7 +1401,7 @@ def main():
high_lr_flag = True
fake_trainable_params_optim[0]['params'].append(param)
if accelerator.is_main_process:
print(f"Set {name} to lr : {args.learning_rate}")
print(f"Set {name} to lr : {args.learning_rate_critic}")
break
if high_lr_flag:
continue
@@ -1290,9 +1410,15 @@ def main():
in_already.append(name)
fake_trainable_params_optim[1]['params'].append(param)
if accelerator.is_main_process:
print(f"Set {name} to lr : {args.learning_rate / 2}")
print(f"Set {name} to lr : {args.learning_rate_critic / 2}")
break
if use_gan:
fake_trainable_params_optim[0]['params'].extend(
param for name, param in fake_score_transformer3d.named_parameters()
if name.startswith("gan_head.") and name not in in_already
)
if args.use_came:
optimizer = optimizer_cls(
trainable_params_optim,
@@ -1334,7 +1460,7 @@ def main():
args.random_hw_adapt = False
# Get the dataset
if args.train_mode != "normal" or args.use_teacher_forcing:
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
train_dataset = ImageVideoDataset(
args.train_data_meta, args.train_data_dir,
video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames,
@@ -1362,7 +1488,7 @@ def main():
return length_to_frame_num
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing:
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan:
aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
batch_sampler_generator = torch.Generator().manual_seed(args.seed)
batch_sampler = AspectRatioBatchImageVideoSampler(
@@ -1672,7 +1798,7 @@ def main():
)
fake_score_lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
optimizer=critic_optimizer,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * accelerator.num_processes,
)
@@ -1798,6 +1924,8 @@ def main():
vae_stream_2 = None
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
min_score_timestep = int(0.02 * noise_scheduler.config.num_train_timesteps)
max_score_timestep = int(0.98 * noise_scheduler.config.num_train_timesteps)
def randomize_denoising_step_indices(
denoising_step_indices_list,
@@ -1852,12 +1980,16 @@ def main():
return result
for epoch in range(first_epoch, args.num_train_epochs):
if global_step >= args.max_train_steps:
break
train_dmd_loss = 0.0
# Number of generator backward contributions since the last log flush; the
# generator only backprops every gen_update_interval batches, so its metrics
# must be averaged by contribution count, not by gradient_accumulation_steps.
train_gen_log_count = 0
train_denoising_loss = 0.0
train_gan_g_loss = 0.0
train_gan_d_loss = 0.0
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
for step, batch in enumerate(train_dataloader):
# Data batch sanity check
@@ -1878,7 +2010,7 @@ def main():
save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.mp4", rescale=True)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
if args.train_mode != "normal" or args.use_teacher_forcing:
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
# Convert images to latent space
pixel_values = batch["pixel_values"].to(weight_dtype)
@@ -1942,8 +2074,9 @@ def main():
temp_n_frames = (temp_n_frames - 1) // sample_n_frames_bucket_interval + 1
pixel_values = pixel_values[:, :temp_n_frames, :, :]
mask_pixel_values = mask_pixel_values[:, :temp_n_frames, :, :]
mask = mask[:, :temp_n_frames, :, :]
if args.train_mode != "normal":
mask_pixel_values = mask_pixel_values[:, :temp_n_frames, :, :]
mask = mask[:, :temp_n_frames, :, :]
# Keep all node same token length to accelerate the traning when resolution grows.
if args.keep_all_node_same_token_length:
@@ -1966,8 +2099,9 @@ def main():
actual_video_length = (actual_video_length - 1) // sample_n_frames_bucket_interval + 1
pixel_values = pixel_values[:, :actual_video_length, :, :]
mask_pixel_values = mask_pixel_values[:, :actual_video_length, :, :]
mask = mask[:, :actual_video_length, :, :]
if args.train_mode != "normal":
mask_pixel_values = mask_pixel_values[:, :actual_video_length, :, :]
mask = mask[:, :actual_video_length, :, :]
if args.low_vram:
torch.cuda.empty_cache()
@@ -1990,7 +2124,7 @@ def main():
pixel_values_bs = pixel_values_bs.sample()
new_pixel_values.append(pixel_values_bs)
return torch.cat(new_pixel_values, dim = 0)
if args.use_teacher_forcing:
if args.use_teacher_forcing or use_gan:
clean_latents = _batch_encode_vae(pixel_values)
else:
clean_latents = None
@@ -2022,7 +2156,7 @@ def main():
clip_context.append(_clip_context)
clip_context = torch.cat(clip_context)
if args.use_teacher_forcing:
if args.use_teacher_forcing or use_gan:
target_shape = clean_latents.size()
else:
target_shape = mask_latents.size()
@@ -2426,17 +2560,7 @@ def main():
# Final output — slice generated frames (may be < num_frames for variable-length)
generator_pred_full = output_pred[:, :, :num_generated_frames]
# Gradient mask: first block gets no gradient when generating > min frames
# (matches Self-Forcing reference: model/base.py L182-L190)
min_num_frames_score = args.score_num_frames
need_gradient_mask = (num_generated_frames != min_num_frames_score)
gradient_mask = None
if need_gradient_mask:
gradient_mask = torch.ones_like(generator_pred_full, dtype=torch.bool)
if args.independent_first_frame:
gradient_mask[:, :, :1] = False
else:
gradient_mask[:, :, :args.num_frame_per_block] = False
need_gradient_mask = (num_generated_frames != args.score_num_frames)
# Slice for score computation: last score_num_frames frames
if num_generated_frames > args.score_num_frames:
@@ -2450,11 +2574,13 @@ def main():
generator_pred_for_score = generator_pred_full
score_num_frames = num_generated_frames
# Compute score_mask for DMD loss (matches Self-Forcing: dmd.py L199-204)
score_mask = None
if gradient_mask is not None:
mask_offset = num_generated_frames - score_num_frames
score_mask = gradient_mask[:, :, mask_offset:mask_offset + score_num_frames]
if need_gradient_mask:
score_mask = torch.ones_like(generator_pred_for_score, dtype=torch.bool)
if args.independent_first_frame:
score_mask[:, :, :1] = False
else:
score_mask[:, :, :args.num_frame_per_block] = False
# generator_pred = the sliced version for DMD loss
generator_pred = generator_pred_for_score
@@ -2530,7 +2656,7 @@ def main():
generator_noise_list = [generator_noise[i] for i in range(bsz)]
clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None
# Use block_mask for causal training (一次性处理整个视频)
# Use block_mask for causal training
generator_pred = generator_transformer3d(
x=generator_noise_list,
context=prompt_embeds,
@@ -2563,6 +2689,8 @@ def main():
# Common code for both KV cache and block mask training
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep)
generator_denoised_input = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
@@ -2668,17 +2796,53 @@ def main():
train_dmd_loss += avg_dmd_loss.item()
train_gen_log_count += 1
# --- One-Forcing Generator GAN Loss ---
generator_loss = dmd_loss
if args.gan_g_weight > 0:
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
gan_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
gan_timestep = gan_timestep.clamp(min_score_timestep, max_score_timestep)
gan_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng)
gan_inputs = [add_noise(generator_pred, gan_noise, gan_timestep)]
if args.relativistic_discriminator:
real_latents = clean_latents[:, :, :generator_pred.size(2)].to(generator_pred.dtype)
gan_inputs.append(add_noise(real_latents, gan_noise, gan_timestep))
num_gan_inputs = len(gan_inputs)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, gan_logits = fake_score_transformer3d(
x=torch.cat(gan_inputs).to(dtype=weight_dtype),
context=repeat_batch(prompt_embeds, num_gan_inputs),
t=gan_timestep.repeat(num_gan_inputs),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, num_gan_inputs) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, num_gan_inputs) if args.train_mode != "normal" else None,
gan_classify=True,
)
gan_logits = gan_logits.float().chunk(num_gan_inputs)
if args.relativistic_discriminator:
gan_g_loss = F.softplus(-(gan_logits[0] - gan_logits[1])).mean()
else:
gan_g_loss = F.softplus(-gan_logits[0]).mean()
gan_g_loss = gan_g_loss * args.gan_g_weight
train_gan_g_loss += accelerator.gather(gan_g_loss.detach().repeat(args.train_batch_size)).mean().item()
generator_loss = generator_loss + gan_g_loss
if args.low_vram:
real_score_transformer3d = real_score_transformer3d.to("cpu")
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
if args.gan_g_weight == 0:
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
torch.cuda.empty_cache()
accelerator.backward(dmd_loss)
accelerator.backward(generator_loss)
if accelerator.sync_gradients:
accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
if args.gan_g_weight > 0:
# Discard discriminator gradients from the generator GAN backward.
critic_optimizer.zero_grad()
if args.low_vram:
fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device)
@@ -2967,6 +3131,8 @@ def main():
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
critic_timestep = critic_timestep.clamp(min_score_timestep, max_score_timestep)
critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng)
fake_score_denoised_input = add_noise(
@@ -2975,15 +3141,25 @@ def main():
critic_timestep
)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
)
# Generated and real latents share critic_noise / critic_timestep. Prepare the
# real input and the R1/R2 perturbations once; they are consumed by the critic
# forward below and by the R1/R2 penalties after the main backward(s).
real_noisy = None
r1_perturbation = None
r2_perturbation = None
if use_gan:
real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype)
real_noisy = add_noise(real_latents, critic_noise, critic_timestep)
del real_latents
if args.r1_weight > 0:
r1_perturbation = args.r1_sigma * torch.randn(
real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng
)
if args.r2_weight > 0:
r2_perturbation = args.r2_sigma * torch.randn(
fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype,
device=fake_score_denoised_input.device, generator=torch_rng
)
def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
noise_pred = noise_pred.float()
@@ -2997,11 +3173,136 @@ def main():
final_loss = masked_loss.mean()
return final_loss
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
accelerator_fake_score_transformer3d.backward(denoising_loss)
# --- One-Forcing Discriminator GAN Loss ---
if use_gan and not args.relativistic_discriminator:
# Memory-optimized critic: the non-relativistic gan_d_loss =
# softplus(-real) + softplus(fake) is separable, so run fake and real as two
# sequential 1x-batch forwards and backward the fake pass first. Backwarding
# fake frees its gradient-checkpointing block inputs before the real forward,
# dropping the packed fake+real+R1+R2 forward peak (up to 4x batch) down to
# 1x batch here; R1/R2 re-forward their own 2x batch below. Both backward(s)
# accumulate into fake_score.grad, giving the identical update as the packed
# single-forward critic.
# -- Fake pass (1x batch): denoising output + fake discriminator logit --
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output, fake_gan_logits = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
gan_classify=True,
)
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
fake_logit = fake_gan_logits.float()
fake_d_loss = F.softplus(fake_logit).mean() * args.gan_d_weight
critic_loss = denoising_loss + fake_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
# Fake graph freed: the transformer blocks' stored inputs are released here.
fake_d_detached = fake_d_loss.detach()
del critic_loss, fake_gan_logits, fake_logit, fake_d_loss, fake_score_denoised_output
# -- Real pass (1x batch): discriminator logit only --
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, real_gan_logits = fake_score_transformer3d(
x=real_noisy,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
gan_classify=True,
)
real_logit = real_gan_logits.float()
real_d_loss = F.softplus(-real_logit).mean() * args.gan_d_weight
accelerator_fake_score_transformer3d.backward(real_d_loss)
# Real graph freed. Combine the two detached halves for logging.
gan_d_loss = fake_d_detached + real_d_loss.detach()
train_gan_d_loss += accelerator.gather(gan_d_loss.repeat(args.train_batch_size)).mean().item()
del fake_d_detached, real_gan_logits, real_logit, real_d_loss
else:
# Combined path: 2x-batch forward for relativistic GAN (both logits coupled by
# softplus(-(real - fake))), or 1x-batch forward when GAN is disabled.
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
if use_gan:
fake_score_denoised_output, gan_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0]
gan_logits = gan_logits.float().chunk(2)
else:
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
)
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
critic_loss = denoising_loss
if use_gan:
fake_logit, real_logit = gan_logits[0], gan_logits[1]
if args.relativistic_discriminator:
gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
else:
gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
gan_d_loss = gan_d_loss * args.gan_d_weight
train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
critic_loss = critic_loss + gan_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
del critic_loss
# R1/R2 gradient penalties: each re-forwards its own 2x batch after the main
# backward(s) and finite-differences the perturbed logit against the clean one,
# matching the packed-formula value while keeping the peak at 2x batch.
if use_gan and args.r1_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, r1_logits = fake_score_transformer3d(
x=torch.cat([real_noisy + r1_perturbation, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
r1_logits = r1_logits.float().chunk(2)
r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean()
accelerator_fake_score_transformer3d.backward(r1_loss)
del r1_loss
if use_gan and args.r2_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, r2_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
r2_logits = r2_logits.float().chunk(2)
r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean()
accelerator_fake_score_transformer3d.backward(r2_loss)
del r2_loss
if accelerator_fake_score_transformer3d.sync_gradients:
accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm)
critic_optimizer.step()
@@ -3019,10 +3320,15 @@ def main():
ema_generator_transformer3d.step(generator_transformer3d.parameters())
progress_bar.update(1)
global_step += 1
accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step)
logs = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}
if use_gan:
logs.update({"train_gan_g_loss": train_gan_g_loss / max(train_gen_log_count, 1), "train_gan_d_loss": train_gan_d_loss})
accelerator.log(logs, step=global_step)
train_dmd_loss = 0.0
train_gen_log_count = 0
train_denoising_loss = 0.0
train_gan_g_loss = 0.0
train_gan_d_loss = 0.0
if global_step % args.checkpointing_steps == 0:
if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
@@ -3092,6 +3398,9 @@ def main():
if global_step >= args.max_train_steps:
break
if global_step >= args.max_train_steps:
break
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
with progress_bar.paused():
if args.use_ema:
@@ -3121,12 +3430,6 @@ def main():
# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
if args.use_ema:
# Copy EMA weights to generator on ALL processes before save_state (collective op).
ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
if accelerator.is_main_process:
generator_transformer3d = unwrap_model(generator_transformer3d)
if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
gc.collect()
torch.cuda.empty_cache()
@@ -3135,8 +3438,6 @@ def main():
fake_score_save_path = os.path.join(save_path, "fake_score")
accelerator.save_state(save_path)
accelerator_fake_score_transformer3d.save_state(fake_score_save_path)
if args.use_ema and accelerator.is_main_process:
ema_generator_transformer3d.save_pretrained(os.path.join(save_path, "transformer_ema"))
logger.info(f"Saved state to {save_path}")
accelerator.end_training()
scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md
+23
View File
@@ -258,6 +258,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| `--adam_epsilon` | AdamW epsilon | 1e-10 |
| `--vae_mini_batch` | VAE encoding mini-batch size | 1 |
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
| `--fsdp_reduce_fp32` | Optional: reduce gradients in fp32 under FSDP mixed precision (default reduces in bf16) | - |
| `--enable_bucket` | Enable bucket training, no cropping, group by resolution | - |
| `--random_hw_adapt` | Auto-scale images/videos to random sizes in `[min_size, max_size]` range | - |
| `--training_with_video_token_length` | Train based on token length, supports arbitrary resolutions | - |
@@ -275,12 +276,17 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| Parameter | Description | Example Value |
|-----|------|-------|
| `--denoising_step_indices_list` | Denoising step indices list (core distillation parameter) | `1000 750 500 250` |
| `--randomize_step_indices` | Randomize the denoising timestep indices during training | - |
| `--real_guidance_scale` | Real guidance scale for scoring | 6.0 |
| `--score_timestep_clamp` | Optional: clamp DMD/critic timesteps to `[0.02, 0.98] * num_train_timesteps` | - |
| `--fake_guidance_scale` | Fake guidance scale for scoring | 0.0 |
| `--gen_update_interval` | Generator update interval | 5 |
| `--negative_prompt` | Negative prompt for distillation | Chinese negative prompt |
| `--train_sampling_steps` | Training sampling steps | 1000 |
| `--ode_transformer_path` | Path to ODE-trained weights to load into generator transformer3d | `models/Diffusion_Transformer/Self-Forcing/checkpoints/ode_init.pt` |
| `--real_score_model_name_or_path` | Optional: separate model used as real score (defaults to `pretrained_model_name_or_path`) | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
| `--score_num_frames` | Latent frames in the score-computation window; `fake_score`/`real_score` always receive this many frames | 21 |
| `--min_length_prob_bias` | Probability bias toward sampling the minimum length (`score_num_frames`); `0.0` = uniform sampling | 0.0 |
**Self-Forcing-Specific Parameters**:
@@ -294,6 +300,23 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| `--use_teacher_forcing` | Enable teacher forcing training (pass clean_x to transformer) | - |
| `--teacher_forcing_prob` | Probability of applying teacher forcing per step | 1.0 |
**GAN / Discriminator Parameters (One-Forcing)**:
> All default to `0` / off, so a plain DMD run is unaffected. Setting **any** positive `--gan_g_weight` / `--gan_d_weight` / `--r1_weight` / `--r2_weight` attaches a discriminator head to `fake_score` and **requires real videos** in the training data.
| Parameter | Description | Example Value |
|-----|------|-------|
| `--gan_g_weight` | One-Forcing: weight of the non-saturating GAN loss added to the DMD generator loss (`0` disables) | 0.0 |
| `--gan_d_weight` | One-Forcing: weight of the GAN discriminator loss added to the `fake_score` denoising loss (`0` disables) | 0.0 |
| `--gan_feature_layers` | `fake_score` block indices whose outputs feed the discriminator head (one register token per layer) | `21 29` |
| `--gan_block_ffn_dim` | FFN dim of the discriminator cross-attention blocks | 2048 |
| `--gan_block_num_heads` | Number of heads of the discriminator cross-attention blocks | 12 |
| `--relativistic_discriminator` | Use the relativistic GAN loss (fake vs. real logit difference) | - |
| `--r1_weight` | Finite-difference R1 penalty on noisy real latents (APT form, not divided by sigma) | 0.0 |
| `--r1_sigma` | Perturbation std of the R1 penalty | 0.01 |
| `--r2_weight` | Finite-difference R2 penalty on noisy generated latents | 0.0 |
| `--r2_sigma` | Perturbation std of the R2 penalty | 0.01 |
**Sample Size Configuration Guide**:
- `video_sample_size` represents the resolution size of videos; when `random_hw_adapt` is True, it represents the minimum value between video and image resolutions.
- `image_sample_size` represents the resolution size of images; when `random_hw_adapt` is True, it represents the maximum value between video and image resolutions.
scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md
+23
View File
@@ -259,6 +259,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| `--adam_epsilon` | AdamW epsilon 值 | 1e-10 |
| `--vae_mini_batch` | VAE 编码时的迷你批次大小 | 1 |
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
| `--fsdp_reduce_fp32` | 可选：FSDP 混精下梯度以 fp32 做跨卡 reduce（默认 bf16） | - |
| `--enable_bucket` | 启用分桶训练，不裁剪图片/视频，按分辨率分组训练 | - |
| `--random_hw_adapt` | 自动缩放图片/视频到 `[min_size, max_size]` 范围内的随机尺寸 | - |
| `--training_with_video_token_length` | 根据 token 长度训练，支持任意分辨率 | - |
@@ -276,12 +277,17 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| 参数 | 说明 | 示例值 |
|-----|------|-------|
| `--denoising_step_indices_list` | 去噪步骤列表（蒸馏核心参数） | `1000 750 500 250` |
| `--randomize_step_indices` | 训练时随机化去噪 timestep 索引 | - |
| `--real_guidance_scale` | 用于评分的真实 guidance scale | 6.0 |
| `--score_timestep_clamp` | 可选：将 DMD/critic 的 timestep 截断到 `[0.02, 0.98] * num_train_timesteps` | - |
| `--fake_guidance_scale` | 用于评分的虚拟 guidance scale | 0.0 |
| `--gen_update_interval` | 生成器更新间隔 | 5 |
| `--negative_prompt` | 用于蒸馏的负向提示词 | 中文负向提示词 |
| `--train_sampling_steps` | 训练采样步数 | 1000 |
| `--ode_transformer_path` | ODE 训练权重路径，加载到 generator transformer3d 中 | `models/Diffusion_Transformer/Self-Forcing/checkpoints/ode_init.pt` |
| `--real_score_model_name_or_path` | 可选：单独指定 real score 模型（默认与 `pretrained_model_name_or_path` 相同） | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
| `--score_num_frames` | 评分窗口的 latent 帧数；`fake_score`/`real_score` 始终接收这么多帧 | 21 |
| `--min_length_prob_bias` | 采样最小长度（`score_num_frames`）的概率偏置；`0.0` 表示均匀采样 | 0.0 |
**Self-Forcing 特有参数**：
@@ -295,6 +301,23 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
| `--use_teacher_forcing` | 启用 teacher forcing 训练（将 clean_x 传给 transformer） | - |
| `--teacher_forcing_prob` | 每步应用 teacher forcing 的概率 | 1.0 |
**GAN / 判别器参数（One-Forcing）**：
> 全部默认为 `0` / 关闭，因此普通 DMD 训练不受影响。只要 `--gan_g_weight` / `--gan_d_weight` / `--r1_weight` / `--r2_weight` 中**任一**为正值，就会在 `fake_score` 上挂一个判别器头，并且训练数据**必须包含真实视频**。
| 参数 | 说明 | 示例值 |
|-----|------|-------|
| `--gan_g_weight` | One-Forcing：加到 DMD 生成器损失上的非饱和 GAN 损失权重（`0` 关闭） | 0.0 |
| `--gan_d_weight` | One-Forcing：加到 `fake_score` 去噪损失上的 GAN 判别器损失权重（`0` 关闭） | 0.0 |
| `--gan_feature_layers` | 输出喂给判别器头的 `fake_score` block 索引（每层一个 register token） | `21 29` |
| `--gan_block_ffn_dim` | 判别器 cross-attention block 的 FFN 维度 | 2048 |
| `--gan_block_num_heads` | 判别器 cross-attention block 的头数 | 12 |
| `--relativistic_discriminator` | 使用 relativistic GAN 损失（fake 与 real logit 之差） | - |
| `--r1_weight` | 对含噪真实 latent 的有限差分 R1 惩罚（APT 形式，不除以 sigma） | 0.0 |
| `--r1_sigma` | R1 惩罚的扰动标准差 | 0.01 |
| `--r2_weight` | 对含噪生成 latent 的有限差分 R2 惩罚 | 0.0 |
| `--r2_sigma` | R2 惩罚的扰动标准差 | 0.01 |
**Sample Size 配置指南**：
- `video_sample_size` 表示视频的分辨率大小；当 `random_hw_adapt` 为 True 时，表示视频和图像分辨率的最小值。
- `image_sample_size` 表示图像的分辨率大小；当 `random_hw_adapt` 为 True 时，表示视频和图像分辨率的最大值。
scripts/wan2.1_self_forcing/README_TRAIN_ODE.md
+3
View File
@@ -277,6 +277,9 @@ bash scripts/wan2.1_self_forcing/train_ode.sh
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
| `--trainable_modules` | Trainable modules (`"."` = all) | `"."` |
| `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` |
| `--save_state` | Save the full accelerator state on checkpoint (required to resume training) | - |
| `--use_kv_cache_training` | Run block-by-block KV cache training that matches `pipeline_wan_self_forcing` inference; otherwise fall back to one-shot causal-mask ODE regression (baseline) | - |
| `--prob_full_zero_start` | Per-sample probability of forcing all frames in all blocks to timestep index `0` (pure-noise start) to bridge the train-inference gap; `0.0` disables | 0.0 |
**ODE-specific parameters** (must match Step 1 unless you understand the consequences):
scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md
+3
View File
@@ -277,6 +277,9 @@ bash scripts/wan2.1_self_forcing/train_ode.sh
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
| `--trainable_modules` | 可训练模块（`"."` 表示全量） | `"."` |
| `--resume_from_checkpoint` | 恢复训练路径或 `"latest"` | `latest` |
| `--save_state` | 在 checkpoint 中保存完整的 accelerator 状态（恢复训练所需） | - |
| `--use_kv_cache_training` | 运行逐块 KV 缓存训练，完全匹配 `pipeline_wan_self_forcing` 推理行为；否则回退到一次性的 causal-mask ODE 回归（baseline） | - |
| `--prob_full_zero_start` | 每个样本强制所有块的所有帧使用 timestep 索引 `0`（纯噪声起点）的概率，用于弥合训练-推理差距；`0.0` 表示关闭 | 0.0 |
**ODE 特有参数**（除非清楚后果，否则需与第一步保持一致）：
scripts/wan2.1_self_forcing/train_distill.py
+322 -22
View File
@@ -17,6 +17,7 @@
import argparse
import contextlib
import dataclasses
import gc
import json
import logging
@@ -74,10 +75,12 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512,
get_random_mask)
from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
WanTransformer3DModel,
WanTransformer3DModel_SelfForcing)
WanTransformer3DModel_SelfForcing,
add_wan_gan_discriminator)
from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
WanSelfForcingPipeline)
from videox_fun.utils.discrete_sampler import DiscreteSampling
from videox_fun.utils.fsdp_ema import FSDPEMA
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
get_image_to_video_latent,
@@ -165,6 +168,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi
return pred_for_score, score_num_frames, True
def repeat_batch(value, repeats):
"""Repeat a batched tensor or per-sample list along the batch dimension."""
if value is None or repeats == 1:
return value
if torch.is_tensor(value):
return torch.cat([value] * repeats)
return list(value) * repeats
def filter_kwargs(cls, kwargs):
import inspect
sig = inspect.signature(cls.__init__)
@@ -492,6 +504,8 @@ def parse_args():
parser.add_argument(
"--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
)
parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.")
parser.add_argument("--ema_decay", type=float, default=0.9999, help="EMA decay applied per generator update (Causal-Forcing uses 0.99).")
parser.add_argument(
"--allow_tf32",
action="store_true",
@@ -523,6 +537,9 @@ def parse_args():
parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true", help="Reduce gradients in fp32 under FSDP mixed precision (params are still gathered in low precision)."
)
parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
parser.add_argument(
@@ -722,6 +739,12 @@ def parse_args():
default=None,
help=("If you want to load the ode-trained weight into generator transformer3d, input its path."),
)
parser.add_argument(
"--real_score_model_name_or_path",
type=str,
default=None,
help=("Path to a separate model (e.g. Wan2.1-T2V-14B) used as real score. Defaults to pretrained_model_name_or_path."),
)
parser.add_argument(
"--vae_path",
type=str,
@@ -782,6 +805,11 @@ def parse_args():
default=4.5,
help="The cfg scale for real score.",
)
parser.add_argument(
"--score_timestep_clamp",
action="store_true",
help="Clamp DMD and critic timesteps to [0.02, 0.98] * num_train_timesteps.",
)
parser.add_argument(
'--denoising_step_indices_list',
nargs='+',
@@ -842,6 +870,67 @@ def parse_args():
default=1.0,
help="Probability of applying teacher forcing per step (1.0 = always)"
)
parser.add_argument(
"--gan_g_weight",
type=float,
default=0.0,
help="One-Forcing: weight of the non-saturating GAN loss added to the DMD generator loss (0 disables). "
"Any positive GAN/R1/R2 weight attaches a discriminator head to fake_score and requires real videos."
)
parser.add_argument(
"--gan_d_weight",
type=float,
default=0.0,
help="One-Forcing: weight of the GAN discriminator loss added to the fake_score denoising loss (0 disables)."
)
parser.add_argument(
"--gan_feature_layers",
nargs="+",
type=int,
default=[21, 29],
help="fake_score block indices whose outputs feed the discriminator head (one register token per layer)."
)
parser.add_argument(
"--gan_block_ffn_dim",
type=int,
default=2048,
help="FFN dim of the discriminator cross-attention blocks."
)
parser.add_argument(
"--gan_block_num_heads",
type=int,
default=12,
help="Number of heads of the discriminator cross-attention blocks."
)
parser.add_argument(
"--relativistic_discriminator",
action="store_true",
help="Use the relativistic GAN loss (fake vs. real logit difference)."
)
parser.add_argument(
"--r1_weight",
type=float,
default=0.0,
help="Finite-difference R1 penalty on noisy real latents (APT form, not divided by sigma)."
)
parser.add_argument(
"--r1_sigma",
type=float,
default=0.01,
help="Perturbation std of the R1 penalty."
)
parser.add_argument(
"--r2_weight",
type=float,
default=0.0,
help="Finite-difference R2 penalty on noisy generated latents."
)
parser.add_argument(
"--r2_sigma",
type=float,
default=0.01,
help="Perturbation std of the R2 penalty."
)
args = parser.parse_args()
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
@@ -914,6 +1003,9 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
@@ -970,6 +1062,9 @@ def main():
args.mixed_precision = accelerator.mixed_precision
args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list]
use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight))
if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1):
raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.")
# Load scheduler, tokenizer and models.
noise_scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
@@ -1030,7 +1125,7 @@ def main():
low_cpu_mem_usage=True,
).to(weight_dtype)
real_score_transformer3d = WanTransformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
os.path.join(args.real_score_model_name_or_path or args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
).to(weight_dtype)
@@ -1062,7 +1157,8 @@ def main():
state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()}
m, u = generator_transformer3d.load_state_dict(state_dict, strict=False)
m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False)
if args.real_score_model_name_or_path is None:
m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False)
m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
assert len(u) == 0
@@ -1116,6 +1212,33 @@ def main():
if trainable_module_name in name:
param.requires_grad = True
break
if use_gan:
add_wan_gan_discriminator(
fake_score_transformer3d, args.gan_feature_layers, args.gan_block_ffn_dim, args.gan_block_num_heads,
weight_dtype, seed=args.seed or 0,
)
# Create EMA for the generator_transformer3d.
if args.use_ema:
if zero_stage == 3:
raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.")
if fsdp_plugin is not None:
if getattr(fsdp_plugin, "fsdp_version", 1) != 1:
raise NotImplementedError("FSDPEMA only supports FSDP1.")
ema_module = WanTransformer3DModel_SelfForcing.from_config(generator_transformer3d.config)
ema_generator_transformer3d = FSDPEMA(
ema_module, source=generator_transformer3d,
accelerator=accelerator, fsdp_plugin=fsdp_plugin,
decay=args.ema_decay,
)
else:
ema_generator_transformer3d = EMAModel(
generator_transformer3d.parameters(),
decay=args.ema_decay,
model_cls=WanTransformer3DModel_SelfForcing,
model_config=generator_transformer3d.config,
)
# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
@@ -1178,6 +1301,32 @@ def main():
accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)
if args.use_ema:
def save_ema_hook(models, weights, output_dir):
if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process:
ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
def load_ema_hook(models, input_dir):
ema_path = os.path.join(input_dir, "transformer_ema")
if isinstance(ema_generator_transformer3d, FSDPEMA):
ema_generator_transformer3d.load_pretrained(ema_path)
else:
_, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True)
load_model = WanTransformer3DModel_SelfForcing.from_pretrained(
input_dir, subfolder="transformer_ema",
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=torch.float32,
)
load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config)
load_model.load_state_dict(ema_kwargs)
ema_generator_transformer3d.load_state_dict(load_model.state_dict())
ema_generator_transformer3d.to(accelerator.device)
del load_model
accelerator.register_save_state_pre_hook(save_ema_hook)
accelerator.register_load_state_pre_hook(load_ema_hook)
if args.gradient_checkpointing:
generator_transformer3d.enable_gradient_checkpointing()
fake_score_transformer3d.enable_gradient_checkpointing()
@@ -1269,6 +1418,11 @@ def main():
if accelerator.is_main_process:
print(f"Set {name} to lr : {args.learning_rate_critic / 2}")
break
if use_gan:
fake_trainable_params_optim[0]['params'].extend(
param for name, param in fake_score_transformer3d.named_parameters()
if name.startswith("gan_head.") and name not in in_already
)
if args.use_came:
optimizer = optimizer_cls(
@@ -1311,7 +1465,7 @@ def main():
args.random_hw_adapt = False
# Get the dataset
if args.train_mode != "normal" or args.use_teacher_forcing:
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
train_dataset = ImageVideoDataset(
args.train_data_meta, args.train_data_dir,
video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames,
@@ -1339,7 +1493,7 @@ def main():
return length_to_frame_num
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing:
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan:
aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
batch_sampler_generator = torch.Generator().manual_seed(args.seed)
batch_sampler = AspectRatioBatchImageVideoSampler(
@@ -1682,6 +1836,9 @@ def main():
if args.train_mode != "normal":
clip_image_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if args.use_ema:
ema_generator_transformer3d.to(accelerator.device)
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
if overrode_max_train_steps:
@@ -1772,6 +1929,8 @@ def main():
vae_stream_2 = None
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
min_score_timestep = int(0.02 * noise_scheduler.config.num_train_timesteps)
max_score_timestep = int(0.98 * noise_scheduler.config.num_train_timesteps)
def randomize_denoising_step_indices(
denoising_step_indices_list,
@@ -1832,6 +1991,8 @@ def main():
# must be averaged by contribution count, not by gradient_accumulation_steps.
train_gen_log_count = 0
train_denoising_loss = 0.0
train_gan_g_loss = 0.0
train_gan_d_loss = 0.0
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
for step, batch in enumerate(train_dataloader):
# Data batch sanity check
@@ -1852,7 +2013,7 @@ def main():
save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.mp4", rescale=True)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
if args.train_mode != "normal" or args.use_teacher_forcing:
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
# Convert images to latent space
pixel_values = batch["pixel_values"].to(weight_dtype)
@@ -1964,7 +2125,7 @@ def main():
pixel_values_bs = pixel_values_bs.sample()
new_pixel_values.append(pixel_values_bs)
return torch.cat(new_pixel_values, dim = 0)
if args.use_teacher_forcing:
if args.use_teacher_forcing or use_gan:
clean_latents = _batch_encode_vae(pixel_values)
else:
clean_latents = None
@@ -1996,7 +2157,7 @@ def main():
clip_context.append(_clip_context)
clip_context = torch.cat(clip_context)
if args.use_teacher_forcing:
if args.use_teacher_forcing or use_gan:
target_shape = clean_latents.size()
else:
target_shape = mask_latents.size()
@@ -2504,7 +2665,7 @@ def main():
generator_noise_list = [generator_noise[i] for i in range(bsz)]
clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None
# Use block_mask for causal training (一次性处理整个视频)
# Use block_mask for causal training (process the whole video at once)
generator_pred = generator_transformer3d(
x=generator_noise_list,
context=prompt_embeds,
@@ -2537,6 +2698,8 @@ def main():
# Common code for both KV cache and block mask training
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep)
generator_denoised_input = add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
@@ -2642,17 +2805,53 @@ def main():
train_dmd_loss += avg_dmd_loss.item()
train_gen_log_count += 1
# One-Forcing GAN generator loss: gradients flow through fake_score into generator_pred
generator_loss = dmd_loss
if args.gan_g_weight > 0:
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
gan_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
gan_timestep = gan_timestep.clamp(min_score_timestep, max_score_timestep)
gan_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng)
gan_inputs = [add_noise(generator_pred, gan_noise, gan_timestep)]
if args.relativistic_discriminator:
real_latents = clean_latents[:, :, :generator_pred.size(2)].to(generator_pred.dtype)
gan_inputs.append(add_noise(real_latents, gan_noise, gan_timestep))
num_gan_inputs = len(gan_inputs)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, gan_logits = fake_score_transformer3d(
x=torch.cat(gan_inputs).to(dtype=weight_dtype),
context=repeat_batch(prompt_embeds, num_gan_inputs),
t=gan_timestep.repeat(num_gan_inputs),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, num_gan_inputs) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, num_gan_inputs) if args.train_mode != "normal" else None,
gan_classify=True,
)
gan_logits = gan_logits.float().chunk(num_gan_inputs)
if args.relativistic_discriminator:
gan_g_loss = F.softplus(-(gan_logits[0] - gan_logits[1])).mean()
else:
gan_g_loss = F.softplus(-gan_logits[0]).mean()
gan_g_loss = gan_g_loss * args.gan_g_weight
train_gan_g_loss += accelerator.gather(gan_g_loss.detach().repeat(args.train_batch_size)).mean().item()
generator_loss = generator_loss + gan_g_loss
if args.low_vram:
real_score_transformer3d = real_score_transformer3d.to("cpu")
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
if args.gan_g_weight == 0:
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
torch.cuda.empty_cache()
accelerator.backward(dmd_loss)
accelerator.backward(generator_loss)
if accelerator.sync_gradients:
accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
if args.gan_g_weight > 0:
# Drop the fake_score grads produced by the generator GAN loss
critic_optimizer.zero_grad()
if args.low_vram:
fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device)
@@ -2941,6 +3140,8 @@ def main():
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.score_timestep_clamp:
critic_timestep = critic_timestep.clamp(min_score_timestep, max_score_timestep)
critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng)
fake_score_denoised_input = add_noise(
@@ -2950,14 +3151,41 @@ def main():
)
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,
)
if use_gan:
# Generated and real latents share critic_noise / critic_timestep. The critic forward is
# split into three passes so the peak activation stays at 2x batch instead of 4x:
# A) fake + real for the denoising and GAN loss, B) real + R1 perturbation, C) fake + R2.
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)
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
)
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,
)
def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
noise_pred = noise_pred.float()
@@ -2974,8 +3202,53 @@ def main():
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)
critic_loss = denoising_loss
if use_gan:
fake_logit, real_logit = gan_logits[0], gan_logits[1]
if args.relativistic_discriminator:
gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
else:
gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
gan_d_loss = gan_d_loss * args.gan_d_weight
train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
critic_loss = critic_loss + gan_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
del critic_loss
if use_gan and args.r1_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, r1_logits = fake_score_transformer3d(
x=torch.cat([real_noisy + r1_perturbation, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
r1_logits = r1_logits.float().chunk(2)
r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean()
accelerator_fake_score_transformer3d.backward(r1_loss)
del r1_loss
if use_gan and args.r2_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, r2_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
r2_logits = r2_logits.float().chunk(2)
r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean()
accelerator_fake_score_transformer3d.backward(r2_loss)
del r2_loss
if accelerator_fake_score_transformer3d.sync_gradients:
accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm)
critic_optimizer.step()
@@ -2988,13 +3261,22 @@ def main():
# Checks if the accelerator has performed an optimization step behind the scenes
if accelerator.sync_gradients:
# EMA tracks the generator only on rounds where it actually ran optimizer.step();
# with gen_update_interval > 1 the critic-only rounds leave the generator unchanged.
if args.use_ema and generator_update:
ema_generator_transformer3d.step(generator_transformer3d.parameters())
progress_bar.update(1)
global_step += 1
accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step)
train_log = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}
if use_gan:
train_log.update({"train_gan_g_loss": train_gan_g_loss / max(train_gen_log_count, 1), "train_gan_d_loss": train_gan_d_loss})
accelerator.log(train_log, step=global_step)
train_dmd_loss = 0.0
train_gen_log_count = 0
train_denoising_loss = 0.0
train_gan_g_loss = 0.0
train_gan_d_loss = 0.0
if global_step % args.checkpointing_steps == 0:
if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
@@ -3038,6 +3320,10 @@ def main():
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
with progress_bar.paused():
if args.use_ema:
# Store the generator parameters temporarily and load the EMA parameters to perform inference.
ema_generator_transformer3d.store(generator_transformer3d.parameters())
ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
log_validation(
vae,
text_encoder,
@@ -3050,8 +3336,15 @@ def main():
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original generator parameters.
ema_generator_transformer3d.restore(generator_transformer3d.parameters())
logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
if use_gan:
logs["gan_d_loss"] = gan_d_loss.detach().item()
if args.gan_g_weight > 0:
logs["gan_g_loss"] = gan_g_loss.detach().item()
progress_bar.set_postfix(**logs)
if global_step >= args.max_train_steps:
@@ -3059,6 +3352,10 @@ def main():
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
with progress_bar.paused():
if args.use_ema:
# Store the generator parameters temporarily and load the EMA parameters to perform inference.
ema_generator_transformer3d.store(generator_transformer3d.parameters())
ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
log_validation(
vae,
text_encoder,
@@ -3071,6 +3368,9 @@ def main():
weight_dtype,
global_step,
)
if args.use_ema:
# Switch back to the original generator parameters.
ema_generator_transformer3d.restore(generator_transformer3d.parameters())
# Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever
# something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto
scripts/wan2.1_self_forcing/train_ode.py
+10
View File
@@ -15,6 +15,7 @@
# See the License for the specific language governing permissions and
import argparse
import dataclasses
import gc
import logging
import math
@@ -362,6 +363,10 @@ def parse_args():
parser.add_argument("--adam_weight_decay", type=float, default=3e-2, help="Weight decay to use.")
parser.add_argument("--adam_epsilon", type=float, default=1e-10, help="Epsilon value for the Adam optimizer.")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument(
"--fsdp_reduce_fp32", action="store_true",
help="Reduce gradients in fp32 under FSDP mixed precision (params are still gathered in low precision).",
)
parser.add_argument(
"--mixed_precision",
type=str,
@@ -581,6 +586,11 @@ def main():
else:
fsdp_stage = 0
print(f"Using FSDP stage: {fsdp_stage}")
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
fsdp_plugin.mixed_precision_policy = dataclasses.replace(
fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32
)
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
args.use_fsdp = True
if fsdp_stage == 3:
videox_fun/dist/fsdp.py
Vendored
+7 -6
View File
@@ -146,6 +146,13 @@ def shard_model(
and p.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2)):
p.data = p.data.to(param_dtype)
# `device_id` only moves the managed params/buffers, but `sync_module_states` checks every buffer of the root
# module, including those of the ignored modules (e.g. a computed rope `inv_freq`). Place the ignored modules
# on the device before wrapping. They are small fp32 heads kept replicated (not sharded), so FSDP's CPUOffload
# does not cover them: they stay GPU-resident even in offload mode, which is the intended numerical-precision pin.
for m in ignored_modules:
m.to(device_id)
# With CPU offload, bound the prefetch depth so at most one extra all-gather is in flight; otherwise FSDP
# would keep the gathered (unsharded) params of upcoming blocks resident and eat the savings.
cpu_offload = CPUOffload(offload_params=True) if offload_to_cpu else None
@@ -164,12 +171,6 @@ def shard_model(
limit_all_gathers=offload_to_cpu,
ignored_states=ignored_modules if ignored_modules else None)
# `device_id`/`sync_module_states` only manage the sharded params; the ignored modules must be placed on
# the device explicitly. They are small fp32 heads kept replicated (not sharded), so FSDP's CPUOffload does
# not cover them: they stay GPU-resident even in offload mode, which is the intended numerical-precision pin.
for m in ignored_modules:
m.to(device_id)
return model
videox_fun/models/__init__.py
+5
View File
@@ -94,7 +94,12 @@ from .spark_transformer3d import (SparkWan2_2Transformer3DModel,
SparkWanSelfAttention,
SparkWanTransformer3DModel)
from .turbowan_transformer3d import TurboWanTransformer3DModel
from .fastvideo_wan_transformer3d import (FastVideoWanSelfAttention,
FastVideoWanTransformer3DModel)
from .fastvideo_minimax_h3_transformer3d import \
FastVideoMiniMaxH3Transformer3DModel
from .wan_audio_encoder import WanAudioEncoder
from .wan_gan_discriminator import add_wan_gan_discriminator
from .wan_image_encoder import CLIPModel
from .wan_latent_upsampler import WanLatentUpsamplerModel
from .wan_tae import AutoencoderTinyWan
videox_fun/models/attention_kernel.py
+565
View File
@@ -14,10 +14,575 @@ Citation (please cite if you use this code):
}
"""
import math
import torch
import triton
import triton.language as tl
# ======================================================================================
# This module holds TWO sparse-attention kernel families, kept in contiguous sections.
#
# SECTION A -- VSA, variable-length block-sparse. This is the path the repo actually runs
# for MiniMax-H3 and Wan sparse attention (attention_utils.VSASparseAttention.forward_h3
# -> sparse_linear_attention_varlen). Everything that path needs lives in this section,
# in call order: the two vendored upstream blocks (the official tile-64 forward kernel,
# then the on-device mask -> key-list compaction), our own varlen backward kernels, the
# autograd Function, and its wrapper.
#
# SECTION B -- the fixed-length SLA original this file was modified from (see the licence
# header above). Still live, NOT dead code: attention_utils' flat sparse path calls
# get_block_map + _sparse_linear_attention.apply, and its tile pooling calls mean_pool.
#
# The two families share no kernels. A symbol suffixed `_varlen` belongs to section A.
# ======================================================================================
# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> VENDORED FROM UPSTREAM (Apache-2.0) -- BLOCK 1/2 BEGIN >>>>>>>>>>>>>>>>>>>>
# Source: hao-ai-lab/fastvideo `fastvideo_kernel` 0.3.5, file
# python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py, lines 11-166 and 703-752.
# It replaces the hand-written `_attn_fwd_varlen` that used to live further down this module.
#
# Why vendored: upstream's compiled sm90 CUDA kernel (`fastvideo_kernel._C`) ships only cp312 /
# torch==2.12.0 wheels, so it is not installable here (we run torch 2.7.0 / py3.10), whereas its
# pure-Triton tile-64 block-sparse forward runs as-is with no build step. A kernel-level micro-bench
# on the Wan VSA workload (seq 33390, 672 tiles, sparsity 0.8, bf16, H800) measured this forward at
# 6.31 ms/call vs `_attn_fwd_varlen` at 13.48 ms/call (2.14x), with identical output (rel_l2
# 6.2e-5). The wins are @triton.autotune over num_stages/num_warps, tl.make_block_ptr + tl.advance
# 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
# two `SPARSE ADDITION BEGIN/END` banner comments upstream wrapped around that import are
# dropped with it (the one around `@triton.autotune` below is kept);
# 2. the inner dot's `p.to(tl.bfloat16)` is `p.to(v.dtype)` here, because our varlen path and
# its fp32 oracle tests must not be forced to bf16 (output identical when v IS bf16).
# Everything else between the BEGIN/END markers is byte-identical to upstream. Our own backward
# (`_attn_bwd_*_varlen`) is untouched by this block: it supplies the forward only.
# BLOCK_M / BLOCK_N are fixed at 64 because they are structural, not tunable:
# the kernel indexes the top-k list per BLOCK_M q-tile and addresses keys as
# kv_idx * BLOCK_N, so both must match the granularity q2k_index and
# variable_block_sizes were built at.
#
# num_stages / num_warps ARE free, and the previous {3, 4, 7} was inherited from
# the upstream tutorial rather than tuned here. It skips 5 and 6; on Blackwell
# (sm_121) the optimum is num_stages=5, so the search could not reach it. Both
# block paths independently select 5 once it is available. Autotune still picks
# per architecture, so other GPUs re-tune rather than inheriting this choice.
configs = [
triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
for BM in [64]\
for BN in [64]\
for s in [2, 3, 4, 5, 6, 7]\
for w in [4, 8]\
]
# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
@triton.autotune(configs, key=["N_CTX_Q", "HEAD_DIM"])
@triton.jit
def _attn_fwd_sparse(
Q,
K,
V,
sm_scale, #
q2k_index,
q2k_num,
max_kv_blks, #
variable_block_sizes,
M,
Out, #
stride_qz,
stride_qh,
stride_qm,
stride_qk,
stride_kz,
stride_kh,
stride_kn,
stride_kk,
stride_vz,
stride_vh,
stride_vk,
stride_vn,
stride_oz,
stride_oh,
stride_om,
stride_on,
Z,
H,
N_CTX_Q, #
N_CTX_KV, #
HEAD_DIM: tl.constexpr, #
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
STAGE: tl.constexpr):
"""
64×64 **block-sparse** forward kernel. Back-prop kernels remain dense
(32×64 and 64×32) – memory footprint unchanged.
"""
# ----- program-id mapping -----
q_blk = tl.program_id(0) # Q-tile index
off_hz = tl.program_id(1) # fused (batch, head)
b = off_hz // H
h = off_hz % H
q_tiles = N_CTX_Q // BLOCK_M
meta_base = ((b * H + h) * q_tiles + q_blk)
kv_blocks = tl.load(q2k_num + meta_base) # int32
kv_ptr = q2k_index + meta_base * max_kv_blks # ptr to list
# ----- base pointers -----
# Note: when q and kv have different sequence lengths, their per-(batch,head)
# strides differ, so we must compute separate base offsets.
q_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
k_off = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
v_off = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
o_off = (b.to(tl.int64) * stride_oz + h.to(tl.int64) * stride_oh)
Q_ptr = tl.make_block_ptr(base=Q + q_off,
shape=(N_CTX_Q, HEAD_DIM),
strides=(stride_qm, stride_qk),
offsets=(q_blk * BLOCK_M, 0),
block_shape=(BLOCK_M, HEAD_DIM),
order=(1, 0))
K_base = tl.make_block_ptr(base=K + k_off,
shape=(HEAD_DIM, N_CTX_KV),
strides=(stride_kk, stride_kn),
offsets=(0, 0),
block_shape=(HEAD_DIM, BLOCK_N),
order=(0, 1))
v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
V_base = tl.make_block_ptr(base=V + v_off,
shape=(N_CTX_KV, HEAD_DIM),
strides=(stride_vk, stride_vn),
offsets=(0, 0),
block_shape=(BLOCK_N, HEAD_DIM),
order=v_order)
O_ptr = tl.make_block_ptr(base=Out + o_off,
shape=(N_CTX_Q, HEAD_DIM),
strides=(stride_om, stride_on),
offsets=(q_blk * BLOCK_M, 0),
block_shape=(BLOCK_M, HEAD_DIM),
order=(1, 0))
# ----- accumulators -----
offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M)
m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
qk_scale = sm_scale * 1.44269504 # 1/ln2
q = tl.load(Q_ptr)
# ----- sparse loop over valid K/V tiles -----
for i in range(0, kv_blocks):
kv_idx = tl.load(kv_ptr + i).to(tl.int32)
block_size = tl.load(variable_block_sizes + kv_idx)
K_ptr = tl.advance(K_base, (0, kv_idx * BLOCK_N))
V_ptr = tl.advance(V_base, (kv_idx * BLOCK_N, 0))
k = tl.load(K_ptr)
qk = tl.dot(q, k)
# mask out invalid columns
mask = tl.arange(0, BLOCK_N) < block_size
qk = tl.where(mask[None, :], qk, -float("inf"))
m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
p = tl.math.exp2(qk * qk_scale - m_ij[:, None])
l_ij = tl.sum(p, 1)
alpha = tl.math.exp2(m_i - m_ij)
l_i = l_i * alpha + l_ij
acc = acc * alpha[:, None]
v = tl.load(V_ptr)
acc = tl.dot(p.to(v.dtype), v, acc) # dtype-agnostic: official hardcodes bf16,
# but our varlen path/tests also run fp32; p.to(v.dtype) keeps bf16 identical.
m_i = m_ij
# ----- epilogue -----
m_i += tl.math.log2(l_i)
acc = acc / l_i[:, None]
tl.store(M + off_hz * N_CTX_Q + offs_m, m_i)
tl.store(O_ptr, acc.to(Out.type.element_ty))
def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, variable_block_sizes):
B, H, Tq, D = q.shape
Tkv = k.shape[2]
sm_scale = 1.0 / math.sqrt(D)
max_kv_blks = q2k_index.shape[-1]
assert Tq % 64 == 0, f"q length must be a multiple of 64, but got {Tq}"
assert Tkv % 64 == 0, f"kv length must be a multiple of 64, but got {Tkv}"
assert q2k_num.shape[
-1] == Tq // 64, f"shape mismatch, Tq // 64 = {Tq // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
assert variable_block_sizes.numel() == Tkv // 64, (
f"shape mismatch, variable_block_sizes must have length {Tkv // 64}, "
f"got {variable_block_sizes.numel()}")
o = torch.empty_like(q)
M = torch.empty((B, H, Tq), dtype=torch.float32, device=q.device)
grid = lambda _: (triton.cdiv(Tq, 64), B * H, 1)
_attn_fwd_sparse[grid](q,
k,
v,
sm_scale,
q2k_index,
q2k_num,
max_kv_blks,
variable_block_sizes,
M,
o,
q.stride(0),
q.stride(1),
q.stride(2),
q.stride(3),
k.stride(0),
k.stride(1),
k.stride(2),
k.stride(3),
v.stride(0),
v.stride(1),
v.stride(2),
v.stride(3),
o.stride(0),
o.stride(1),
o.stride(2),
o.stride(3),
B,
H,
Tq,
Tkv,
HEAD_DIM=D,
STAGE=3)
return o, M
# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> VENDORED FROM UPSTREAM (Apache-2.0) -- BLOCK 1/2 END >>>>>>>>>>>>>>>>>>>>
# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> VENDORED FROM UPSTREAM (Apache-2.0) -- BLOCK 2/2 BEGIN >>>>>>>>>>>>>>>>>>>>
# Source: hao-ai-lab/fastvideo `fastvideo_kernel` 0.3.5, file
# python/fastvideo_kernel/triton_kernels/index.py, lines 32-61 and 106-144.
#
# Why vendored: same install blocker as block 1/2, and this helper is pure Triton with no
# dependency on the compiled `_C` extension.
#
# What it replaces: our mask -> key-list conversion used `mask.float().topk(int(counts.max()))`,
# whose `int(...)` forces a GPU->CPU sync on every layer of every step and materializes an fp32
# copy of the [B,H,n,n] mask. `map_to_index` compacts each row on-device
# (grid = (bs, h, num_q_blocks)), pads with -1, and returns the per-row counts as an int32 tensor
# -- no host sync, no fp32 temp.
#
# Semantics relied on by the caller: for each query block, the first `index_num` entries of
# `index` are its selected kv-block ids in ASCENDING order; the rest are -1 and are never read
# (both the official forward kernel above and our backward loop bound their iteration by it).
#
# Deviation from upstream: its `import torch / triton / triton.language` are hoisted to this
# module's top, and its unused `from typing import Tuple` (kept upstream only for parity) is
# dropped. Between the BEGIN/END markers the block is otherwise byte-identical.
@triton.jit
def map_to_index_kernel(
map_ptr,
index_ptr,
index_num_ptr,
map_bs_stride,
map_h_stride,
map_q_stride,
map_kv_stride,
index_bs_stride,
index_h_stride,
index_q_stride,
index_kv_stride,
index_num_bs_stride,
index_num_h_stride,
index_num_q_stride,
num_kv_blocks,
):
b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
num = 0
for i in tl.range(num_kv_blocks):
map_entry = tl.load(map_ptr_base + i * map_kv_stride)
if map_entry:
tl.store(index_ptr_base + num * index_kv_stride, i)
num += 1
tl.store(index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride + q * index_num_q_stride, num)
def map_to_index(block_map: torch.Tensor):
"""
Convert a block map to indices and counts.
Args:
block_map: [bs, h, num_q_blocks, num_kv_blocks]
The block map tensor.
Returns:
index: [bs, h, num_q_blocks, num_kv_blocks]
The indices of the blocks.
index_num: [bs, h, num_q_blocks]
The number of blocks for each q block.
"""
bs, h, num_q_blocks, num_kv_blocks = block_map.shape
index = torch.full((block_map.shape), -1, dtype=torch.int32, device=block_map.device)
index_num = torch.empty((bs, h, num_q_blocks), dtype=torch.int32, device=block_map.device)
grid = (bs, h, num_q_blocks)
map_to_index_kernel[grid](
block_map,
index,
index_num,
block_map.stride(0),
block_map.stride(1),
block_map.stride(2),
block_map.stride(3),
index.stride(0),
index.stride(1),
index.stride(2),
index.stride(3),
index_num.stride(0),
index_num.stride(1),
index_num.stride(2),
num_kv_blocks=num_kv_blocks,
)
return index, index_num
# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> VENDORED FROM UPSTREAM (Apache-2.0) -- BLOCK 2/2 END >>>>>>>>>>>>>>>>>>>>
@triton.jit
def _attn_bwd_preprocess_varlen(
O, DO, DELTAS, QSIZE,
L_PAD,
D: tl.constexpr,
BLOCK_M: tl.constexpr,
):
# delta[row] = sum_d o[row,d] * do[row,d], per query tile, over the tile-order buffer.
idx_m = tl.program_id(0).to(tl.int64)
idx_bh = tl.program_id(1).to(tl.int64)
qkv_offset = idx_bh * L_PAD * D
offs_m = tl.arange(0, BLOCK_M)
offs_d = tl.arange(0, D)
qsize = tl.load(QSIZE + idx_m)
q_mask = offs_m < qsize
ptrs = qkv_offset + idx_m * BLOCK_M * D + offs_m[:, None] * D + offs_d[None, :]
o = tl.load(O + ptrs, mask=q_mask[:, None], other=0.0)
do = tl.load(DO + ptrs, mask=q_mask[:, None], other=0.0)
delta = tl.sum(o.to(tl.float32) * do.to(tl.float32), axis=1)
tl.store(DELTAS + idx_bh * L_PAD + idx_m * BLOCK_M + offs_m, delta, mask=q_mask)
@triton.jit
def _attn_bwd_dq_varlen(
Q, K, V, LSE, DELTAS, DO, DQ, LUT, LUTN,
QSIZE, KSIZE,
qk_scale,
MAXK,
N_TILES,
L_PAD,
D: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
):
# dQ for one query tile: loop over the LUTN[row] key tiles that tile actually selected (ids in
# LUT, ascending, -1 beyond the count) and recompute p from LSE. The bound is a runtime per-row
# value, so a row that selected few tiles does no extra work and never touches the -1 pads.
idx_m = tl.program_id(0).to(tl.int64)
idx_bh = tl.program_id(1).to(tl.int64)
qkv_offset = idx_bh * L_PAD * D
lse_offset = idx_bh * L_PAD
offs_m = tl.arange(0, BLOCK_M)
offs_n = tl.arange(0, BLOCK_N)
offs_d = tl.arange(0, D)
qsize = tl.load(QSIZE + idx_m)
q_mask = offs_m < qsize
m_ptrs = qkv_offset + idx_m * BLOCK_M * D + offs_m[:, None] * D + offs_d[None, :]
q = tl.load(Q + m_ptrs, mask=q_mask[:, None], other=0.0)
do = tl.load(DO + m_ptrs, mask=q_mask[:, None], other=0.0)
delta = tl.load(DELTAS + lse_offset + idx_m * BLOCK_M + offs_m, mask=q_mask, other=0.0)
lse = tl.load(LSE + lse_offset + idx_m * BLOCK_M + offs_m, mask=q_mask, other=float('inf'))
dq = tl.zeros([BLOCK_M, D], dtype=tl.float32)
lut_base = LUT + (idx_bh * N_TILES + idx_m) * MAXK
kv_blocks = tl.load(LUTN + idx_bh * N_TILES + idx_m)
for i in tl.range(0, kv_blocks):
idx_n = tl.load(lut_base + i)
ksize = tl.load(KSIZE + idx_n)
n_mask = offs_n < ksize
n_ptrs = qkv_offset + idx_n * BLOCK_N * D + offs_n[:, None] * D + offs_d[None, :]
k = tl.load(K + n_ptrs, mask=n_mask[:, None], other=0.0) # [BLOCK_N, D]
v = tl.load(V + n_ptrs, mask=n_mask[:, None], other=0.0) # [BLOCK_N, D]
qk = tl.dot(q, k.T) * (qk_scale * 1.4426950408889634) # [BLOCK_M, BLOCK_N]
p = tl.math.exp2(qk - lse[:, None])
p = tl.where(n_mask[None, :], p, 0.0)
dp = tl.dot(do, v.T).to(tl.float32) # [BLOCK_M, BLOCK_N]
ds = p * (dp - delta[:, None])
dq += tl.dot(ds.to(k.dtype), k) # [BLOCK_M, D]
tl.store(DQ + m_ptrs, (dq * qk_scale).to(DQ.type.element_ty), mask=q_mask[:, None])
@triton.jit
def _attn_bwd_dkdv_varlen(
Q, K, V, DO, DK, DV, KBID, LSE, DELTAS,
QSIZE, KSIZE,
qk_scale,
N_TILES,
L_PAD,
D: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
):
# dK/dV for one key tile: loop over ALL query tiles, keep those whose KBID incidence is 1.
idx_n = tl.program_id(0).to(tl.int64)
idx_bh = tl.program_id(1).to(tl.int64)
qkv_offset = idx_bh * L_PAD * D
lse_offset = idx_bh * L_PAD
offs_m = tl.arange(0, BLOCK_M)
offs_n = tl.arange(0, BLOCK_N)
offs_d = tl.arange(0, D)
ksize = tl.load(KSIZE + idx_n)
n_mask = offs_n < ksize
n_ptrs = qkv_offset + idx_n * BLOCK_N * D + offs_n[:, None] * D + offs_d[None, :]
k = tl.load(K + n_ptrs, mask=n_mask[:, None], other=0.0) # [BLOCK_N, D]
v = tl.load(V + n_ptrs, mask=n_mask[:, None], other=0.0) # [BLOCK_N, D]
dk = tl.zeros([BLOCK_N, D], dtype=tl.float32)
dv = tl.zeros([BLOCK_N, D], dtype=tl.float32)
kbid_base = KBID + idx_bh * N_TILES * N_TILES + idx_n
for idx_m in tl.range(0, N_TILES):
kbid = tl.load(kbid_base + idx_m * N_TILES)
if kbid == 1:
qsize = tl.load(QSIZE + idx_m)
m_mask = offs_m < qsize
m_ptrs = qkv_offset + idx_m * BLOCK_M * D + offs_m[:, None] * D + offs_d[None, :]
q = tl.load(Q + m_ptrs, mask=m_mask[:, None], other=0.0) # [BLOCK_M, D]
do = tl.load(DO + m_ptrs, mask=m_mask[:, None], other=0.0)
lse = tl.load(LSE + lse_offset + idx_m * BLOCK_M + offs_m, mask=m_mask, other=float('inf'))
delta = tl.load(DELTAS + lse_offset + idx_m * BLOCK_M + offs_m, mask=m_mask, other=0.0)
qkT = tl.dot(k, q.T) * (qk_scale * 1.4426950408889634) # [BLOCK_N, BLOCK_M]
pT = tl.math.exp2(qkT - lse[None, :])
pT = tl.where(m_mask[None, :], pT, 0.0)
dv += tl.dot(pT.to(do.dtype), do) # [BLOCK_N, D]
dpT = tl.dot(v, do.T).to(tl.float32) # [BLOCK_N, BLOCK_M]
dsT = pT * (dpT - delta[None, :])
dk += tl.dot(dsT.to(q.dtype), q) # [BLOCK_N, D]
tl.store(DK + n_ptrs, (dk * qk_scale).to(DK.type.element_ty), mask=n_mask[:, None])
tl.store(DV + n_ptrs, dv.to(DV.type.element_ty), mask=n_mask[:, None])
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.
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,
block_sizes, kb_id, lse, out)
ctx.qk_scale = qk_scale
return out
@staticmethod
def backward(ctx, do):
(q_tile, k_tile, v_tile, q2k_idx, q2k_num,
block_sizes, kb_id, lse, out) = ctx.saved_tensors
do = do.contiguous()
B, H, L_BUF, D = q_tile.shape
num_tiles = L_BUF // 64
MAXK = q2k_idx.shape[-1]
dq = torch.zeros_like(q_tile)
dk = torch.zeros_like(k_tile)
dv = torch.zeros_like(v_tile)
delta = torch.empty((B, H, L_BUF), device=q_tile.device, dtype=torch.float32)
grid = (num_tiles, B * H)
_attn_bwd_preprocess_varlen[grid](
out, do, delta, block_sizes, L_BUF, D, 64,
num_warps=8, num_stages=3,
)
_attn_bwd_dq_varlen[grid](
q_tile, k_tile, v_tile, lse, delta, do, dq, q2k_idx, q2k_num,
block_sizes, block_sizes, ctx.qk_scale,
MAXK, num_tiles, L_BUF, D, 64, 64,
num_warps=8, num_stages=3,
)
_attn_bwd_dkdv_varlen[grid](
q_tile, k_tile, v_tile, do, dk, dv, kb_id, lse, delta,
block_sizes, block_sizes, ctx.qk_scale,
num_tiles, L_BUF, D, 64, 64,
num_warps=8, num_stages=3,
)
return dq, dk, dv, None, None, None, None, None
def sparse_linear_attention_varlen(q_tile, k_tile, v_tile, block_mask, block_sizes,
qk_scale=None):
"""Tile-order block-sparse attention for the H3/Wan VSA path (forward + backward).
q_tile/k_tile/v_tile: [B, H, L_BUF, D] contiguous, rows in tile order, pad slots zero.
block_mask: [B, H, num_tiles, num_tiles] bool query->key incidence at the PHYSICAL 64-token
block granularity the kernels run at (a logical 128/256-tile selection is expanded down to
it by ``_h3_expand_mask_and_sizes_to_64``). It drives both stages: map_to_index compacts it
into the per-query key list used by the forward and by dQ, and it is the dK/dV reverse-loop
incidence. block_sizes: [num_tiles] int valid token count per 64-token block; 0 is legal
(a partial logical tile's trailing sub-blocks) and makes that block contribute nothing.
num_tiles is derived as L_BUF // 64, so the buffer carries no separate dummy tile.
Returns a NEW [B, H, L_BUF, D] output buffer (pad rows are garbage; the caller untiles only the
real rows). Differentiable -- gradients flow to q_tile/k_tile/v_tile.
"""
assert q_tile.is_contiguous() and k_tile.is_contiguous() and v_tile.is_contiguous()
B, H, L_BUF, D = q_tile.shape
assert L_BUF > 0 and L_BUF % 64 == 0, "VSA tile 缓冲长度必须是正的 64 倍数。"
assert k_tile.shape == q_tile.shape and v_tile.shape == q_tile.shape
num_tiles = L_BUF // 64
assert tuple(block_mask.shape) == (B, H, num_tiles, num_tiles), \
f"block_mask {tuple(block_mask.shape)} vs expected {(B, H, num_tiles, num_tiles)}"
assert block_sizes.numel() == num_tiles, \
f"block_sizes has {block_sizes.numel()} entries, expected {num_tiles}"
default_scale = D ** -0.5
if qk_scale is not None and not math.isclose(qk_scale, default_scale, rel_tol=1e-7):
raise ValueError("VSA 前向核仅支持 head_dim**-0.5；不接受自定义 qk_scale。")
qk_scale = default_scale
# On-device compaction by the vendored official Triton kernel: no host sync (the topk over
# int(counts.max()) it replaces synced once per layer per step). `q2k_idx` keeps the full
# num_tiles width and pads with -1; both kernels bound their loops by q2k_num.
q2k_idx, q2k_num = map_to_index(block_mask.contiguous())
block_sizes = block_sizes.to(torch.int32).contiguous()
kb_id = block_mask.to(torch.int8).contiguous()
return _sparse_linear_attention_varlen.apply(
q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes, kb_id, qk_scale)
# ======================================================================================
# SECTION B -- fixed-length SLA original (NOT varlen): mean_pool / get_block_map / _attn_fwd /
# _attn_bwd_{preprocess,dq,dkdv} / class _sparse_linear_attention. Do not confuse these with
# section A's `*_varlen` kernels; nothing in section A calls into this section.
# ======================================================================================
@triton.jit
def compress_kernel(
videox_fun/models/attention_utils.py
+499 -2
View File
@@ -1,3 +1,5 @@
import functools
import math
import os
import warnings
@@ -42,7 +44,8 @@ except Exception:
SAGE_ATTENTION_AVAILABLE = False
from .attention_kernel import _sparse_linear_attention, get_block_map
from .attention_kernel import (_sparse_linear_attention, get_block_map,
mean_pool, sparse_linear_attention_varlen)
def convert_qkv_dtype(q, k, v):
@@ -550,4 +553,498 @@ class RoLASparseAttention(nn.Module):
[B, H, L, 1] gate in (0, 1).
'''
gb = self.gate_bias if head_slice is None else self.gate_bias[:, head_slice]
return torch.sigmoid(self.gate_proj(x_bhld) + gb)
return torch.sigmoid(self.gate_proj(x_bhld) + gb)
# Logical tile geometries keyed by tile volume -- the official VSA_H3_TILE_SHAPES. 64 is what the
# official defaults actually select for the checkpoints this repo runs: the Wan backend's
# VSA_TILE_SIZE is (4,4,4), and both released FastH3 inference contracts pin "vsa_tile_size": 64.
# (256/(4,8,8) is the H3 backend's *code* default, aimed at the Blackwell CuTe fastpath.) 128 and
# 256 stay selectable and reach the same 64-block kernel through the route-A expansion below.
_H3_TILE_SHAPES = {64: (4, 4, 4), 128: (4, 4, 8), 256: (4, 8, 8)}
_H3_TILE_ELEMS = 64
# Physical block size the Triton kernels are instantiated at: a logical tile of N elems is N // 64
# CONTIGUOUS 64-blocks in the padded tile buffer.
_H3_KERNEL_BLOCK = 64
def _h3_compute_topk(sparsity, num_blocks):
"""Key tiles to KEEP for a sparsity level, clamped to [1, num_blocks]."""
return max(1, min(math.ceil((1.0 - sparsity) * num_blocks), num_blocks))
@functools.lru_cache(maxsize=32)
def _h3_tile_partition_indices(grid, tile_shape, device):
"""Region tile-order -> original flat region index (port of get_tile_partition_indices)."""
T, H, W = grid
ts, hs, ws = tile_shape
indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W)
parts = []
for t in range(math.ceil(T / ts)):
for h in range(math.ceil(H / hs)):
for w in range(math.ceil(W / ws)):
parts.append(indices[t * ts:min(t * ts + ts, T),
h * hs:min(h * hs + hs, H),
w * ws:min(w * ws + ws, W)].flatten())
return torch.cat(parts, dim=0)
@functools.lru_cache(maxsize=32)
def _h3_variable_block_sizes(grid, num_tiles, device, tile_shape):
"""Valid (non-pad) token count inside every 3-D tile (port of construct_variable_block_sizes)."""
t, h, w = grid
ts_t, ts_h, ts_w = tile_shape
n_t, n_h, n_w = num_tiles
def _sizes(dim_len, tile, n):
sizes = torch.full((n,), tile, dtype=torch.long, device=device)
rem = dim_len - (n - 1) * tile
sizes[-1] = rem if rem > 0 else tile
return sizes
t_s = _sizes(t, ts_t, n_t)
h_s = _sizes(h, ts_h, n_h)
w_s = _sizes(w, ts_w, n_w)
return (t_s[:, None, None] * h_s[None, :, None] * w_s[None, None, :]).reshape(-1)
def _h3_non_pad_index(variable_block_sizes, tile_elems):
"""Padded-buffer slot index of every real row, in tile order (port of get_non_pad_index)."""
n_win = variable_block_sizes.shape[0]
device = variable_block_sizes.device
starts = torch.arange(n_win, device=device) * tile_elems
index_pad = starts[:, None] + torch.arange(tile_elems, device=device)[None, :]
index_mask = torch.arange(tile_elems, device=device)[None, :] < variable_block_sizes[:, None]
return index_pad[index_mask]
def _h3_validate_segment_geometry(segments, variable_block_sizes, untile_combined_index, tile_elems):
R'''Fail synchronously on out-of-bounds tile geometry (port of the official guard).
Invariants the block-sparse kernels trust without checking: every tile's valid size is in
(0, tile_elems]; the sizes sum to the packed sequence length; and ``untile_combined_index``
maps each packed row to exactly one non-pad slot of the padded tile buffer. A violation would
otherwise surface only as an async device fault at some later kernel or collective (e.g. an
FSDP all-gather), which is unattributable -- so raise here, once per cached geometry.
'''
total = sum(math.prod(seg) if isinstance(seg, tuple) else seg for seg in segments)
n_pad = variable_block_sizes.numel() * tile_elems
sizes_min = int(variable_block_sizes.min())
sizes_max = int(variable_block_sizes.max())
sizes_sum = int(variable_block_sizes.sum())
if sizes_min < 1 or sizes_max > tile_elems or sizes_sum != total:
raise ValueError(
f"VSA-H3 tile sizes out of bounds for segments={segments}, tile_elems={tile_elems}: "
f"min={sizes_min}, max={sizes_max}, sum={sizes_sum}, expected sum={total}.")
if untile_combined_index.numel() != total:
raise ValueError(
f"VSA-H3 untile index has {untile_combined_index.numel()} entries for a packed "
f"sequence of {total} rows (segments={segments}).")
idx_min = int(untile_combined_index.min())
idx_max = int(untile_combined_index.max())
if idx_min < 0 or idx_max >= n_pad:
# Range first: the pad-slot gather below would itself index out of bounds (the very async
# fault this guard exists to preempt).
raise ValueError(
f"VSA-H3 untile index is not an injective map into non-pad slots: range "
f"[{idx_min}, {idx_max}] vs padded length {n_pad} (segments={segments}).")
in_tile_offset = untile_combined_index % tile_elems
maps_into_pad = bool(
(in_tile_offset >= variable_block_sizes[untile_combined_index // tile_elems]).any())
if maps_into_pad or int(torch.unique(untile_combined_index).numel()) != total:
raise ValueError(
f"VSA-H3 untile index is not an injective map into non-pad slots: "
f"pad-slot hit={maps_into_pad} (segments={segments}).")
@functools.lru_cache(maxsize=32)
def build_h3_tile_geometry(segments, device, tile_shape=_H3_TILE_SHAPES[_H3_TILE_ELEMS]):
R'''Tile a packed sequence of interleaved dense segments and video regions.
Port of the official ``_h3_segment_tile_geometry``. ``segments`` lists, in packed
order, dense segment row counts (``int``) and video-region token grids (``(t,h,w)``
triples). Returns ``(tile_partition_indices, variable_block_sizes,
untile_combined_index, num_prefix_tiles, num_video_tiles, video_tile_spans,
tile_elems)`` where ``video_tile_spans`` holds one ``[start, end)`` tile range per
video region (prefix tiles included in the numbering).
'''
tile_elems = math.prod(tile_shape)
ts_t, ts_h, ts_w = tile_shape
prefix_sizes, prefix_rows, video_entries = [], [], []
cursor = 0
for seg in segments:
if isinstance(seg, tuple):
video_entries.append((cursor, seg))
cursor += math.prod(seg)
continue
if seg <= 0:
continue
full, rem = divmod(seg, tile_elems)
prefix_sizes.extend([tile_elems] * full)
if rem:
prefix_sizes.append(rem)
prefix_rows.append(torch.arange(cursor, cursor + seg, device=device, dtype=torch.long))
cursor += seg
num_prefix_tiles = len(prefix_sizes)
video_index_parts, video_size_parts, video_tile_spans = [], [], []
tile_cursor = num_prefix_tiles
for start, grid in video_entries:
t, h, w = grid
num_tiles = (math.ceil(t / ts_t), math.ceil(h / ts_h), math.ceil(w / ts_w))
region_sizes = _h3_variable_block_sizes(grid, num_tiles, device, tile_shape)
video_index_parts.append(_h3_tile_partition_indices(grid, tile_shape, device) + start)
video_size_parts.append(region_sizes)
video_tile_spans.append((tile_cursor, tile_cursor + int(region_sizes.numel())))
tile_cursor += int(region_sizes.numel())
num_video_tiles = tile_cursor - num_prefix_tiles
tile_partition_indices = torch.cat(prefix_rows + video_index_parts)
variable_block_sizes = torch.cat(
[torch.tensor(prefix_sizes, dtype=torch.long, device=device)] + video_size_parts)
non_pad_index = _h3_non_pad_index(variable_block_sizes, tile_elems)
untile_combined_index = non_pad_index[torch.argsort(tile_partition_indices)]
# One-time (lru-cached) synchronous bounds check; see _h3_validate_segment_geometry.
_h3_validate_segment_geometry(segments, variable_block_sizes, untile_combined_index, tile_elems)
return (tile_partition_indices, variable_block_sizes, untile_combined_index,
num_prefix_tiles, num_video_tiles, tuple(video_tile_spans), tile_elems)
# Per-process reusable padded tile scratch for forward_h3's no-grad (inference) path.
# Mirrors the official FastVideo builder-owned tile_buf_holder: one buffer per
# (role, shape, dtype, device) shared across all layers and steps, so the per-layer
# torch.zeros((B,H,L_buf,D)) allocation+zeroing churn is paid once instead of every call.
# Keyed by role because q/k/v/gate are all live simultaneously within one forward_h3.
_H3_TILE_BUFS = {}
def _h3_tile_into_buf(t, dtype, idx, shape, device, role, reuse):
R'''Scatter packed [B,L,H,D] ``t`` into a [B,H,L_buf,D] tile buffer (pad/dummy stay zero).
Under ``reuse`` (no-grad inference) the buffer is cached in ``_H3_TILE_BUFS`` and only
re-zeroed when the tile geometry (``idx`` identity) changes, exactly like the official
holder: pad slots are never written and non-pad slots are fully overwritten per call, so
a reused buffer stays valid. Callers consume the result before the next layer overwrites
it. Under grad the buffer is fresh -- an in-place write into a buffer that outlives the
step would anchor that step's autograd graph (the official tile-cache OOM).
'''
B, H, L_buf, D = shape
if not reuse:
buf = torch.zeros(shape, device=device, dtype=dtype)
buf[:, :, idx, :] = t.transpose(1, 2).to(dtype)
return buf
key = (role, B, H, L_buf, D, dtype, str(device))
ent = _H3_TILE_BUFS.get(key)
if ent is None or ent[0].shape != shape or ent[0].dtype != dtype:
ent = [torch.zeros(shape, device=device, dtype=dtype), None]
_H3_TILE_BUFS[key] = ent
if ent[1] is not idx: # geometry identity changed -> clear stale valid rows once
ent[0].zero_()
ent[1] = idx
ent[0][:, :, idx, :] = t.transpose(1, 2).to(dtype)
return ent[0]
def _h3_pool_tiles(x, variable_block_sizes, tile_elems, n_tiles):
R'''fp32 masked mean over each tile. ``x``: [B,H,L_buf,D] (pad slots zero) -> [B,H,n_tiles,D].
Pad slots are guaranteed zero (zeros-init buffer, only non-pad slots written), so a
plain fp32 sum divided by the true tile size is exactly the masked mean.
'''
B, H, _, D = x.shape
xt = x[:, :, :n_tiles * tile_elems].reshape(B, H, n_tiles, tile_elems, D)
pooled = xt.sum(dim=3, dtype=torch.float32)
return pooled / variable_block_sizes.view(1, 1, -1, 1).to(torch.float32)
def _h3_build_block_mask(scores, num_prefix_tiles, sparsity, exempt,
video_tile_spans, span_sparsities):
R'''scores [B,H,n_tiles,n_tiles] -> bool mask (port of official ``_build_block_mask``).
Prefix rows/cols: with ``exempt`` every query keeps every prefix column, and every
query keeps its own top-k of EACH video region under that region's sparsity.
``sparsity <= 0`` keeps every tile (dense).
'''
n_tiles = scores.shape[-1]
if sparsity <= 0.0:
span_sparsities = tuple(sparsity for _ in video_tile_spans)
span_topk = [_h3_compute_topk(s, end - start)
for (start, end), s in zip(video_tile_spans, span_sparsities)]
if all(kk == end - start for kk, (start, end) in zip(span_topk, video_tile_spans)):
return torch.ones_like(scores, dtype=torch.bool)
mask = torch.zeros_like(scores, dtype=torch.bool)
if exempt or num_prefix_tiles == 0:
for (start, end), k_span in zip(video_tile_spans, span_topk):
if k_span == end - start:
mask[..., start:end] = True
continue
idx = scores[..., start:end].topk(k_span, dim=-1).indices + start
mask.scatter_(-1, idx, True)
mask[..., :num_prefix_tiles] = True
else:
k_total = min(span_topk[0] + num_prefix_tiles, n_tiles)
idx = scores.topk(k_total, dim=-1).indices
mask.scatter_(-1, idx, True)
mask[:, :, :num_prefix_tiles, :] = True
return mask
def _h3_expand_mask_and_sizes_to_64(block_mask, variable_block_sizes, tile_elems):
R'''Logical tile selection -> the physical 64-block map the Triton kernels run at.
Port of the official ``_expand_mask_and_sizes_256_to_64`` / ``_128_to_64`` (their "route A"),
generalized over ``r = tile_elems // 64``. A logical tile occupies ``r`` CONTIGUOUS 64-blocks in
the padded buffer, so selecting it selects all r sub-blocks on both the query and the key axis
(``repeat_interleave``), and its valid-token count chops into 64-token strides. The trailing
sub-blocks of a partial logical tile legitimately get size 0, which every kernel treats as
"contributes nothing" (the forward masks those columns to -inf, dQ skips them, dK/dV never
stores them). ``tile_elems == 64`` is the identity.
'''
r = tile_elems // _H3_KERNEL_BLOCK
if r == 1:
return block_mask, variable_block_sizes
mask_64 = block_mask.repeat_interleave(r, dim=2).repeat_interleave(r, dim=3)
sizes = variable_block_sizes.to(torch.int32)
offsets = torch.arange(r, device=sizes.device, dtype=torch.int32) * _H3_KERNEL_BLOCK
sizes_64 = torch.clamp(sizes[:, None] - offsets[None, :],
min=0, max=_H3_KERNEL_BLOCK).reshape(-1)
return mask_64, sizes_64
class VSASparseAttention(nn.Module):
R'''
Video Sparse Attention (VSA) -- the FastVideo/hao-ai-lab coarse-to-fine block-sparse
attention (arXiv:2505.13389) used by the official FastWan2.1 sparse-distilled students.
Two branches, mirroring the released FastWan ``to_gate_compress`` design:
- sparse branch : mean-pool q/k into blocks -> coarse block scores -> keep the top-k
key blocks per query block -> block-sparse EXACT softmax over the
selected blocks.
- compressed branch: a global mean-pooled ("compressed") attention over ALL key blocks,
broadcast back to token resolution and scaled per token/head by the
learned gate ``to_gate_compress(x)`` (x = attention-input hidden
states). This is FastVideo's learned compensation for the context
lost to sparsification. ``to_gate_compress`` is zero-init, so a
checkpoint without those weights reduces to sparse-only::
out = o_sparse + to_gate_compress(x) * o_compressed
New learnable params per layer (all under ``attn_op.*``): ``to_gate_compress`` weight+bias.
Absent from a dense checkpoint (loaded with strict=False); the official FastWan release
provides them.
NOTE on fidelity: the block-sparse forward IS the official FastVideo Triton kernel (vendored),
the tile geometry / exempt mask / top-k are ports of the official H3 backend, and the tile
volume is selectable over the official {64, 128, 256} -- 128/256 reach the same 64-block kernel
through the official route-A expansion.
``forward`` accepts BSND (``[B, L, H, D]``) q/k/v like the sibling attention modules;
internal helpers use BHLD (``[B, H, L, *]``).
'''
def __init__(self, head_dim, num_heads, dim, topk=0.2,
BLKQ=64, BLKK=64, use_bf16=True, tile_elems=_H3_TILE_ELEMS):
R'''
Args:
head_dim: dimension of each head.
num_heads: number of attention heads.
dim: hidden size of the attention input (sizes ``to_gate_compress``).
topk: keep-ratio of key blocks for the sparse branch (0.2 -> 80% sparse,
matching the official ``--VSA-sparsity 0.8``).
BLKQ/BLKK: query/key block sizes for the flat (non-segment) sparse kernel.
use_bf16: compute the sparse kernel in bfloat16 (default) or float16.
tile_elems: logical tile volume for the segment-aware path -- one of the official
``VSA_H3_TILE_SHAPES`` keys (64 / 128 / 256). Must match the volume the
checkpoint was distilled at; the released FastWan and FastH3 both use 64.
'''
super().__init__()
self.head_dim = head_dim
self.num_heads = num_heads
self.dim = dim
self.topk = topk
self.BLKQ = BLKQ
self.BLKK = BLKK
self.dtype = torch.bfloat16 if use_bf16 else torch.float16
if tile_elems not in _H3_TILE_SHAPES:
raise ValueError(
f"Unsupported VSA tile volume {tile_elems}; supported volumes are "
f"{sorted(_H3_TILE_SHAPES)} (the official VSA_H3_TILE_SHAPES).")
self.tile_elems = tile_elems
self.tile_shape = _H3_TILE_SHAPES[tile_elems]
inner = num_heads * head_dim
# Learned gate for the compressed global branch (FastVideo ``to_gate_compress``).
self.to_gate_compress = nn.Linear(dim, inner, bias=True)
self.init_weights_()
def init_weights_(self):
with torch.no_grad():
nn.init.zeros_(self.to_gate_compress.weight)
nn.init.zeros_(self.to_gate_compress.bias)
def sparse(self, q, k, v):
R'''Block-sparse exact softmax (VSA fine stage).
Args:
q, k: [B, H, L, D] POST-RoPE queries/keys.
v: [B, H, L, D] values.
Returns:
[B, H, L, D], or ``None`` when the sequence is shorter than one block.
'''
L = q.shape[2]
if L < max(self.BLKQ, self.BLKK):
return None
orig_dtype = q.dtype
with torch.no_grad():
pq = self._pool_flat(q, self.BLKQ)
pk = self._pool_flat(k, self.BLKK)
scores = pq @ pk.transpose(-1, -2)
real_topk = max(1, min(scores.shape[-1], math.ceil(self.topk * scores.shape[-1])))
lut = scores.topk(real_topk, dim=-1, sorted=False).indices.contiguous()
sparse_map = torch.zeros_like(scores, dtype=torch.bool)
sparse_map.scatter_(-1, lut, True)
if self.BLKQ % 64 or self.BLKK % 64:
raise ValueError("VSA 的 flat block 大小必须是 64 的正整数倍。")
B, H, L, D = q.shape
n_tiles = math.ceil(L / 64)
padding = n_tiles * 64 - L
q, k, v = (F.pad(t.to(self.dtype), (0, 0, 0, padding)).contiguous() for t in (q, k, v))
sizes = (L - torch.arange(n_tiles, device=q.device) * 64).clamp(max=64).to(torch.int32)
mask = sparse_map.repeat_interleave(self.BLKQ // 64, dim=-2).repeat_interleave(
self.BLKK // 64, dim=-1)[..., :n_tiles, :n_tiles].contiguous()
o_s = sparse_linear_attention_varlen(q, k, v, mask, sizes)
return o_s[:, :, :L].to(orig_dtype)
@staticmethod
def _pool_flat(x, block_size):
B, H, L, D = x.shape
n_blocks = math.ceil(L / block_size)
padded = F.pad(x.float(), (0, 0, 0, n_blocks * block_size - L))
counts = (L - torch.arange(n_blocks, device=x.device) * block_size).clamp(max=block_size)
return padded.reshape(B, H, n_blocks, block_size, D).sum(3) / counts.view(1, 1, -1, 1)
def compressed(self, q, k, v):
R'''Mean-pooled global ("compressed") attention, broadcast to token resolution.
Args:
q, k, v: [B, H, L, D] POST-RoPE queries/keys/values.
Returns:
[B, H, L, D] compressed global context per token.
'''
B, H, L, D = q.shape
pq = self._pool_flat(q, self.BLKQ)
pk = self._pool_flat(k, self.BLKK)
pv = self._pool_flat(v, self.BLKK)
scale = D ** -0.5
scores = (pq.float() @ pk.float().transpose(-1, -2)) * scale
attn = torch.softmax(scores, dim=-1).to(pv.dtype)
oc_blocks = attn @ pv # [B,H,Lq,D]
# Broadcast each query block's context to its BLKQ tokens, then trim to L.
oc = oc_blocks.repeat_interleave(self.BLKQ, dim=2)[:, :, :L]
return oc
def forward_h3(self, q, k, v, x, segments, return_sparsity=False):
R'''Segment-aware H3 VSA -- the official FastVideo segment/tile mask semantics.
Args:
q, k, v: [B, L, H, D] POST-RoPE queries/keys/values in PACKED order.
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.
Returns:
[B, L, H, D] attention output (before the output projection).
Tiles the packed sequence into [prefix tiles][video cubic tiles] of ``self.tile_shape``,
builds the exempt block mask on pooled tile scores, expands that logical selection down to
the 64-block granularity the kernels run at, runs the block-sparse Triton kernel for the
fine branch, adds the gated compressed branch on the same logical geometry, then untiles
back to packed order.
'''
B, L, H, D = q.shape
device = q.device
(tile_partition_indices, variable_block_sizes, untile_combined_index,
num_prefix_tiles, num_video_tiles, video_tile_spans, tile_elems) = \
build_h3_tile_geometry(tuple(segments), device, self.tile_shape)
n_tiles = num_prefix_tiles + num_video_tiles
n_regions = len(video_tile_spans)
# Logical tiles are laid end to end with no spare: the kernels take a per-64-block valid
# count, so a partial tile needs no zero dummy partner tile any more.
L_buf = n_tiles * tile_elems
sparsity = 1.0 - self.topk # self.topk is the KEEP ratio
vbs = variable_block_sizes.to(device)
reuse = not torch.is_grad_enabled()
buf_shape = (B, H, L_buf, D)
def _tile(t, dtype, role):
# [B,L,H,D] packed -> [B,H,L_buf,D] tile order (pad slots stay zero); under no-grad the
# buffer is shared across layers/steps (see _h3_tile_into_buf).
return _h3_tile_into_buf(t, dtype, untile_combined_index, buf_shape,
device, role, reuse)
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)
gb = _tile(gate, gate.dtype, 'g')
pq = _h3_pool_tiles(qb, vbs, tile_elems, n_tiles) # [B,H,n_tiles,D] fp32
pk = _h3_pool_tiles(kb, vbs, tile_elems, n_tiles)
pv = _h3_pool_tiles(vb, vbs, tile_elems, n_tiles)
# Block selection is a non-differentiable constant (exactly like the flat path's
# get_block_map): build the mask under no_grad on detached pooled scores so no graph is
# retained for the top-k.
with torch.no_grad():
scores = torch.matmul(pq.detach(), pk.detach().transpose(-2, -1)) / (D ** 0.5)
mask = _h3_build_block_mask(scores, num_prefix_tiles, sparsity, True,
video_tile_spans, (sparsity,) * n_regions)
# Logical (tile_elems) selection -> the physical 64-block map + per-block counts.
# Identity when tile_elems == 64.
block_mask, block_sizes = _h3_expand_mask_and_sizes_to_64(mask, vbs, tile_elems)
o_s = sparse_linear_attention_varlen(
qb, kb, vb, block_mask, block_sizes, qk_scale=D ** -0.5)
# gated compressed branch on the same logical tile geometry
cs = torch.matmul(pq, pk.transpose(-2, -1)) * (D ** -0.5)
attn = torch.softmax(cs, dim=-1).to(pv.dtype)
oc = (attn @ pv).repeat_interleave(tile_elems, dim=2) # [B,H,L_buf,D]
out_buf = o_s.to(gb.dtype) + gb * oc.to(gb.dtype)
out = out_buf[:, :, untile_combined_index].transpose(1, 2) # [B,L,H,D] packed
out = out.to(q.dtype)
if return_sparsity:
return out, self.topk
return out
def forward(self, q, k, v, x, return_sparsity=False, segments=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.
Returns:
[B, L, H, D] attention output (before the output projection).
'''
if segments is not None:
# Segment-aware tile geometry with a differentiable variable-length block-sparse
# 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)
B, L, H, D = q.shape
qb = q.transpose(1, 2).contiguous() # [B,H,L,D]
kb = k.transpose(1, 2).contiguous()
vb = v.transpose(1, 2).contiguous()
o_s = self.sparse(qb, kb, vb)
if o_s is None:
# 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]
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]
if return_sparsity:
return out, self.topk
return out
videox_fun/models/fastvideo_minimax_h3_transformer3d.py
+256
View File
@@ -0,0 +1,256 @@
# FastVideo VSA (Video Sparse Attention) student backbone for MiniMax-H3. Mirrors the Wan-side
# integration in ``fastvideo_wan_transformer3d.py``: it swaps the dense softmax self-attention of
# ``MiniMaxH3Transformer3DModel`` for the coarse-to-fine block-sparse VSA kernel (reusing the
# in-repo SLA Triton kernel) plus FastVideo's learned ``to_gate_compress`` compressed-global
# compensation branch.
#
# MiniMax-H3 runs full self-attention over ONE packed multi-modal sequence (video + text + audio
# rows) through a swappable attention *processor*, so the VSA op is wired in at the processor
# level (``FastVideoMiniMaxH3AttnProcessor``) and the trainable ``attn_op`` gate is attached to
# each main block's ``MiniMaxH3Attention`` module. The token refiner keeps the dense processor.
#
# VSA reference: Zhang et al., "VSA: Faster Video Diffusion with Trainable Sparse Attention",
# arXiv:2505.13389 (FastVideo / hao-ai-lab).
import torch
import torch.nn.functional as F
from .attention_utils import VSASparseAttention, attention
from .minimax_h3_transformer3d import (MiniMaxH3Transformer3DModel,
apply_minimax_h3_rotary_emb)
class FastVideoMiniMaxH3AttnProcessor:
r"""MiniMax-H3 attention processor whose dense softmax kernel is swapped for VSA block-sparse
attention (sparse fine branch + gated mean-pooled compressed branch).
Mirrors ``MiniMaxH3AttnProcessor.__call__`` exactly, except the padless single-document path
``attention(q, k, v, causal=False)`` becomes the VSA op ``attn.attn_op(q, k, v, x)``. Two
cases keep the original DENSE attention:
- ``attention_mask is not None``: the packed sequence carries padding rows kept in their own
attention document; the VSA kernel does not accept an arbitrary mask, so fall back to the
masked SDPA path (identical to the base processor).
- ``attn.topk >= 1.0``: sparsity disabled -> the base dense full-softmax path.
"""
def __call__(
self,
attn,
hidden_states: torch.Tensor,
rotary_emb=None,
attention_mask=None,
valid_length=None,
vsa_segments=None,
) -> torch.Tensor:
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)
query = query.unflatten(-1, (attn.heads, -1))
key = key.unflatten(-1, (attn.heads, -1))
value = value.unflatten(-1, (attn.heads, -1))
query = attn.norm_q(query)
key = attn.norm_k(key)
if rotary_emb is not None:
query = apply_minimax_h3_rotary_emb(query, *rotary_emb)
key = apply_minimax_h3_rotary_emb(key, *rotary_emb)
# VSA on the padless single-document path only. query/key/value are [B, L, H, D] POST-RoPE
# and hidden_states is the [B, L, C] attention input that drives the compress gate; the op
# returns [B, L, H, D], matching the dense paths below so the common tail is unchanged.
if attention_mask is None and getattr(attn, "topk", 1.0) < 1.0:
hidden_states = attn.attn_op(
query, key, value, hidden_states, segments=vsa_segments)
elif attention_mask is None:
hidden_states = attention(query, key, value, causal=False)
else:
hidden_states = F.scaled_dot_product_attention(
query.transpose(1, 2),
key.transpose(1, 2),
value.transpose(1, 2),
attn_mask=attention_mask[None, None],
).transpose(1, 2)
hidden_states = hidden_states.flatten(2, 3).type_as(query)
hidden_states = attn.to_out[0](hidden_states)
hidden_states = attn.to_out[1](hidden_states)
return hidden_states
class FastVideoMiniMaxH3Transformer3DModel(MiniMaxH3Transformer3DModel):
r"""MiniMax-H3 transformer (joint video + audio) with a VSA sparse self-attention on the main
block stack. Block type stays ``MiniMaxH3TransformerBlock``, so the existing
``--fsdp_transformer_layer_cls_to_wrap=MiniMaxH3TransformerBlock`` policy keeps working.
"""
# NOTE: diffusers from_pretrained filters config.json keys by __init__ signature, so every
# parent parameter must be declared explicitly (a *args/**kwargs passthrough would silently
# drop hidden_size/num_layers/... and build a wrongly-sized model).
def __init__(
self,
num_attention_heads: int = 56,
attention_head_dim: int = 128,
hidden_size: int = 5376,
num_layers: int = 50,
num_refiner_layers: int = 2,
ffn_dim: int = 14336,
in_channels: int = 24,
audio_in_channels: int = 32,
patch_size=(1, 2, 2),
text_dim: int = 5120,
freq_dim: int = 256,
time_embed_hidden_dim: int = 5376,
time_embed_dim: int = 2688,
rope_freq_dim: int = 16,
rope_theta: float = 10000.0,
norm_eps: float = 1e-5,
qk_norm_eps: float = 1e-5,
final_norm_eps: float = 1e-5,
student_topk=0.2,
):
super().__init__(
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
hidden_size=hidden_size,
num_layers=num_layers,
num_refiner_layers=num_refiner_layers,
ffn_dim=ffn_dim,
in_channels=in_channels,
audio_in_channels=audio_in_channels,
patch_size=patch_size,
text_dim=text_dim,
freq_dim=freq_dim,
time_embed_hidden_dim=time_embed_hidden_dim,
time_embed_dim=time_embed_dim,
rope_freq_dim=rope_freq_dim,
rope_theta=rope_theta,
norm_eps=norm_eps,
qk_norm_eps=qk_norm_eps,
final_norm_eps=final_norm_eps,
)
self.student_topk = student_topk
self.register_to_config(student_topk=student_topk)
self._apply_sparse_self_attention(student_topk)
def _apply_sparse_self_attention(self, topk):
# Sparsify only the main block stack (the heavy video+audio attention); the text token
# refiner keeps the dense processor. The VSA op is attached to the existing attention
# module (an nn.Module) so its trainable gate lands in the model's parameter tree, and the
# processor is swapped to route the padless path through it.
for block in self.transformer_blocks:
attn = block.attn
attn.topk = topk
attn.set_processor(FastVideoMiniMaxH3AttnProcessor())
if topk >= 1.0:
# Dense fallback: drop any residual VSA submodule (e.g. one left on meta by an
# earlier low_cpu_mem_usage __init__) so no stale meta params reach the offloader.
if hasattr(attn, "attn_op"):
del attn.attn_op
continue
# ``dim`` sizes to_gate_compress's input, which is the attention-input hidden states
# (hidden_size = to_q.in_features), NOT inner_dim (heads * head_dim).
attn_op = VSASparseAttention(
attn.head_dim, attn.heads, attn.to_q.in_features, topk=topk)
# Under low_cpu_mem_usage from_pretrained the backbone params are meta tensors and the
# freshly built gate is materialized by a later set_student_topk (called after
# from_pretrained) or the checkpoint load; only align device/dtype when concrete, and
# match the backbone dtype so bf16 weights and an fp32 gate never mix under FSDP.
p0 = attn.to_q.weight
if not p0.is_meta:
attn_op = attn_op.to(device=p0.device, dtype=p0.dtype)
attn.attn_op = attn_op
def set_student_topk(self, topk):
self.student_topk = topk
self.register_to_config(student_topk=topk)
if topk < 1.0 and all(hasattr(block.attn, "attn_op") and
not any(p.is_meta for p in block.attn.attn_op.parameters())
for block in self.transformer_blocks):
for block in self.transformer_blocks:
block.attn.topk = topk
block.attn.attn_op.topk = topk
else:
self._apply_sparse_self_attention(topk)
def forward(self, hidden_states, audio_hidden_states, encoder_hidden_states,
timestep, timestep_indices, token_tags, position_ids,
video_indices, audio_indices, text_indices,
attention_kwargs=None, return_dict=True):
attention_kwargs = dict(attention_kwargs or {})
attention_kwargs["vsa_segments"] = self._compute_h3_segments(token_tags, position_ids)
return super().forward(
hidden_states, audio_hidden_states, encoder_hidden_states, timestep,
timestep_indices, token_tags, position_ids, video_indices, audio_indices,
text_indices, attention_kwargs=attention_kwargs, return_dict=return_dict)
@staticmethod
def _video_grid_from_positions(pos):
"""Recover the (T, H, W) token grid of one contiguous video run from its rotary
``position_ids`` [n, 3]. The video rows are frame-major then (h, w) raster with w
fastest: frames carry a distinct temporal coord, and within a frame the leading run
of an identical height coord spans exactly one width row. Returns None on mismatch."""
n = pos.shape[0]
if n == 0:
return None
t_col = pos[:, 0]
rows_per_frame = int((t_col == t_col[0]).sum()) # frame 0's temporal coord is unique
if rows_per_frame <= 0 or n % rows_per_frame != 0:
return None
h_col = pos[:rows_per_frame, 1]
W = int((h_col == h_col[0]).to(torch.int32).cumprod(0).sum()) # one height row
if W <= 0 or rows_per_frame % W != 0:
return None
H = rows_per_frame // W
T = n // rows_per_frame
if T * H * W != n:
return None
return (T, H, W)
def _compute_h3_segments(self, token_tags, position_ids):
"""Packed-order segment list for ``build_h3_tile_geometry``: dense run lengths
(int) plus one (t, h, w) triple per video region, in packed order. Supports
multiple video regions (Ref2VA / PDD reference videos sit as extra tag-0 runs and
each becomes its own sparse region, exempt top-k per region). Returns None (flat
fallback) when the sequence is padded or any video run's grid cannot be recovered."""
if token_tags.numel() == 0 or bool((token_tags < 0).any()):
return None
tags = token_tags.detach().to(torch.int64).cpu()
change = torch.nonzero(tags[1:] != tags[:-1]).flatten() + 1
bounds = torch.cat([torch.zeros(1, dtype=change.dtype), change,
torch.tensor([tags.numel()], dtype=change.dtype)])
bounds = [int(b) for b in bounds.tolist()]
segments = []
video_runs = 0
for i in range(len(bounds) - 1):
s, e = bounds[i], bounds[i + 1]
if e <= s:
continue
if int(tags[s]) == 0: # video run
video_runs += 1
grid = self._video_grid_from_positions(position_ids[s:e])
if grid is None:
return None
segments.append(grid)
else: # dense (text / audio) run
segments.append(e - s)
if video_runs < 1:
return None
return tuple(segments)
def enable_multi_gpus_inference(self,):
"""VSA has no sequence-parallel path in this repo.
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).
"""
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.")
def sparse_trainable_params(self):
"""Names of the extra VSA-branch parameters, for logging / freezing."""
return [name for name, _ in self.named_parameters() if "attn_op." in name]
videox_fun/models/fastvideo_wan_transformer3d.py
+204
View File
@@ -0,0 +1,204 @@
# FastVideo VSA (Video Sparse Attention) student backbone for the Wan2.1 single-expert
# models. Mirrors the SparkDiffusion/RoLa integration in ``spark_transformer3d.py`` and the
# TurboDiffusion/SLA integration in ``turbowan_transformer3d.py``: it swaps the dense softmax
# self-attention of ``WanTransformer3DModel`` for the coarse-to-fine block-sparse VSA kernel
# (reusing the in-repo SLA Triton kernel) plus FastVideo's learned ``to_gate_compress``
# compressed-global compensation branch.
#
# VSA reference: Zhang et al., "VSA: Faster Video Diffusion with Trainable Sparse Attention",
# arXiv:2505.13389 (FastVideo / hao-ai-lab). Official weights: FastVideo/FastWan2.1-T2V-1.3B.
import torch
from .attention_utils import VSASparseAttention, attention
from .wan_transformer3d import (WanSelfAttention, WanTransformer3DModel,
rope_apply_qk)
class FastVideoWanSelfAttention(WanSelfAttention):
"""Wan self-attention whose dense softmax kernel is swapped for VSA block-sparse attention
(sparse fine branch + gated mean-pooled compressed branch)."""
def __init__(self,
dim,
num_heads,
window_size=(-1, -1),
qk_norm=True,
eps=1e-6,
topk=0.2):
super().__init__(dim, num_heads, window_size, qk_norm, eps)
# topk: keep-ratio of key blocks per query (0.2 -> 80% sparse, the official FastWan
# ``--VSA-sparsity 0.8``). >= 1.0 falls back to the original dense attention.
self.topk = topk
if topk < 1.0:
self.attn_op = VSASparseAttention(
self.head_dim, num_heads, dim, topk=topk)
def forward(self, x, seq_lens, grid_sizes, freqs, dtype=torch.bfloat16, t=0):
r"""
Args:
x(Tensor): Shape [B, L, num_heads * head_dim]
seq_lens(Tensor): Shape [B]
grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
"""
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
# 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)
# Dense fallback (topk >= 1.0): original full-softmax attention.
if self.topk >= 1.0:
q, k = rope_apply_qk(q, k, grid_sizes, freqs)
x = attention(
q.to(dtype),
k.to(dtype),
v=v.to(dtype),
k_lens=seq_lens,
window_size=self.window_size
)
x = x.to(dtype).flatten(2)
return self.o(x)
q, k = rope_apply_qk(q, k, grid_sizes, freqs)
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)
class FastVideoWanTransformer3DModel(WanTransformer3DModel):
"""Wan2.1 transformer (T2V / I2V, single expert) with a VSA sparse self-attention. Block
type stays WanAttentionBlock, so the existing
`--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock` policy keeps working."""
# NOTE: diffusers from_pretrained filters config.json keys by __init__ signature, so
# every parent parameter must be declared explicitly (a *args/**kwargs passthrough would
# silently drop dim/ffn_dim/... and build a wrongly-sized model).
def __init__(
self,
model_type='t2v',
patch_size=(1, 2, 2),
text_len=512,
in_dim=16,
dim=2048,
ffn_dim=8192,
freq_dim=256,
text_dim=4096,
out_dim=16,
num_heads=16,
num_layers=32,
window_size=(-1, -1),
qk_norm=True,
cross_attn_norm=True,
eps=1e-6,
in_channels=16,
hidden_size=2048,
add_control_adapter=False,
in_dim_control_adapter=24,
downscale_factor_control_adapter=8,
add_ref_conv=False,
in_dim_ref_conv=16,
cross_attn_type=None,
student_topk=0.2,
):
super().__init__(
model_type=model_type,
patch_size=patch_size,
text_len=text_len,
in_dim=in_dim,
dim=dim,
ffn_dim=ffn_dim,
freq_dim=freq_dim,
text_dim=text_dim,
out_dim=out_dim,
num_heads=num_heads,
num_layers=num_layers,
window_size=window_size,
qk_norm=qk_norm,
cross_attn_norm=cross_attn_norm,
eps=eps,
in_channels=in_channels,
hidden_size=hidden_size,
add_control_adapter=add_control_adapter,
in_dim_control_adapter=in_dim_control_adapter,
downscale_factor_control_adapter=downscale_factor_control_adapter,
add_ref_conv=add_ref_conv,
in_dim_ref_conv=in_dim_ref_conv,
cross_attn_type=cross_attn_type,
)
self.student_topk = student_topk
self.register_to_config(student_topk=student_topk)
self._apply_sparse_self_attention(student_topk)
def _apply_sparse_self_attention(self, topk):
for block in self.blocks:
dense_attn = block.self_attn
sparse_attn = FastVideoWanSelfAttention(
dense_attn.dim, dense_attn.num_heads, dense_attn.window_size,
dense_attn.qk_norm, dense_attn.eps, topk=topk,
)
# Under low_cpu_mem_usage from_pretrained the replaced params are meta tensors, so
# init_weights_' zeros carry no data: wan_transformer3d.py's loader re-materialises
# every key the checkpoint lacks, and `to_gate_compress` is on its zero-init
# whitelist, so a cold start from a dense base still gets a zero gate. Only upgrade
# in place (copy the real tensors over) when the weights are concrete.
if not next(dense_attn.parameters()).is_meta:
old_sd = {k: v for k, v in dense_attn.state_dict().items() if not v.is_meta}
# Match the live params' device AND dtype: freshly built VSA params are fp32,
# and mixing them with bf16 backbone weights breaks FSDP flat-param flattening.
p0 = next(dense_attn.parameters())
sparse_attn = sparse_attn.to(device=p0.device, dtype=p0.dtype)
sparse_attn.load_state_dict(old_sd, strict=False)
block.self_attn = sparse_attn
def set_student_topk(self, topk):
self.student_topk = topk
self.register_to_config(student_topk=topk)
if topk < 1.0 and all(hasattr(block.self_attn, "attn_op") and
not any(p.is_meta for p in block.self_attn.attn_op.parameters())
for block in self.blocks):
for block in self.blocks:
block.self_attn.topk = topk
block.self_attn.attn_op.topk = topk
else:
self._apply_sparse_self_attention(topk)
def enable_multi_gpus_inference(self,):
"""VSA has no sequence-parallel path in this repo.
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).
"""
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.")
def sparse_trainable_params(self):
"""Names of the extra VSA-branch parameters, for logging / freezing."""
return [name for name, _ in self.named_parameters() if "attn_op." in name]
videox_fun/models/wan_gan_discriminator.py
+108
View File
@@ -0,0 +1,108 @@
# Modified from https://github.com/Aurora-edu/One-Forcing/blob/main/wan/modules/model.py
r"""One-Forcing GAN discriminator head for the Wan fake score.
The discriminator shares the fake-score backbone: one learnable register token per selected block cross-attends to
that block's hidden states, and the concatenated tokens go through a small classification branch to a single logit.
[`add_wan_gan_discriminator`] attaches the head as the `gan_head` submodule of a [`WanTransformer3DModel`] and wraps its
forward, so calling the model with `gan_classify=True` returns `(flow_pred, logits)`. The head runs inside the forward,
so FSDP shards and trains it together with the backbone.
"""
import math
import types
import torch
import torch.nn as nn
import torch.nn.functional as F
from .wan_transformer3d import WanLayerNorm, WanRMSNorm
class GanAttentionBlock(nn.Module):
"""Register-token cross-attention block of the One-Forcing discriminator head."""
def __init__(self, dim, ffn_dim, num_heads, eps=1e-6):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.q = nn.Linear(dim, dim)
self.k = nn.Linear(dim, dim)
self.v = nn.Linear(dim, dim)
self.o = nn.Linear(dim, dim)
self.norm_q = WanRMSNorm(dim, eps=eps)
self.norm_k = WanRMSNorm(dim, eps=eps)
self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True)
self.norm2 = WanLayerNorm(dim, eps)
self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim))
def forward(self, x, token):
b, n, d = x.size(0), self.num_heads, self.head_dim
x = self.norm3(x)
q = self.norm_q(self.q(token)).view(b, -1, n, d).transpose(1, 2)
k = self.norm_k(self.k(x)).view(b, -1, n, d).transpose(1, 2)
v = self.v(x).view(b, -1, n, d).transpose(1, 2)
token = token + self.o(F.scaled_dot_product_attention(q, k, v).transpose(1, 2).flatten(2))
return token + self.ffn(self.norm2(token))
class GanDiscriminatorHead(nn.Module):
"""One-Forcing discriminator head: one register token per fake_score feature layer."""
def __init__(self, dim, num_layers, ffn_dim, num_heads):
super().__init__()
self.register_tokens = nn.Parameter(torch.randn(num_layers, dim) * 0.02)
self.register_norm = WanRMSNorm(dim, eps=1e-6)
self.blocks = nn.ModuleList([GanAttentionBlock(dim, ffn_dim, num_heads) for _ in range(num_layers)])
self.cls_pred_branch = nn.Sequential(
nn.LayerNorm(dim * num_layers), nn.Linear(dim * num_layers, dim), nn.SiLU(), nn.Linear(dim, 1)
)
def forward(self, features):
registers = self.register_norm(self.register_tokens)
tokens = []
for i, (block, feature) in enumerate(zip(self.blocks, features)):
token = registers[i:i + 1].unsqueeze(0).expand(feature.size(0), -1, -1)
# Checkpoint each cross-attention block so its q/k/v activations are recomputed during
# backward instead of being retained; one register token keeps the recompute cost tiny.
if torch.is_grad_enabled() and self.training:
token = torch.utils.checkpoint.checkpoint(block, feature, token, use_reentrant=False)
else:
token = block(feature, token)
tokens.append(token)
return self.cls_pred_branch(torch.cat(tokens, dim=1).flatten(1))
def add_wan_gan_discriminator(transformer, feature_layers, ffn_dim, num_heads, dtype, seed=0):
"""
Attach the One-Forcing discriminator head to fake_score before accelerator.prepare.
Calling the model with gan_classify=True returns (flow_pred, logits); the head runs
inside the wrapped forward so it is sharded and trained together with fake_score.
"""
feature_layers = sorted(set(feature_layers))
with torch.random.fork_rng(devices=[]):
torch.manual_seed(seed)
transformer.gan_head = GanDiscriminatorHead(transformer.dim, len(feature_layers), ffn_dim, num_heads)
transformer.gan_head.to(device=next(transformer.parameters()).device, dtype=dtype)
transformer.gan_head.requires_grad_(True)
transformer._gan_features = None
def capture_feature(module, inputs, output):
if transformer._gan_features is not None:
transformer._gan_features.append(output)
for layer in feature_layers:
transformer.blocks[layer].register_forward_hook(capture_feature)
def forward(self, *args, gan_classify=False, **kwargs):
if not gan_classify:
return type(self).forward(self, *args, **kwargs)
self._gan_features = []
try:
flow = type(self).forward(self, *args, **kwargs)
features = self._gan_features
finally:
self._gan_features = None
num_tokens = math.prod(flow.shape[2:]) // math.prod(self.patch_size)
return flow, self.gan_head([feature[:, :num_tokens] for feature in features])
transformer.forward = types.MethodType(forward, transformer)
videox_fun/models/wan_tae.py
+302 -66
View File
@@ -13,8 +13,7 @@ import torch.nn.functional as F
from collections import namedtuple
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.loaders.single_file_model import FromOriginalModelMixin
from diffusers.models.autoencoders.vae import (DecoderOutput,
DiagonalGaussianDistribution)
from diffusers.models.autoencoders.vae import DecoderOutput, DiagonalGaussianDistribution
from diffusers.models.modeling_outputs import AutoencoderKLOutput
from diffusers.models.modeling_utils import ModelMixin
from diffusers.utils.accelerate_utils import apply_forward_hook
@@ -39,7 +38,9 @@ class Clamp(nn.Module):
class MemBlock(nn.Module):
def __init__(self, n_in, n_out):
super().__init__()
self.conv = nn.Sequential(conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out))
self.conv = nn.Sequential(
conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out)
)
self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
self.act = nn.ReLU(inplace=True)
@@ -49,11 +50,13 @@ class MemBlock(nn.Module):
class SuperMemBlock(nn.Module):
"""MemBlock variant used by the Super decoder (ConvNeXt-style: 7x7 depthwise conv + inverted bottleneck)."""
def __init__(self, n_f):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(n_f * 2, n_f * 2, 7, padding=3, groups=n_f * 2, bias=False),
nn.Conv2d(n_f * 2, n_f * 4, 1), nn.ReLU(inplace=True),
nn.Conv2d(n_f * 2, n_f * 4, 1),
nn.ReLU(inplace=True),
nn.Conv2d(n_f * 4, n_f, 1, bias=False),
)
@@ -208,7 +211,15 @@ class TAEWan(nn.Module):
(~Gaussian) latents, matching the reference taehv implementation.
"""
def __init__(self, encoder_time_downscale=(True, True, False), decoder_time_upscale=(False, True, True), decoder_space_upscale=(True, True, True), patch_size=1, latent_channels=16, arch_variant=None):
def __init__(
self,
encoder_time_downscale=(True, True, False),
decoder_time_upscale=(False, True, True),
decoder_space_upscale=(True, True, True),
patch_size=1,
latent_channels=16,
arch_variant=None,
):
"""
Args:
encoder_time_downscale: whether temporal downsampling is enabled for each block.
@@ -226,29 +237,78 @@ class TAEWan(nn.Module):
decoder_time_upscale = (False, *decoder_time_upscale)
assert arch_variant in (None, "super"), f"unrecognized arch_variant {arch_variant!r}"
self.encoder = nn.Sequential(
conv(self.image_channels * self.patch_size ** 2, 64), nn.ReLU(inplace=True),
TPool(64, 2 if encoder_time_downscale[0] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
TPool(64, 2 if encoder_time_downscale[1] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
TPool(64, 2 if encoder_time_downscale[2] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
conv(self.image_channels * self.patch_size ** 2, 64),
nn.ReLU(inplace=True),
TPool(64, 2 if encoder_time_downscale[0] else 1),
conv(64, 64, stride=2, bias=False),
MemBlock(64, 64),
MemBlock(64, 64),
MemBlock(64, 64),
TPool(64, 2 if encoder_time_downscale[1] else 1),
conv(64, 64, stride=2, bias=False),
MemBlock(64, 64),
MemBlock(64, 64),
MemBlock(64, 64),
TPool(64, 2 if encoder_time_downscale[2] else 1),
conv(64, 64, stride=2, bias=False),
MemBlock(64, 64),
MemBlock(64, 64),
MemBlock(64, 64),
conv(64, self.latent_channels),
)
if arch_variant == "super":
n_f = [512, 256, 128, 64]
self.decoder = nn.Sequential(
nn.Conv2d(self.latent_channels, n_f[0], 1, bias=False),
SuperMemBlock(n_f[0]), SuperMemBlock(n_f[0]), SuperMemBlock(n_f[0]), conv(n_f[0], n_f[1] * (2 if decoder_space_upscale[0] else 1) ** 2), nn.ReLU(inplace=True), nn.PixelShuffle(2 if decoder_space_upscale[0] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1),
SuperMemBlock(n_f[1]), SuperMemBlock(n_f[1]), SuperMemBlock(n_f[1]), conv(n_f[1], n_f[2] * (2 if decoder_space_upscale[1] else 1) ** 2), nn.ReLU(inplace=True), nn.PixelShuffle(2 if decoder_space_upscale[1] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1),
SuperMemBlock(n_f[2]), SuperMemBlock(n_f[2]), SuperMemBlock(n_f[2]), conv(n_f[2], n_f[3] * (2 if decoder_space_upscale[2] else 1) ** 2), nn.ReLU(inplace=True), nn.PixelShuffle(2 if decoder_space_upscale[2] else 1), TGrow(n_f[3], 2 if decoder_time_upscale[2] else 1),
SuperMemBlock(n_f[0]),
SuperMemBlock(n_f[0]),
SuperMemBlock(n_f[0]),
conv(n_f[0], n_f[1] * (2 if decoder_space_upscale[0] else 1) ** 2),
nn.ReLU(inplace=True),
nn.PixelShuffle(2 if decoder_space_upscale[0] else 1),
TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1),
SuperMemBlock(n_f[1]),
SuperMemBlock(n_f[1]),
SuperMemBlock(n_f[1]),
conv(n_f[1], n_f[2] * (2 if decoder_space_upscale[1] else 1) ** 2),
nn.ReLU(inplace=True),
nn.PixelShuffle(2 if decoder_space_upscale[1] else 1),
TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1),
SuperMemBlock(n_f[2]),
SuperMemBlock(n_f[2]),
SuperMemBlock(n_f[2]),
conv(n_f[2], n_f[3] * (2 if decoder_space_upscale[2] else 1) ** 2),
nn.ReLU(inplace=True),
nn.PixelShuffle(2 if decoder_space_upscale[2] else 1),
TGrow(n_f[3], 2 if decoder_time_upscale[2] else 1),
conv(n_f[3], self.image_channels * self.patch_size ** 2),
)
else:
n_f = [256, 128, 64, 64]
self.decoder = nn.Sequential(
Clamp(), conv(self.latent_channels, n_f[0]), nn.ReLU(inplace=True),
MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 2 if decoder_time_upscale[0] else 1), conv(n_f[0], n_f[1], bias=False),
MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[1] else 1), conv(n_f[1], n_f[2], bias=False),
MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[2] else 1), conv(n_f[2], n_f[3], bias=False),
nn.ReLU(inplace=True), conv(n_f[3], self.image_channels * self.patch_size ** 2),
Clamp(),
conv(self.latent_channels, n_f[0]),
nn.ReLU(inplace=True),
MemBlock(n_f[0], n_f[0]),
MemBlock(n_f[0], n_f[0]),
MemBlock(n_f[0], n_f[0]),
nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1),
TGrow(n_f[0], 2 if decoder_time_upscale[0] else 1),
conv(n_f[0], n_f[1], bias=False),
MemBlock(n_f[1], n_f[1]),
MemBlock(n_f[1], n_f[1]),
MemBlock(n_f[1], n_f[1]),
nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1),
TGrow(n_f[1], 2 if decoder_time_upscale[1] else 1),
conv(n_f[1], n_f[2], bias=False),
MemBlock(n_f[2], n_f[2]),
MemBlock(n_f[2], n_f[2]),
MemBlock(n_f[2], n_f[2]),
nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1),
TGrow(n_f[2], 2 if decoder_time_upscale[2] else 1),
conv(n_f[2], n_f[3], bias=False),
nn.ReLU(inplace=True),
conv(n_f[3], self.image_channels * self.patch_size ** 2),
)
# computed properties
self.t_downscale = 2 ** sum(t.stride == 2 for t in self.encoder if isinstance(t, TPool))
@@ -268,7 +328,7 @@ class TAEWan(nn.Module):
key = f"decoder.{i}.conv.weight"
if key in sd and sd[key].shape[0] > new_sd[key].shape[0]:
# take the last-timestep output channels
sd[key] = sd[key][-new_sd[key].shape[0]:]
sd[key] = sd[key][-new_sd[key].shape[0] :]
return sd
def preprocess_input_frames(self, x):
@@ -293,6 +353,8 @@ class TAEWan(nn.Module):
if False, frames will be processed sequentially.
Returns NTCHW latent tensor with ~Gaussian values.
"""
if isinstance(self.encoder, nn.Identity):
raise RuntimeError("仅 Decoder 的 TAE 权重不支持 encode；请加载包含 Encoder 的完整权重。")
x = self.preprocess_input_frames(x)
if x.shape[1] % self.t_downscale != 0:
# pad at end to multiple of self.t_downscale
@@ -313,7 +375,7 @@ class TAEWan(nn.Module):
"""
x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar)
x = self.postprocess_output_frames(x)
return x[:, self.frames_to_trim:]
return x[:, self.frames_to_trim :]
class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
@@ -330,24 +392,142 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
"""
# per-channel statistics of the full-size Wan VAE latent spaces
WAN2_1_MEAN = [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921]
WAN2_1_STD = [2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160]
WAN2_1_MEAN = [
-0.7571,
-0.7089,
-0.9113,
0.1075,
-0.1745,
0.9653,
-0.1517,
1.5508,
0.4134,
-0.0715,
0.5517,
-0.3632,
-0.1922,
-0.9497,
0.2503,
-0.2921,
]
WAN2_1_STD = [
2.8184,
1.4541,
2.3275,
2.6558,
1.2196,
1.7708,
2.6052,
2.0743,
3.2687,
2.1526,
2.8652,
1.5579,
1.6382,
1.1253,
2.8251,
1.9160,
]
WAN2_2_MEAN = [
-0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557,
-0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825,
-0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502,
-0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230,
-0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748,
0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667]
-0.2289,
-0.0052,
-0.1323,
-0.2339,
-0.2799,
0.0174,
0.1838,
0.1557,
-0.1382,
0.0542,
0.2813,
0.0891,
0.1570,
-0.0098,
0.0375,
-0.1825,
-0.2246,
-0.1207,
-0.0698,
0.5109,
0.2665,
-0.2108,
-0.2158,
0.2502,
-0.2055,
-0.0322,
0.1109,
0.1567,
-0.0729,
0.0899,
-0.2799,
-0.1230,
-0.0313,
-0.1649,
0.0117,
0.0723,
-0.2839,
-0.2083,
-0.0520,
0.3748,
0.0152,
0.1957,
0.1433,
-0.2944,
0.3573,
-0.0548,
-0.1681,
-0.0667,
]
WAN2_2_STD = [
0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013,
0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978,
0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659,
0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093,
0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887,
0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744]
0.4765,
1.0364,
0.4514,
1.1677,
0.5313,
0.4990,
0.4818,
0.5013,
0.8158,
1.0344,
0.5894,
1.0901,
0.6885,
0.6165,
0.8454,
0.4978,
0.5759,
0.3523,
0.7135,
0.6804,
0.5833,
1.4146,
0.8986,
0.5659,
0.7069,
0.5338,
0.4889,
0.4917,
0.4069,
0.4999,
0.6866,
0.4093,
0.5709,
0.6065,
0.6415,
0.4944,
0.5726,
1.2042,
0.5458,
1.6887,
0.3971,
1.0600,
0.3943,
0.5537,
0.5444,
0.4089,
0.7468,
0.7744,
]
@register_to_config
def __init__(
@@ -361,6 +541,7 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
temporal_compression_ratio=4,
spatial_compression_ratio=None,
parallel=True,
decoder_only=False,
):
super().__init__()
if latent_channels == 16:
@@ -368,7 +549,9 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
elif latent_channels == 48:
mean, std = self.WAN2_2_MEAN, self.WAN2_2_STD
else:
raise ValueError(f"AutoencoderTinyWan supports 16ch (taew2_1) or 48ch (taew2_2) latents, got {latent_channels}")
raise ValueError(
f"AutoencoderTinyWan supports 16ch (taew2_1) or 48ch (taew2_2) latents, got {latent_channels}"
)
assert len(mean) == latent_channels and len(std) == latent_channels
self.mean = torch.tensor(mean, dtype=torch.float32)
self.std = torch.tensor(std, dtype=torch.float32)
@@ -387,6 +570,8 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
latent_channels=latent_channels,
arch_variant=arch_variant,
)
if decoder_only:
self.model.encoder = nn.Identity()
self.parallel = parallel
self.gradient_checkpointing = False
@@ -406,12 +591,12 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
# latents_mean=zeros/latents_std=ones and passes pipeline latents
# through unchanged), so no mean/std scaling is applied here: the
# encoder output is already interchangeable with full-VAE latents.
x = (x.add(1).div_(2)).clamp_(0, 1) # [-1, 1] -> [0, 1]
x = x.permute(0, 2, 1, 3, 4) # BCTHW -> NTCHW
x = (x.add(1).div_(2)).clamp_(0, 1) # [-1, 1] -> [0, 1]
x = x.permute(0, 2, 1, 3, 4) # BCTHW -> NTCHW
zs = []
for u in x:
z = self.model.encode_video(u.unsqueeze(0), parallel=self.parallel).squeeze(0)
zs.append(z.permute(1, 0, 2, 3)) # TCHW -> CTHW
zs.append(z.permute(1, 0, 2, 3)) # TCHW -> CTHW
return torch.stack(zs)
@apply_forward_hook
@@ -432,7 +617,7 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
# zs: BCLHW normalized latents -> BCTHW in [-1, 1].
# NOTE: no denormalization here, for the same reason as in _encode:
# the official TAE decoder consumes normalized latents directly.
zs = zs.permute(0, 2, 1, 3, 4) # BCTHW -> NTCHW
zs = zs.permute(0, 2, 1, 3, 4) # BCTHW -> NTCHW
dec = []
for u in zs:
d = self.model.decode_video(u.unsqueeze(0), parallel=self.parallel).squeeze(0)
@@ -450,48 +635,99 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
return DecoderOutput(sample=decoded)
@classmethod
def from_pretrained(cls, pretrained_model_path, additional_kwargs={}, **kwargs):
def from_pretrained(cls, pretrained_model_path, additional_kwargs=None, **kwargs):
def filter_kwargs(cls, kwargs):
import inspect
sig = inspect.signature(cls.__init__)
valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
return filtered_kwargs
# standard diffusers directory kwargs (e.g. subfolder= from load_model_hook)
subfolder = kwargs.get("subfolder", None)
pretrained_model_path = os.fspath(pretrained_model_path)
subfolder = kwargs.pop("subfolder", None)
if subfolder:
pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
kwargs = filter_kwargs(cls, additional_kwargs)
# support diffusers directory checkpoints (save_pretrained output)
model_kwargs = filter_kwargs(cls, additional_kwargs or {})
model_kwargs.update(filter_kwargs(cls, kwargs))
if os.path.isdir(pretrained_model_path):
if os.path.exists(os.path.join(pretrained_model_path, "config.json")):
return super().from_pretrained(pretrained_model_path, **kwargs)
candidates = sorted(f for f in os.listdir(pretrained_model_path)
if f.endswith((".pth", ".safetensors")))
if not candidates:
raise ValueError(f"No TAE weights found in {pretrained_model_path}")
loader_kwargs = {k: v for k, v in kwargs.items() if k not in model_kwargs}
return_loading_info = loader_kwargs.pop("output_loading_info", False)
model, info = super().from_pretrained(
pretrained_model_path, **model_kwargs, **loader_kwargs, output_loading_info=True
)
if any(info.get(k) for k in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")):
raise RuntimeError(f"TAE 配置与权重不匹配或权重不完整：{info}")
return (model, info) if return_loading_info else model
candidates = sorted(
f for f in os.listdir(pretrained_model_path) if f.endswith((".pt", ".pth", ".bin", ".safetensors"))
)
if len(candidates) != 1:
raise ValueError(f"请指定唯一的 TAE 权重文件，目录内找到：{candidates}")
pretrained_model_path = os.path.join(pretrained_model_path, candidates[0])
# infer model flavour from the checkpoint filename when not given explicitly
checkpoint_name = pretrained_model_path.lower()
if "latent_channels" not in kwargs:
if "taew2_2" in checkpoint_name:
kwargs["latent_channels"], kwargs["patch_size"] = 48, 2
else:
kwargs["latent_channels"], kwargs["patch_size"] = 16, 1
if kwargs.get("arch_variant") is None and "_super" in checkpoint_name:
kwargs["arch_variant"] = "super"
model = cls(**kwargs)
if pretrained_model_path.endswith(".safetensors"):
from safetensors.torch import load_file
state_dict = load_file(pretrained_model_path)
else:
state_dict = torch.load(pretrained_model_path, map_location="cpu", weights_only=True)
# TAE checkpoints already use "encoder.*" / "decoder.*" keys, no prefix needed
state_dict = model.model.patch_tgrow_layers(state_dict)
m, u = model.model.load_state_dict(state_dict, strict=False)
print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
print(m, u)
if isinstance(state_dict, dict):
if isinstance(state_dict.get("state_dict"), dict):
state_dict = state_dict["state_dict"]
elif isinstance(state_dict.get("decoder"), dict):
state_dict = state_dict["decoder"]
if (
not isinstance(state_dict, dict)
or not state_dict
or not all(isinstance(k, str) and isinstance(v, torch.Tensor) for k, v in state_dict.items())
):
raise ValueError("无效的 TAE 权重：需要 tensor state_dict 或 checkpoint['decoder']。")
if all(k.startswith("model.") for k in state_dict):
state_dict = {k[len("model.") :]: v for k, v in state_dict.items()}
if "decoder.0.weight" in state_dict and "decoder.1.weight" not in state_dict:
arch_variant, first_key, block_key = "super", "decoder.0.weight", "decoder.1.conv.0.weight"
first_channels, first_kernel, block_shape = 512, (1, 1), (1024, 1, 7, 7)
elif "decoder.1.weight" in state_dict and "decoder.0.weight" not in state_dict:
arch_variant, first_key, block_key = None, "decoder.1.weight", "decoder.3.conv.0.weight"
first_channels, first_kernel, block_shape = 256, (3, 3), (256, 512, 3, 3)
else:
raise ValueError("无法识别 TAE 架构：Decoder 首层权重缺失或存在冲突。")
first = state_dict[first_key]
block = state_dict.get(block_key)
last = state_dict.get("decoder.22.weight")
if (
first.ndim != 4
or first.shape[0] != first_channels
or tuple(first.shape[2:]) != first_kernel
or block is None
or tuple(block.shape) != block_shape
or last is None
or last.ndim != 4
):
raise ValueError("无法识别 TAE 架构：首层、记忆块或输出层的形状不符合普通/Super TAE。")
latent_channels = first.shape[1]
patch_size = {16: 1, 48: 2}.get(latent_channels)
if patch_size is None or tuple(last.shape) != (3 * patch_size ** 2, 64, 3, 3):
raise ValueError("TAE latent 通道或输出 patch_size 不符合 Wan 2.1/2.2 架构。")
inferred = {
"arch_variant": arch_variant,
"latent_channels": latent_channels,
"patch_size": patch_size,
"decoder_only": not any(k.startswith("encoder.") for k in state_dict),
}
for key, value in inferred.items():
if key in model_kwargs and model_kwargs[key] != value:
raise ValueError(f"TAE 参数 {key}={model_kwargs[key]!r} 与权重识别值 {value!r} 冲突。")
model_kwargs[key] = value
model = cls(**model_kwargs)
state_dict = model.model.patch_tgrow_layers(dict(state_dict))
model.model.load_state_dict(state_dict, strict=True)
torch_dtype = kwargs.get("torch_dtype")
if torch_dtype == "auto":
torch_dtype = first.dtype
if torch_dtype is not None:
model.to(dtype=torch_dtype)
return model
videox_fun/models/wan_transformer3d.py
+7 -3
View File
@@ -1325,10 +1325,14 @@ class WanTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
elif (
('audio_injector' in key and '.o.' in key)
or ('injector_adain_layers' in key and '.linear.' in key)
or any(name in key for name in ('trainable_cond_mask', 'cond_encoder', 'zip_motion_out.1', 'proj_l'))
or any(name in key for name in ('trainable_cond_mask', 'cond_encoder', 'zip_motion_out.1', 'proj_l', 'to_gate_compress'))
):
# S2V zero_init_weights() / turbowan SparseLinearAttention.init_weights_()
# zero these modules in __init__, which is lost under init_empty_weights.
# S2V zero_init_weights() / turbowan SparseLinearAttention.init_weights_() /
# VSASparseAttention.init_weights_() zero these modules in __init__, which is
# lost under init_empty_weights. A VSA student cold-started from a dense base
# has no attn_op.to_gate_compress keys in the checkpoint, so without this
# branch the gate falls through to xavier_uniform_ below and the compressed
# global branch starts at full random magnitude instead of as a no-op.
initialized_dict[key] = torch.zeros(param_shape, dtype=param_dtype)
print(f"Initializing missing parameter '{key}' with zero.")
elif 'weight' in key: