Update Sparkwan, Reward Model Training, FastVideo (#525)
This commit is contained in:
Bubbliiiing
@@ -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()
|
||||
@@ -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()
|
||||
@@ -46,7 +46,7 @@ ring_degree = 1
|
||||
# Use FSDP to save more GPU memory in multi gpus. The Qwen3-VL conditioner is ~62 GB, so with fsdp_dit alone every
|
||||
# rank still replicates it; fsdp_text_encoder shards it too. Note it must wrap the inner `text_encoder.model`
|
||||
# (Qwen3VLModel): encode_prompt calls that submodule directly, so a wrap on the top-level module would never fire.
|
||||
fsdp_dit = False
|
||||
fsdp_dit = True
|
||||
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 sequential_cpu_offload.
|
||||
@@ -207,6 +207,11 @@ 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
|
||||
transformer.enable_multi_gpus_inference()
|
||||
@@ -215,12 +220,13 @@ if ulysses_degree > 1 or ring_degree > 1:
|
||||
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)
|
||||
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))
|
||||
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")
|
||||
|
||||
|
||||
@@ -47,7 +47,7 @@ ring_degree = 1
|
||||
# Use FSDP to save more GPU memory in multi gpus. The Qwen3-VL conditioner is ~62 GB, so with fsdp_dit alone every
|
||||
# rank still replicates it; fsdp_text_encoder shards it too. Note it must wrap the inner `text_encoder.model`
|
||||
# (Qwen3VLModel): encode_prompt calls that submodule directly, so a wrap on the top-level module would never fire.
|
||||
fsdp_dit = False
|
||||
fsdp_dit = True
|
||||
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 sequential_cpu_offload.
|
||||
@@ -213,6 +213,11 @@ 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
|
||||
transformer.enable_multi_gpus_inference()
|
||||
@@ -221,12 +226,13 @@ if ulysses_degree > 1 or ring_degree > 1:
|
||||
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)
|
||||
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))
|
||||
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")
|
||||
|
||||
@@ -263,6 +269,7 @@ def parse_reference(entry: str):
|
||||
return MiniMaxH3AudioReference.from_file(media)
|
||||
raise ValueError(f"A reference entry must start with `image=`, `video=` or `audio=`, got {entry!r}.")
|
||||
|
||||
|
||||
# Decode every reference at the rate its container carries, which the pipeline's setup resamples onto MiniMax-H3's
|
||||
# own 24 fps and the audio VAE's sample rate.
|
||||
parsed_references = [parse_reference(entry) for entry in references]
|
||||
|
||||
@@ -44,7 +44,7 @@ ring_degree = 1
|
||||
# Use FSDP to save more GPU memory in multi gpus. The Qwen3-VL conditioner is ~62 GB, so with fsdp_dit alone every
|
||||
# rank still replicates it; fsdp_text_encoder shards it too. Note it must wrap the inner `text_encoder.model`
|
||||
# (Qwen3VLModel): encode_prompt calls that submodule directly, so a wrap on the top-level module would never fire.
|
||||
fsdp_dit = False
|
||||
fsdp_dit = True
|
||||
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 sequential_cpu_offload.
|
||||
@@ -199,6 +199,11 @@ 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
|
||||
transformer.enable_multi_gpus_inference()
|
||||
@@ -207,12 +212,13 @@ if ulysses_degree > 1 or ring_degree > 1:
|
||||
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)
|
||||
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))
|
||||
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")
|
||||
|
||||
|
||||
@@ -21,6 +21,7 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
|
||||
get_video_to_video_latent, 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.
|
||||
#
|
||||
@@ -51,7 +52,7 @@ ring_degree = 1
|
||||
# Use FSDP to save more GPU memory in multi gpus. The Qwen3-VL conditioner is ~62 GB, so with fsdp_dit alone every
|
||||
# rank still replicates it; fsdp_text_encoder shards it too. Note it must wrap the inner `text_encoder.model`
|
||||
# (Qwen3VLModel): encode_prompt calls that submodule directly, so a wrap on the top-level module would never fire.
|
||||
fsdp_dit = False
|
||||
fsdp_dit = True
|
||||
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 sequential_cpu_offload.
|
||||
@@ -215,6 +216,11 @@ 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
|
||||
transformer.enable_multi_gpus_inference()
|
||||
@@ -223,12 +229,13 @@ if ulysses_degree > 1 or ring_degree > 1:
|
||||
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) + list(transformer.control_blocks),
|
||||
ignored_modules=fp32_modules)
|
||||
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))
|
||||
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")
|
||||
|
||||
|
||||
@@ -21,6 +21,7 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
|
||||
get_video_to_video_latent, 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.
|
||||
#
|
||||
@@ -51,7 +52,7 @@ ring_degree = 1
|
||||
# Use FSDP to save more GPU memory in multi gpus. The Qwen3-VL conditioner is ~62 GB, so with fsdp_dit alone every
|
||||
# rank still replicates it; fsdp_text_encoder shards it too. Note it must wrap the inner `text_encoder.model`
|
||||
# (Qwen3VLModel): encode_prompt calls that submodule directly, so a wrap on the top-level module would never fire.
|
||||
fsdp_dit = False
|
||||
fsdp_dit = True
|
||||
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 sequential_cpu_offload.
|
||||
@@ -218,6 +219,11 @@ if use_qfloat8:
|
||||
# Scale-aware fp8 must run before the FSDP wrapping below so the flat buffers hold the fp8 values.
|
||||
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
|
||||
transformer.enable_multi_gpus_inference()
|
||||
@@ -239,12 +245,13 @@ if ulysses_degree > 1 or ring_degree > 1:
|
||||
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) + list(transformer.control_blocks),
|
||||
ignored_modules=fp32_modules)
|
||||
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))
|
||||
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")
|
||||
|
||||
|
||||
@@ -0,0 +1,304 @@
|
||||
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,
|
||||
SparkWanTransformer3DModel, 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 = "sequential_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 sparse (RoLa) kernel only has an ulysses head-parallel path (no ring attention), so a
|
||||
# sparse student requires ring_degree=1; put every GPU into ulysses_degree instead.
|
||||
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"
|
||||
|
||||
# ---- SparkDiffusion sparse-student config -------------------------------------------------
|
||||
# Official released checkpoint: SparkWan2.1-T2V-1.3B-480P-0.90Sparsity (RoLA sparse attention +
|
||||
# CrossDistill 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.10 -> 90% sparsity. It selects how
|
||||
# many key blocks the sparse kernel keeps, so a mismatch silently mis-attends.
|
||||
student_topk = 0.10
|
||||
|
||||
# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++".
|
||||
# The few-step distilled student uses "Flow" (FlowMatchEulerDiscreteScheduler).
|
||||
sampler_name = "Flow"
|
||||
# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics.
|
||||
shift = 5
|
||||
|
||||
# Official released checkpoint (single .pth, bf16, flat 'net.'-prefixed state_dict).
|
||||
# Download:
|
||||
# HF_ENDPOINT=https://hf-mirror.com huggingface-cli download \
|
||||
# alibabagroup/SparkWan2.1-T2V-1.3B-480P-0.90Sparsity --local-dir <dir>
|
||||
transformer_path = "models/Personalized_Model/SparkWan2.1-T2V-1.3B-480P-0.90Sparsity/SparkWan2.1-T2V-1.3B-480P-0.90Sparsity.pth"
|
||||
vae_path = None
|
||||
lora_path = None
|
||||
|
||||
# Other params
|
||||
sample_size = [480, 832]
|
||||
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 = "一只棕色的狗摇着头，坐在舒适房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里柔和温暖的灯光营造出舒适的氛围。"
|
||||
negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
|
||||
guidance_scale = 1.0 # the CrossDistill student is CFG-free
|
||||
seed = 43
|
||||
num_inference_steps = 4
|
||||
lora_weight = 0.55
|
||||
save_path = "samples/sparkdiffusion-wan2.1-official"
|
||||
|
||||
|
||||
def remap_spark_ckpt_to_vxf(state_dict):
|
||||
"""Official SparkDiffusion .pth -> VideoX-Fun SparkWanTransformer3DModel key layout.
|
||||
|
||||
The released checkpoint is a flat dict whose keys are all prefixed with ``net.`` and whose
|
||||
RoLa sparse-attention params sit DIRECTLY on ``self_attn`` (proj_q / proj_k / gate_proj /
|
||||
gate_bias). VideoX-Fun keeps the same native Wan backbone naming but nests those four params
|
||||
under a ``self_attn.attn_op`` submodule. Two edits turn the 945 released keys into an exact
|
||||
1:1 match for the model (825 backbone + 30 layers x 4 RoLa = 120):
|
||||
1) drop the leading ``net.``
|
||||
2) ``blocks.N.self_attn.{proj_q,proj_k,gate_proj,gate_bias}`` -> ``...self_attn.attn_op.{...}``
|
||||
"""
|
||||
ROLA = {"proj_q.weight", "proj_k.weight", "gate_proj.weight", "gate_bias"}
|
||||
out = {}
|
||||
for k, v in state_dict.items():
|
||||
nk = k[len("net."):] if k.startswith("net.") else k
|
||||
m = re.match(r"(blocks\.\d+\.self_attn\.)(.+)$", nk)
|
||||
if m and m.group(2) in ROLA:
|
||||
nk = m.group(1) + "attn_op." + m.group(2)
|
||||
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 sparse kernels at the
|
||||
# released keep-ratio BEFORE loading weights. Call order matters: set_student_topk recreates
|
||||
# attn_op (resetting the RoLa params to their init), so the trained checkpoint values must be
|
||||
# loaded afterwards to overwrite them.
|
||||
transformer = SparkWanTransformer3DModel.from_pretrained(
|
||||
os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
|
||||
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
||||
low_cpu_mem_usage=True,
|
||||
torch_dtype=weight_dtype,
|
||||
)
|
||||
transformer.set_student_topk(student_topk)
|
||||
|
||||
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 release -> VideoX-Fun key layout (strip 'net.', route RoLa params under attn_op).
|
||||
state_dict = remap_spark_ckpt_to_vxf(state_dict)
|
||||
|
||||
# The released checkpoint flattens the Conv3d patch_embedding to 2D (out, in*kT*kH*kW);
|
||||
# restore the 5D Conv3d weight exactly like the turbodiffusion loader does.
|
||||
if len(state_dict["patch_embedding.weight"].size()) == 2:
|
||||
out_channels = state_dict["patch_embedding.weight"].shape[0]
|
||||
in_channels = state_dict["patch_embedding.weight"].shape[1] // transformer.config.patch_size[1] // transformer.config.patch_size[-2]
|
||||
state_dict["patch_embedding.weight"] = state_dict["patch_embedding.weight"].reshape(
|
||||
out_channels, in_channels, 1, transformer.config.patch_size[1], transformer.config.patch_size[2]
|
||||
)
|
||||
|
||||
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 sparse student MUST provide every attn_op param; missing ones mean the RoLa branch
|
||||
# would run on freshly-init values -> garbage output. Fail loudly instead of silently.
|
||||
raise RuntimeError(
|
||||
f"{len(attn_op_missing)} sparse 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 0.90-sparsity 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:
|
||||
# The sparse (RoLa) kernel only supports ulysses head-parallel (no ring-attention path); a
|
||||
# sparse student under ring_degree>1 would silently mis-attend.
|
||||
if ring_degree > 1 and transformer_path is not None:
|
||||
raise ValueError(
|
||||
"SparkDiffusion sparse multi-GPU requires ring_degree=1; put every GPU into "
|
||||
"ulysses_degree instead (num_heads must divide ulysses_degree).")
|
||||
from functools import partial
|
||||
transformer.enable_multi_gpus_inference()
|
||||
if fsdp_dit:
|
||||
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)
|
||||
|
||||
if ulysses_degree * ring_degree > 1:
|
||||
import torch.distributed as dist
|
||||
if dist.get_rank() == 0:
|
||||
save_results()
|
||||
else:
|
||||
save_results()
|
||||
@@ -0,0 +1,412 @@
|
||||
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, AutoencoderKLWan3_8,
|
||||
AutoTokenizer, SparkWan2_2Transformer3DModel,
|
||||
Wan2_2Transformer3DModel, WanT5EncoderModel)
|
||||
from videox_fun.models.cache_utils import get_teacache_coefficients
|
||||
from videox_fun.pipeline import Wan2_2Pipeline
|
||||
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 = "sequential_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 sparse (RoLa) kernel only has an ulysses head-parallel path (no ring attention), so a
|
||||
# sparse student requires ring_degree=1; put every GPU into ulysses_degree instead.
|
||||
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
|
||||
|
||||
# TeaCache config
|
||||
# Kept off: TeaCache / cfg-skip skip whole denoising steps, which does not fit the few-step
|
||||
# CFG-free sparse student.
|
||||
enable_teacache = False
|
||||
# Recommended to be set between 0.05 and 0.30. A larger threshold can cache more steps, speeding up the inference process,
|
||||
# but it may cause slight differences between the generated content and the original content.
|
||||
# # --------------------------------------------------------------------------------------------------- #
|
||||
# | Model Name | threshold | Model Name | threshold |
|
||||
# | Wan2.2-T2V-A14B | 0.10~0.15 | Wan2.2-I2V-A14B | 0.15~0.20 |
|
||||
# # --------------------------------------------------------------------------------------------------- #
|
||||
teacache_threshold = 0.10
|
||||
# The number of steps to skip TeaCache at the beginning of the inference process, which can
|
||||
# reduce the impact of TeaCache on generated video quality.
|
||||
num_skip_start_steps = 5
|
||||
# Whether to offload TeaCache tensors to cpu to save a little bit of GPU memory.
|
||||
teacache_offload = False
|
||||
|
||||
# Skip some cfg steps in inference
|
||||
# Recommended to be set between 0.00 and 0.25
|
||||
cfg_skip_ratio = 0
|
||||
|
||||
# Riflex config
|
||||
enable_riflex = False
|
||||
# Index of intrinsic frequency
|
||||
riflex_k = 6
|
||||
|
||||
# Config and model path
|
||||
config_path = "config/wan2.2/wan_civitai_t2v.yaml"
|
||||
# Base model (Wan-AI/Wan2.2-T2V-A14B): provides config.json + dense low/high experts + VAE + T5.
|
||||
model_name = "models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
|
||||
# ---- SparkDiffusion sparse-student config -------------------------------------------------
|
||||
# Official released checkpoints: SparkWan2.2-T2V-14B-480P-0.95Sparsity-{Low,High}.pth (RoLA sparse
|
||||
# attention + CrossDistill few-step distillation, dual-expert MoE). student_topk is the keep-ratio
|
||||
# of key blocks per query and MUST equal the value the checkpoints were trained at: topk=0.05 ->
|
||||
# 95% sparsity. A mismatch silently mis-attends.
|
||||
student_topk = 0.05
|
||||
|
||||
# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++".
|
||||
# The few-step distilled student uses "Flow" (FlowMatchEulerDiscreteScheduler).
|
||||
sampler_name = "Flow"
|
||||
# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics.
|
||||
shift = 12
|
||||
|
||||
# Official released checkpoints (single .pth per expert, bf16, flat 'net.'-prefixed state_dict).
|
||||
# transformer_path -> low-noise expert, transformer_high_path -> high-noise expert.
|
||||
# Download:
|
||||
# HF_ENDPOINT=https://hf-mirror.com HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download \
|
||||
# alibabagroup/SparkWan2.2-T2V-14B-480P-0.95Sparsity --local-dir <dir>
|
||||
transformer_path = "models/Personalized_Model/SparkWan2.2-T2V-14B-480P-0.95Sparsity/SparkWan2.2-T2V-14B-480P-0.95Sparsity-Low.pth"
|
||||
transformer_high_path = "models/Personalized_Model/SparkWan2.2-T2V-14B-480P-0.95Sparsity/SparkWan2.2-T2V-14B-480P-0.95Sparsity-High.pth"
|
||||
vae_path = None
|
||||
lora_path = None
|
||||
lora_high_path = None
|
||||
|
||||
# Other params
|
||||
sample_size = [480, 832]
|
||||
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 = "一只棕色的狗摇着头，坐在舒适房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里柔和温暖的灯光营造出舒适的氛围。"
|
||||
negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
|
||||
guidance_scale = 1.0 # the CrossDistill student is CFG-free
|
||||
seed = 43
|
||||
num_inference_steps = 4
|
||||
lora_weight = 0.55
|
||||
lora_high_weight = 0.55
|
||||
save_path = "samples/sparkdiffusion-wan2.2-official"
|
||||
|
||||
|
||||
def remap_spark_ckpt_to_vxf(state_dict):
|
||||
"""Official SparkDiffusion .pth -> VideoX-Fun SparkWan2_2Transformer3DModel key layout.
|
||||
|
||||
The released checkpoint is a flat dict whose keys are all prefixed with ``net.`` and whose
|
||||
RoLa sparse-attention params sit DIRECTLY on ``self_attn`` (proj_q / proj_k / gate_proj /
|
||||
gate_bias). VideoX-Fun keeps the same native Wan backbone naming but nests those four params
|
||||
under a ``self_attn.attn_op`` submodule. Two edits make the keys match 1:1:
|
||||
1) drop the leading ``net.``
|
||||
2) ``blocks.N.self_attn.{proj_q,proj_k,gate_proj,gate_bias}`` -> ``...self_attn.attn_op.{...}``
|
||||
"""
|
||||
ROLA = {"proj_q.weight", "proj_k.weight", "gate_proj.weight", "gate_bias"}
|
||||
out = {}
|
||||
for k, v in state_dict.items():
|
||||
nk = k[len("net."):] if k.startswith("net.") else k
|
||||
m = re.match(r"(blocks\.\d+\.self_attn\.)(.+)$", nk)
|
||||
if m and m.group(2) in ROLA:
|
||||
nk = m.group(1) + "attn_op." + m.group(2)
|
||||
out[nk] = v
|
||||
return out
|
||||
|
||||
|
||||
def load_spark_ckpt(path):
|
||||
"""Load an official .pth and return it in VideoX-Fun key layout (remap + patch reshape)."""
|
||||
if path.endswith("safetensors"):
|
||||
from safetensors.torch import load_file
|
||||
state_dict = load_file(path)
|
||||
else:
|
||||
state_dict = torch.load(path, map_location="cpu")
|
||||
state_dict = state_dict["state_dict"] if isinstance(state_dict, dict) and "state_dict" in state_dict else state_dict
|
||||
return remap_spark_ckpt_to_vxf(state_dict)
|
||||
|
||||
|
||||
device = set_multi_gpus_devices(ulysses_degree, ring_degree)
|
||||
config = OmegaConf.load(config_path)
|
||||
boundary = config['transformer_additional_kwargs'].get('boundary', 0.875)
|
||||
|
||||
# Build one transformer per MoE expert, then REBUILD the sparse kernels at the released keep-ratio
|
||||
# BEFORE loading weights. Call order matters: set_student_topk recreates attn_op (resetting the
|
||||
# RoLa params to their init), so the trained checkpoint values must be loaded afterwards.
|
||||
low_subpath = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
|
||||
high_subpath = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
|
||||
|
||||
if transformer_path is not None:
|
||||
transformer = SparkWan2_2Transformer3DModel.from_pretrained(
|
||||
os.path.join(model_name, low_subpath),
|
||||
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)
|
||||
else:
|
||||
transformer = Wan2_2Transformer3DModel.from_pretrained(
|
||||
os.path.join(model_name, low_subpath),
|
||||
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
||||
low_cpu_mem_usage=True,
|
||||
torch_dtype=weight_dtype,
|
||||
)
|
||||
|
||||
if config['transformer_additional_kwargs'].get('transformer_combination_type', 'single') == "moe":
|
||||
if transformer_high_path is not None:
|
||||
transformer_2 = SparkWan2_2Transformer3DModel.from_pretrained(
|
||||
os.path.join(model_name, high_subpath),
|
||||
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
||||
low_cpu_mem_usage=True,
|
||||
torch_dtype=weight_dtype,
|
||||
)
|
||||
transformer_2.set_student_topk(student_topk)
|
||||
else:
|
||||
transformer_2 = Wan2_2Transformer3DModel.from_pretrained(
|
||||
os.path.join(model_name, high_subpath),
|
||||
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
||||
low_cpu_mem_usage=True,
|
||||
torch_dtype=weight_dtype,
|
||||
)
|
||||
else:
|
||||
transformer_2 = None
|
||||
|
||||
# Load the per-expert official checkpoints AFTER set_student_topk so the trained RoLa params
|
||||
# overwrite the freshly-init ones; a missing attn_op.* key means the remap / topk is wrong.
|
||||
if transformer_path is not None:
|
||||
print(f"From checkpoint: {transformer_path}")
|
||||
state_dict = load_spark_ckpt(transformer_path)
|
||||
if len(state_dict["patch_embedding.weight"].size()) == 2:
|
||||
out_channels = state_dict["patch_embedding.weight"].shape[0]
|
||||
in_channels = state_dict["patch_embedding.weight"].shape[1] // transformer.config.patch_size[1] // transformer.config.patch_size[-2]
|
||||
state_dict["patch_embedding.weight"] = state_dict["patch_embedding.weight"].reshape(
|
||||
out_channels, in_channels, 1, transformer.config.patch_size[1], transformer.config.patch_size[2]
|
||||
)
|
||||
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:
|
||||
raise RuntimeError(
|
||||
f"{len(attn_op_missing)} sparse attn_op keys missing on the low-noise expert "
|
||||
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]}")
|
||||
|
||||
if transformer_2 is not None:
|
||||
if transformer_high_path is not None:
|
||||
print(f"From checkpoint: {transformer_high_path}")
|
||||
state_dict = load_spark_ckpt(transformer_high_path)
|
||||
if len(state_dict["patch_embedding.weight"].size()) == 2:
|
||||
out_channels = state_dict["patch_embedding.weight"].shape[0]
|
||||
in_channels = state_dict["patch_embedding.weight"].shape[1] // transformer_2.config.patch_size[1] // transformer_2.config.patch_size[-2]
|
||||
state_dict["patch_embedding.weight"] = state_dict["patch_embedding.weight"].reshape(
|
||||
out_channels, in_channels, 1, transformer_2.config.patch_size[1], transformer_2.config.patch_size[2]
|
||||
)
|
||||
m, u = transformer_2.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:
|
||||
raise RuntimeError(
|
||||
f"{len(attn_op_missing)} sparse attn_op keys missing on the high-noise expert "
|
||||
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]}")
|
||||
|
||||
# Get Vae
|
||||
Chosen_AutoencoderKL = {
|
||||
"AutoencoderKLWan": AutoencoderKLWan,
|
||||
"AutoencoderKLWan3_8": AutoencoderKLWan3_8
|
||||
}[config['vae_kwargs'].get('vae_type', 'AutoencoderKLWan')]
|
||||
vae = Chosen_AutoencoderKL.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 = Wan2_2Pipeline(
|
||||
transformer=transformer,
|
||||
transformer_2=transformer_2,
|
||||
vae=vae,
|
||||
tokenizer=tokenizer,
|
||||
text_encoder=text_encoder,
|
||||
scheduler=scheduler,
|
||||
)
|
||||
if ulysses_degree > 1 or ring_degree > 1:
|
||||
# The sparse (RoLa) kernel only supports ulysses head-parallel (no ring-attention path); a
|
||||
# sparse expert under ring_degree>1 would silently mis-attend.
|
||||
if ring_degree > 1 and (transformer_path is not None or transformer_high_path is not None):
|
||||
raise ValueError(
|
||||
"SparkDiffusion sparse multi-GPU requires ring_degree=1; put every GPU into "
|
||||
"ulysses_degree instead (num_heads must divide ulysses_degree).")
|
||||
from functools import partial
|
||||
transformer.enable_multi_gpus_inference()
|
||||
if transformer_2 is not None:
|
||||
transformer_2.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)
|
||||
if transformer_2 is not None:
|
||||
pipeline.transformer_2 = shard_fn(pipeline.transformer_2)
|
||||
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])
|
||||
if transformer_2 is not None:
|
||||
for i in range(len(pipeline.transformer_2.blocks)):
|
||||
pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.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)
|
||||
|
||||
coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
|
||||
if coefficients is not None:
|
||||
print(f"Enable TeaCache with threshold {teacache_threshold} and skip the first {num_skip_start_steps} steps.")
|
||||
pipeline.transformer.enable_teacache(
|
||||
coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload
|
||||
)
|
||||
if transformer_2 is not None:
|
||||
pipeline.transformer_2.share_teacache(transformer=pipeline.transformer)
|
||||
|
||||
if cfg_skip_ratio is not None:
|
||||
print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.")
|
||||
pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps)
|
||||
if transformer_2 is not None:
|
||||
pipeline.transformer_2.share_cfg_skip(transformer=pipeline.transformer)
|
||||
|
||||
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)
|
||||
if lora_high_path is not None and transformer_2 is not None:
|
||||
pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2")
|
||||
|
||||
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)
|
||||
if transformer_2 is not None:
|
||||
pipeline.transformer_2.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,
|
||||
boundary = boundary,
|
||||
shift = shift,
|
||||
).videos
|
||||
|
||||
if lora_path is not None:
|
||||
pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
|
||||
if lora_high_path is not None and transformer_2 is not None:
|
||||
pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2")
|
||||
|
||||
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)
|
||||
|
||||
if ulysses_degree * ring_degree > 1:
|
||||
import torch.distributed as dist
|
||||
if dist.get_rank() == 0:
|
||||
save_results()
|
||||
else:
|
||||
save_results()
|
||||
@@ -67,12 +67,13 @@ sampler_name = "Flow"
|
||||
# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics.
|
||||
# Used when the sampler is in "Flow_Unipc", "Flow_DPM++".
|
||||
shift = 5
|
||||
stochastic_sampling = False
|
||||
|
||||
# Load pretrained model if need
|
||||
# Any Wan2.1 / CausVid / Self-Forcing checkpoint loads as-is: the Flex-Forcing
|
||||
# backbone inherits every parameter name and only the new `flex_kproj.*` tensors
|
||||
# are reported missing (they are identity-initialised, so step 0 is unchanged).
|
||||
transformer_path = "output_dir_wan2.1_flex_forcing_distill/checkpoint-3000/diffusion_pytorch_model.safetensors"
|
||||
transformer_path = "output_dir_wan2.1_flex_forcing_distill/checkpoint-10000/diffusion_pytorch_model.safetensors"
|
||||
vae_path = None
|
||||
lora_path = None
|
||||
|
||||
@@ -252,7 +253,7 @@ apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
|
||||
print(f"[Flex-Forcing] denoise_mode={denoise_mode}, "
|
||||
f"min_num_frame_per_block={min_num_frame_per_block}, "
|
||||
f"local_attn_size={local_attn_size}, sink_size={sink_size}")
|
||||
print(f"[Flex-Forcing] partitions measured in the paper: "
|
||||
print(f"[Flex-Forcing] partitions evaluated in the paper: "
|
||||
f"{[list(c) for c in PAPER_CHUNK_CONFIGS]}")
|
||||
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
@@ -281,6 +282,7 @@ with torch.no_grad():
|
||||
context_noise = context_noise,
|
||||
denoise_mode = denoise_mode,
|
||||
min_num_frame_per_block = min_num_frame_per_block,
|
||||
stochastic_sampling = stochastic_sampling,
|
||||
).videos
|
||||
torch.cuda.synchronize()
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from itertools import filterfalse
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -67,12 +68,13 @@ sampler_name = "Flow"
|
||||
# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics.
|
||||
# Used when the sampler is in "Flow_Unipc", "Flow_DPM++".
|
||||
shift = 5
|
||||
stochastic_sampling = False
|
||||
|
||||
# Load pretrained model if need
|
||||
# Any Wan2.1 / CausVid / Self-Forcing checkpoint loads as-is: the Flex-Forcing
|
||||
# backbone inherits every parameter name and only the new `flex_kproj.*` tensors
|
||||
# are reported missing (they are identity-initialised, so step 0 is unchanged).
|
||||
transformer_path = "output_dir_wan2.1_flex_forcing_distill/checkpoint-1000/diffusion_pytorch_model.safetensors"
|
||||
transformer_path = "output_dir_wan2.1_flex_forcing_distill/checkpoint-3000/diffusion_pytorch_model.safetensors"
|
||||
vae_path = None
|
||||
lora_path = None
|
||||
|
||||
@@ -292,15 +294,13 @@ with torch.no_grad():
|
||||
edit_steps = edit_steps,
|
||||
shift = shift,
|
||||
context_noise = context_noise,
|
||||
stochastic_sampling = stochastic_sampling,
|
||||
num_frame_per_block = num_frame_per_block,
|
||||
generator = generator,
|
||||
).videos
|
||||
torch.cuda.synchronize()
|
||||
elapsed = time.time() - start_time
|
||||
print(f"[Timing] edited span {edit_span} in {elapsed:.2f}s")
|
||||
# Same diagnostic as predict_t2v.py, read after the edit: the cache now holds
|
||||
# the whole clip committed as clean context, so `kv_tokens` is the full-clip
|
||||
# width the edited span was able to attend over.
|
||||
if getattr(pipeline, "kv_cache_pos", None) is not None:
|
||||
kv_tokens = pipeline.kv_cache_pos[0]["k"].shape[1]
|
||||
kv_mib = sum(c["k"].numel() + c["v"].numel()
|
||||
@@ -325,7 +325,6 @@ def save_results():
|
||||
image = (image * 255).numpy().astype(np.uint8)
|
||||
image = Image.fromarray(image)
|
||||
image.save(image_path)
|
||||
print(f"Saved image to: {image_path}")
|
||||
else:
|
||||
video_path = os.path.join(save_path, prefix + "-edited.mp4")
|
||||
save_videos_grid(sample, video_path, fps=fps)
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -0,0 +1,452 @@
|
||||
# SparkDiffusion Sparse Distillation Training Guide (Wan2.1-T2V-1.3B)
|
||||
|
||||
This document provides the complete workflow for **SparkDiffusion** on Wan2.1-T2V-1.3B: a two-stage pipeline that first warms up a **sparse (RoLa) student** and then **distills** it into a few-step, CFG-free generator. It covers environment setup, data preparation, the two launchers (Stage-1 warm-up + Stage-2 DMD distill), and sparse few-step inference. Wan2.1-T2V-1.3B is a **single full-range expert**, so this is the leaner counterpart of the Wan2.2 MoE guide ([README_TRAIN_DISTILL_wan2.2.md](./README_TRAIN_DISTILL_wan2.2.md)): no MoE boundary, no high/low expert split, and no init-rollout relay.
|
||||
|
||||
> **What is SparkDiffusion?**
|
||||
>
|
||||
> SparkDiffusion (arXiv:2609.23153, *Mitigating the High-Sparsity Trap*) combines a **compensating sparse attention** (RoLa) with a **DMD** few-step distillation recipe. Wan2.1-T2V-1.3B is a single-expert diffusion transformer that spans the whole noise range, so the student is warmed up and distilled **once over the full rf-time range `[0, 1]`** — there is no MoE boundary to split on.
|
||||
>
|
||||
> - **Stage 1 — sparse warm-up** ([`train_finetune_sparse_wan2.1.py`](./train_finetune_sparse_wan2.1.py)): insert the RoLa sparse self-attention into the dense backbone and briefly fine-tune it with the plain step-local flow-matching objective, so distillation starts from a usable coarse prior instead of a zero-initialized sparse branch. This stage is intentionally **short**.
|
||||
> - **Stage 2 — DMD distill** ([`train_distill_wan2.1.py`](./train_distill_wan2.1.py)): distill the sparse student into a few-step generator against the **dense** teacher. It is plain Distribution Matching Distillation (DMD) with three co-resident transformers — a trainable sparse `generator`, a frozen dense `real_score` teacher, and a trainable dense `fake_score` critic — rolled out over a discrete few-step grid. An optional **CrossDistill** PCM diversity term (`--loss_scale_div > 0`, see §3.3.1) can be layered on top to preserve sample diversity.
|
||||
>
|
||||
> | Stage | Launcher | Output |
|
||||
> |-------|----------|--------|
|
||||
> | **1 warm-up** | [`train_finetune_sparse_wan2.1.sh`](./train_finetune_sparse_wan2.1.sh) | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola` |
|
||||
> | **2 distill** | [`train_distill_wan2.1.sh`](./train_distill_wan2.1.sh) | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill` |
|
||||
>
|
||||
> Run the two launchers in order: the Stage-2 generator is seeded from the Stage-1 warm-up product, and the single distilled checkpoint is the whole few-step student at inference.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
- [1. Environment Setup](#1-environment-setup)
|
||||
- [2. Data Preparation](#2-data-preparation)
|
||||
- [3. Training](#3-training)
|
||||
- [3.1 Download Pretrained Models](#31-download-pretrained-models)
|
||||
- [3.2 Stage 1 — Sparse Warm-up](#32-stage-1--sparse-warm-up)
|
||||
- [3.3 Stage 2 — DMD Distill](#33-stage-2--dmd-distill)
|
||||
- [3.4 Common Training Parameters](#34-common-training-parameters)
|
||||
- [3.5 FSDP Backend and Memory Notes](#35-fsdp-backend-and-memory-notes)
|
||||
- [3.6 Other Backends](#36-other-backends)
|
||||
- [3.7 Multi-Node Distributed Training](#37-multi-node-distributed-training)
|
||||
- [4. Inference — Sparse Few-Step Student](#4-inference--sparse-few-step-student)
|
||||
- [5. Additional Resources](#5-additional-resources)
|
||||
|
||||
---
|
||||
|
||||
## 1. Environment Setup
|
||||
|
||||
**Method 1: Using requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**Method 2: Manual Installation**
|
||||
|
||||
```bash
|
||||
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
|
||||
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
|
||||
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
|
||||
pip install yunchang xfuser modelscope openpyxl
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
pip install deepspeed==0.17.0 numpy==1.26.4
|
||||
```
|
||||
|
||||
> 💡 The RoLa sparse attention relies on the in-repo Triton block-sparse kernel (`videox_fun/models/attention_utils.py`), so a working Triton install is required for the sparse student.
|
||||
|
||||
**Method 3: Using Docker**
|
||||
|
||||
```bash
|
||||
# pull image
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
|
||||
# enter image
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Data Preparation
|
||||
|
||||
SparkDiffusion consumes ordinary video-caption data (the same format as the standard Wan2.1 training).
|
||||
|
||||
### 2.1 Quick Test Dataset
|
||||
|
||||
```bash
|
||||
# Download official example dataset
|
||||
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
|
||||
```
|
||||
|
||||
### 2.2 metadata.json Format
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"file_path": "train/video001.mp4",
|
||||
"text": "A beautiful sunset over the ocean, golden hour lighting",
|
||||
"type": "video",
|
||||
"width": 1024,
|
||||
"height": 1024
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
- `file_path`: video path (relative to `DATASET_NAME`, or absolute).
|
||||
- `text`: caption / prompt.
|
||||
- `type`: fixed as `"video"`.
|
||||
- `width` / `height`: recommended for bucket training; add them with `scripts/process_json_add_width_and_height.py` if absent.
|
||||
|
||||
### 2.3 Relative vs Absolute Paths
|
||||
|
||||
```bash
|
||||
# relative paths
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
|
||||
# absolute paths (data on NAS/OSS shared across machines)
|
||||
export DATASET_NAME=""
|
||||
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. Training
|
||||
|
||||
The full order is **Stage-1 warm-up → Stage-2 DMD distill**. Run the two launchers by hand, in that order.
|
||||
|
||||
### 3.1 Download Pretrained Models
|
||||
|
||||
```bash
|
||||
# Create model directory
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# Download Wan2.1 T2V official weights
|
||||
modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
|
||||
```
|
||||
|
||||
The base checkpoint seeds all three Stage-2 transformers (the sparse `generator` is additionally initialized from the Stage-1 warm-up product). Its layout is the civitai single-file form read by `config/wan2.1/wan_civitai.yaml`: the transformer sits at the model root (`transformer_subpath: ./`), alongside `Wan2.1_VAE.pth`, the UMT5 encoder `models_t5_umt5-xxl-enc-bf16.pth` and the `google/umt5-xxl` tokenizer. The scheduler `shift` is `5.0`, supplied by the config `scheduler_kwargs`; the distill launcher does not override it.
|
||||
|
||||
### 3.2 Stage 1 — Sparse Warm-up
|
||||
|
||||
Stage 1 inserts the RoLa sparse self-attention and briefly warms it up with the step-local flow-matching objective. It is a **short** stage — a brief warm-up is enough to establish a usable prior for distillation, so the launcher caps it at `--max_train_steps=2000`. Do **not** judge this stage by its own sample quality; the sharpness is established in Stage 2.
|
||||
|
||||
Run [scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh](./train_finetune_sparse_wan2.1.sh):
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.1.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--max_train_steps=2000 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=1e-4 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--random_hw_adapt \
|
||||
--enable_bucket \
|
||||
--uniform_sampling \
|
||||
--trainable_modules "attn_op" \
|
||||
--student_topk=0.10 \
|
||||
--train_mode="normal" \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
Or simply:
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh
|
||||
```
|
||||
|
||||
Stage-1 keys:
|
||||
|
||||
| Flag | Meaning |
|
||||
|------|---------|
|
||||
| `--student_topk` | Ratio of key blocks kept per query in the sparse attention (sparsity = `1 - topk`). `0.10` ≈ 90% sparse. **Must match** the value used in Stage 2. |
|
||||
| `--trainable_modules` | `"attn_op"` = Stage A (train only the sparse branch); `"."` = Stage B (full sparse finetune, load Stage A via `--transformer_path`). |
|
||||
| `--max_train_steps` | Warm-up budget; keep it short. |
|
||||
|
||||
On completion the warm-up product is `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`.
|
||||
|
||||
### 3.3 Stage 2 — DMD Distill
|
||||
|
||||
Stage 2 distills the sparse student into a few-step, CFG-free generator against the dense teacher. `--sparse_student` turns the `generator` into `SparkWanTransformer3DModel` while the teacher and critic stay dense. `--generator_transformer_path` loads the Stage-1 warm-up checkpoint. Being a single full-range expert, the student distills over the whole rf-time range `[0, 1]` — there is no boundary segment and no init-rollout relay. This is the plain-DMD baseline: the generator rolls out a discrete few-step grid (`--denoising_step_indices_list 1000 750 500 250` with `--flow_euler_rollout`), and the generator/critic alternate via `--gen_update_interval`. On top of this discrete DMD host, an optional CrossDistill PCM diversity term (`--loss_scale_div > 0`, see §3.3.1 below) adds a phase-0 boundary-alignment loss to preserve sample diversity. Two further opt-in switches (§3.3.2) replace parts of this baseline — `--tdm` (the official SparkDiffusion TDM student loss) and `--joint_backward_timesteps` (define the few-step grid directly with rf knots) — but neither is enabled in the launcher, so the default run stays plain DMD.
|
||||
|
||||
Run [scripts/sparkdiffusion/train_distill_wan2.1.sh](./train_distill_wan2.1.sh):
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# Stage-1 warm-up product. Env-overridable; --student_topk below MUST match Stage 1.
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.1.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--random_hw_adapt \
|
||||
--enable_bucket \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--uniform_sampling \
|
||||
--denoising_step_indices_list 1000 750 500 250 \
|
||||
--flow_euler_rollout \
|
||||
--randomize_step_indices \
|
||||
--gen_update_interval 5 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 5.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 5.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--sparse_student \
|
||||
--student_topk=0.10 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
Or simply:
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_distill_wan2.1.sh
|
||||
```
|
||||
|
||||
> ⚠️ **The distillation phase has no step cap.** It runs until you stop it manually (kill the `train_distill_wan2.1.py` processes); checkpoints are written every `--checkpointing_steps=200`.
|
||||
|
||||
#### 3.3.1 CrossDistill Diversity (PCM)
|
||||
|
||||
Setting `--loss_scale_div > 0` layers the official SparkDiffusion **CrossDistill** PCM diversity term on top of the discrete DMD host, to counter the diversity loss that pure distribution matching can cause. It is a self-contained additive loss that leaves the DMD rollout untouched:
|
||||
|
||||
- **What it does** — a phase-0 boundary-alignment MSE: the student's one-step jump from a noisy anchor to the phase-0 boundary (`--div_rf_boundary`, the official `backward_timesteps[0]`) must match a teacher K-step Euler rollout (`--dcm_skipping_interval_steps` sub-steps) followed by the student's own target jump. It runs on its **own continuous rf grid** (`--dcm_total_steps` subdivisions shifted by `--teacher_timestep_shift`), decoupled from the discrete DMD grid.
|
||||
- **Anchor** — the paired **real video latent** (faithful to the official `_compute_pcm_div_loss`, which distills on real `x0`), encoded on demand only when diversity is on.
|
||||
- **Merge** — `generator_loss = dmd_loss + loss_scale_div * div_loss` in a **single** `accelerator.backward` (FSDP allows one backward per wrapped module per iteration); the teacher rollout and the student target run under `no_grad` before the graded student forward.
|
||||
- **Scale & tuning** — `--loss_scale_div` is in **per-element-mean** space (aligned with the mean-reduced `dmd_loss`); the official uses `10.0` in sum space, so `5.0` is the recommended start. On the first run watch `train_div_loss` vs `train_dmd_loss` in the logs and adjust so diversity does not overwhelm DMD; `--div_teacher_guidance` is the diversity teacher's CFG scale, independent of `--real_guidance_scale`. Set `--loss_scale_div 0` to recover the exact plain-DMD baseline.
|
||||
|
||||
#### 3.3.2 TDM Loss & Direct rf Grid (opt-in)
|
||||
|
||||
Two opt-in switches replace parts of the plain-DMD baseline above; **neither is enabled in `train_distill_wan2.1.sh`**, so the default run stays plain DMD.
|
||||
|
||||
- **`--tdm`** — replaces the plain-DMD student loss with **TDM** (Trajectory Distribution Matching, the official SparkDiffusion student loss). Three deltas from DMD: (1) the last rollout segment does not collapse to `x0` but takes one more rf-Euler step to `s_{i-1}` (the next grid knot, `0.0` when exhausted), so the anchor is `G_x_{t_{i-1}}`; (2) `D_time` is clamped strictly above `s_{i-1}`; (3) the fake/teacher `x0` preds are mapped back to `s_{i-1}` via the rf PF-ODE analytic step and the surrogate is compared in the `x_{t_{i-1}}` domain. When `s_{i-1}==0` (grid exhausted) it reduces **exactly** to DMD. Requires `--flow_euler_rollout`; mutually exclusive with `--dfd` / `--decoupled_dmd`; compatible with the diversity term.
|
||||
- **`--joint_backward_timesteps`** — descending rf knots (in `[0,1]`) that define the few-step grid **directly**, replacing `--denoising_step_indices_list`. Wan2.1 is a single expert (no MoE boundary split), so the whole list is used: the grid becomes `[1.0] + [k for k in knots if 0<k<1]` and the rf values serve directly as Euler sigmas (`sigma == rf == timestep/1000`). Official Wan2.1 rCM knots: `0.933781 0.852895 0.608979`. When omitted, the grid falls back to `--denoising_step_indices_list`.
|
||||
|
||||
### 3.4 Common Training Parameters
|
||||
|
||||
| Parameter | Description | Example Value |
|
||||
|-----------|-------------|---------------|
|
||||
| `--config_path` | Wan2.1 config (single expert, no `boundary`) | `config/wan2.1/wan_civitai.yaml` |
|
||||
| `--pretrained_model_name_or_path` | Wan2.1-T2V-1.3B base (transformer + VAE + UMT5) | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
|
||||
| `--train_data_dir` / `--train_data_meta` | Dataset root / metadata JSON | see §2 |
|
||||
| `--train_mode` | SparkDiffusion supports `normal` (T2V) only | `normal` |
|
||||
| `--student_topk` | Sparse attention keep-ratio (sparsity = 1 − topk); Stage 1 and Stage 2 must match | `0.10` |
|
||||
| `--sparse_student` | (Stage 2) build the generator as the sparse student | - |
|
||||
| `--generator_transformer_path` | (Stage 2) Stage-1 warm-up checkpoint loaded into the generator | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` |
|
||||
| `--denoising_step_indices_list` | Discrete few-step grid the generator rolls out; validation reuses it (`num_inference_steps` = its length) | `1000 750 500 250` |
|
||||
| `--flow_euler_rollout` | Roll the generator in flow/velocity space with deterministic Euler steps (LightX2V-style), no x0 re-noising | - |
|
||||
| `--tdm` | (opt-in) Use TDM — the official SparkDiffusion student loss — instead of plain DMD; reduces exactly to DMD when `s_{i-1}==0`. Requires `--flow_euler_rollout`, mutually exclusive with `--dfd` / `--decoupled_dmd` | - |
|
||||
| `--joint_backward_timesteps` | (opt-in) Descending rf knots that define the few-step grid directly, replacing `--denoising_step_indices_list` (single expert: the whole list, no MoE split) | `0.933781 0.852895 0.608979` |
|
||||
| `--uniform_sampling` | Sample the DMD noise level uniformly over the training schedule | - |
|
||||
| `--randomize_step_indices` | Jitter the grid indices each step for robustness | - |
|
||||
| `--gen_update_interval` | Update the generator once every N steps (the critic updates every step) | `5` |
|
||||
| `--real_guidance_scale` | CFG scale of the frozen dense teacher (not overridden; uses the script default) | `6.0` |
|
||||
| `--div_mode` | CrossDistill diversity mode: `pcm_random` (random phase-0 offset) or `pcm_fixed` | `pcm_random` |
|
||||
| `--loss_scale_div` | Weight of the CrossDistill PCM diversity term (per-element-mean space); `0` disables it and recovers plain DMD | `5.0` |
|
||||
| `--dcm_total_steps` | Subdivision of the diversity term's continuous rf grid (`dw = 1 / N`) | `48` |
|
||||
| `--dcm_skipping_interval_steps` | Teacher Euler sub-steps `K` within phase 0 | `1` |
|
||||
| `--teacher_timestep_shift` | rf timestep shift for the diversity grid (matches the Wan2.1 scheduler shift) | `5.0` |
|
||||
| `--div_rf_boundary` | Phase-0 right boundary `rf_b` (official `backward_timesteps[0]`) | `0.933781` |
|
||||
| `--div_teacher_guidance` | CFG scale of the diversity teacher rollout (independent of `--real_guidance_scale`) | `5.0` |
|
||||
| `--learning_rate` / `--learning_rate_critic` | Generator / critic learning rate | `2e-06` / `2e-06` |
|
||||
| `--checkpointing_steps` | Save checkpoint every N steps | `200` |
|
||||
| `--max_grad_norm` | Gradient clipping threshold | `0.05` |
|
||||
| `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` |
|
||||
|
||||
> 💡 The few-step grid is set by `--denoising_step_indices_list` (indices into the `--train_sampling_steps=1000` schedule); validation reuses the same grid, so `num_inference_steps` equals its length (`4` here).
|
||||
|
||||
### 3.5 FSDP Backend and Memory Notes
|
||||
|
||||
Stage 2 keeps **three** transformers co-resident (sparse generator + dense teacher + dense critic). At 1.3B each this is far lighter than the Wan2.2 14B experts, but the optimizer states still dominate, so the distill launcher keeps FSDP `FULL_SHARD` plus `--gradient_checkpointing`:
|
||||
|
||||
- `--fsdp_sharding_strategy "FULL_SHARD"` — shard params/grads/optimizer across GPUs.
|
||||
- `--fsdp_backward_prefetch "BACKWARD_POST"` — lower median residency during backward.
|
||||
- `--gradient_checkpointing` — recompute activations in backward.
|
||||
- **Resolution, not frame count, is the OOM lever.** Memory scales with `seq_len`; the launchers bucket at `--image_sample_size=640 --video_sample_size=640 --token_sample_size=640` with `--random_hw_adapt --enable_bucket` at `--video_sample_n_frames=81`. If it OOMs, lower the sample sizes (e.g. `640 → 480`) rather than dropping frames.
|
||||
- Frame count must satisfy `(N − 1) % 4 == 0`.
|
||||
|
||||
> 💡 At 1.3B the three transformers usually fit without CPU offload; if it OOMs, add `--fsdp_offload_params true` to the `accelerate launch` line to offload params/grads/optimizer to CPU (slower). `--low_vram` only parks the frozen side on CPU; it complements, but does not replace, sharding.
|
||||
|
||||
### 3.6 Other Backends
|
||||
|
||||
The Stage-2 launcher defaults to FSDP `FULL_SHARD`. The same [`train_distill_wan2.1.py`](./train_distill_wan2.1.py) also runs under DeepSpeed Zero-2 or without a sharding backend; at 1.3B these are more viable than for the Wan2.2 14B experts. Keep every script argument identical to §3.3 — only the `accelerate launch` header changes.
|
||||
|
||||
#### 3.6.1 Training with DeepSpeed Zero-2
|
||||
|
||||
DeepSpeed Zero-2 shards optimizer states and gradients but **not** parameters. For three co-resident 1.3B transformers it fits comfortably at the default resolution:
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/sparkdiffusion/train_distill_wan2.1.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--random_hw_adapt \
|
||||
--enable_bucket \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--uniform_sampling \
|
||||
--denoising_step_indices_list 1000 750 500 250 \
|
||||
--flow_euler_rollout \
|
||||
--randomize_step_indices \
|
||||
--gen_update_interval 5 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 5.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 5.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--sparse_student \
|
||||
--student_topk=0.10 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
#### 3.6.2 Training Without DeepSpeed and FSDP
|
||||
|
||||
Drop the sharding backend entirely — three co-resident 1.3B transformers fit on a single high-memory GPU, though FSDP is still recommended for multi-GPU. Replace the `accelerate launch` header with the plain form and keep the same script arguments as §3.6.1:
|
||||
|
||||
```bash
|
||||
accelerate launch --mixed_precision="bf16" \
|
||||
scripts/sparkdiffusion/train_distill_wan2.1.py \
|
||||
<the same --flag block as §3.6.1>
|
||||
```
|
||||
|
||||
### 3.7 Multi-Node Distributed Training
|
||||
|
||||
The launchers are single-node. For multiple machines, add the following to the `accelerate launch` line (all machines must share the same dataset / base weights / `output_dir` over NFS/shared storage):
|
||||
|
||||
```bash
|
||||
--main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT \
|
||||
--num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK
|
||||
```
|
||||
|
||||
Without RDMA, also set `NCCL_IB_DISABLE=1` and `NCCL_P2P_DISABLE=1`.
|
||||
|
||||
---
|
||||
|
||||
## 4. Inference — Sparse Few-Step Student
|
||||
|
||||
The Stage-2 product is a few-step, CFG-free sparse student. Being single-expert, inference loads **one** checkpoint (there is no high-noise expert). Build the wan2.1 sparse class `SparkWanTransformer3DModel`, set the training top-k, then load the distilled weights — the order matters because `set_student_topk` (re)creates the `attn_op` kernels that receive the trained RoLa params:
|
||||
|
||||
```python
|
||||
from videox_fun.models import SparkWanTransformer3DModel
|
||||
transformer = SparkWanTransformer3DModel.from_pretrained(
|
||||
os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
|
||||
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
||||
low_cpu_mem_usage=True, torch_dtype=weight_dtype,
|
||||
)
|
||||
transformer.set_student_topk(student_topk) # BEFORE load_state_dict
|
||||
# transformer.load_state_dict(load_file(transformer_path), strict=False)
|
||||
```
|
||||
|
||||
The Wan2.1 T2V entry [examples/wan2.1/predict_t2v.py](../../examples/wan2.1/predict_t2v.py) already runs few-step CFG-free sampling (`sampler_name="Flow"`, `num_inference_steps=4`, `guidance_scale=1.0`) under `WanPipeline`; adapt it by swapping the dense `WanTransformer3DModel` for the sparse build above. Set:
|
||||
|
||||
| Parameter | Description | Example Value |
|
||||
|-----------|-------------|---------------|
|
||||
| `GPU_memory_mode` | Memory mode | `model_full_load` |
|
||||
| `config_path` | Wan2.1 config | `config/wan2.1/wan_civitai.yaml` |
|
||||
| `model_name` | Wan2.1-T2V-1.3B base (transformer / VAE / UMT5) | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
|
||||
| `student_topk` | **Must equal** the training `--student_topk` | `0.10` |
|
||||
| `transformer_path` | Sparse student checkpoint (single expert) | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
|
||||
| `sampler_name` | `Flow` matches the few-step CFG-free student | `Flow` |
|
||||
| `num_inference_steps` | The student's few-step count | `4` |
|
||||
| `guidance_scale` | Distilled student is CFG-free | `1.0` |
|
||||
| `sample_size` / `video_length` / `fps` | `[height, width]` / frames / fps | `[480, 832]` / `81` / `16` |
|
||||
|
||||
```bash
|
||||
python examples/wan2.1/predict_t2v.py
|
||||
```
|
||||
|
||||
> 💡 Load order is `from_pretrained → set_student_topk → load_state_dict`; reversing it re-zeroes the trained RoLa weights. A sparse checkpoint is confirmed when the `attn_op.*` keys load with **0** unexpected keys.
|
||||
|
||||
> 📌 The bundled [examples/sparkdiffusion/predict_t2v.py](../../examples/sparkdiffusion/predict_t2v.py) targets the Wan2.2 two-expert MoE student (`SparkWan2_2Transformer3DModel` + `Wan2_2Pipeline`); the single-expert Wan2.1 student uses `SparkWanTransformer3DModel` + `WanPipeline` with one checkpoint.
|
||||
|
||||
> 📌 **Stage 3** of the paper (the FP8/INT8 fused quantized inference kernel) is intentionally not integrated in this repository.
|
||||
|
||||
---
|
||||
|
||||
## 5. Additional Resources
|
||||
|
||||
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
|
||||
- **Paper**: SparkDiffusion, *Mitigating the High-Sparsity Trap* (arXiv:2609.23153)
|
||||
- **Sparse student model**: `videox_fun/models/spark_transformer3d.py` (`SparkWanTransformer3DModel`)
|
||||
- **RoLa sparse attention kernel**: `videox_fun/models/attention_utils.py`
|
||||
- **Wan2.2 MoE counterpart**: [README_TRAIN_DISTILL_wan2.2.md](./README_TRAIN_DISTILL_wan2.2.md)
|
||||
@@ -0,0 +1,452 @@
|
||||
# SparkDiffusion 稀疏蒸馏训练指南（Wan2.1-T2V-1.3B）
|
||||
|
||||
本文档给出 Wan2.1-T2V-1.3B 上 **SparkDiffusion** 的完整流程：一条两阶段管线，先暖启一个**稀疏（RoLa）学生**，再把它**蒸馏**成少步、CFG-free 的生成器。内容涵盖环境配置、数据准备、两个启动脚本（Stage-1 暖启 + Stage-2 DMD 蒸馏）、以及稀疏少步推理。Wan2.1-T2V-1.3B 是**单专家、覆盖全噪声区间**的模型，因此本文是 Wan2.2 MoE 指南（[README_TRAIN_DISTILL_wan2.2_zh-CN.md](./README_TRAIN_DISTILL_wan2.2_zh-CN.md)）的精简版：没有 MoE 分界、没有高/低噪专家拆分、也没有 init-rollout 接力。
|
||||
|
||||
> **什么是 SparkDiffusion？**
|
||||
>
|
||||
> SparkDiffusion（arXiv:2609.23153，《Mitigating the High-Sparsity Trap》）把**补偿式稀疏注意力**（RoLa）与 **DMD** 少步蒸馏配方结合起来。Wan2.1-T2V-1.3B 是覆盖整个噪声区间的单专家 diffusion transformer，因此学生在**整个 rf-time 区间 `[0, 1]` 上一次性**完成暖启与蒸馏——没有可供拆分的 MoE 分界。
|
||||
>
|
||||
> - **Stage 1 — 稀疏暖启**（[`train_finetune_sparse_wan2.1.py`](./train_finetune_sparse_wan2.1.py)）：把 RoLa 稀疏自注意力插入 dense 主干，用标准的 step-local flow-matching 目标短暂微调，使蒸馏从一个可用的粗先验出发，而不是从零初始化的稀疏分支出发。该阶段刻意保持**短**。
|
||||
> - **Stage 2 — DMD 蒸馏**（[`train_distill_wan2.1.py`](./train_distill_wan2.1.py)）：对着 **dense** teacher 把稀疏学生蒸馏成少步生成器。它是最普通的分布匹配蒸馏（DMD），含三个同驻网络——可训练的稀疏 `generator`、冻结的 dense `real_score` teacher、可训练的 dense `fake_score` critic——在离散的少步网格上 rollout。其上可叠加一个可选的 **CrossDistill** PCM 多样性项（`--loss_scale_div > 0`，见 §3.3.1），用于保持样本多样性。
|
||||
>
|
||||
> | 阶段 | 启动脚本 | 产物 |
|
||||
> |-------|----------|--------|
|
||||
> | **1 暖启** | [`train_finetune_sparse_wan2.1.sh`](./train_finetune_sparse_wan2.1.sh) | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola` |
|
||||
> | **2 蒸馏** | [`train_distill_wan2.1.sh`](./train_distill_wan2.1.sh) | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill` |
|
||||
>
|
||||
> 按顺序运行两个启动脚本：Stage-2 的 generator 由 Stage-1 暖启产物播种，而单个蒸馏 checkpoint 在推理时即为完整的少步学生。
|
||||
|
||||
---
|
||||
|
||||
## 目录
|
||||
- [一、环境配置](#一环境配置)
|
||||
- [二、数据准备](#二数据准备)
|
||||
- [三、训练](#三训练)
|
||||
- [3.1 下载预训练模型](#31-下载预训练模型)
|
||||
- [3.2 Stage 1 — 稀疏暖启](#32-stage-1--稀疏暖启)
|
||||
- [3.3 Stage 2 — DMD 蒸馏](#33-stage-2--dmd-蒸馏)
|
||||
- [3.4 通用训练参数](#34-通用训练参数)
|
||||
- [3.5 FSDP 后端与显存说明](#35-fsdp-后端与显存说明)
|
||||
- [3.6 其他后端](#36-其他后端)
|
||||
- [3.7 多机分布式训练](#37-多机分布式训练)
|
||||
- [四、推理 — 稀疏少步学生](#四推理--稀疏少步学生)
|
||||
- [五、更多资源](#五更多资源)
|
||||
|
||||
---
|
||||
|
||||
## 一、环境配置
|
||||
|
||||
**方式一：使用 requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**方式二：手动安装**
|
||||
|
||||
```bash
|
||||
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
|
||||
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
|
||||
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
|
||||
pip install yunchang xfuser modelscope openpyxl
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
pip install deepspeed==0.17.0 numpy==1.26.4
|
||||
```
|
||||
|
||||
> 💡 RoLa 稀疏注意力依赖仓库内的 Triton block-sparse 内核（`videox_fun/models/attention_utils.py`），因此稀疏学生需要一个可用的 Triton 安装。
|
||||
|
||||
**方式三：使用 Docker**
|
||||
|
||||
```bash
|
||||
# 拉取镜像
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
|
||||
# 进入镜像
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 二、数据准备
|
||||
|
||||
SparkDiffusion 使用普通的「视频-字幕」数据（与标准 Wan2.1 训练格式一致）。
|
||||
|
||||
### 2.1 快速测试数据集
|
||||
|
||||
```bash
|
||||
# 下载官方示例数据集
|
||||
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
|
||||
```
|
||||
|
||||
### 2.2 metadata.json 格式
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"file_path": "train/video001.mp4",
|
||||
"text": "A beautiful sunset over the ocean, golden hour lighting",
|
||||
"type": "video",
|
||||
"width": 1024,
|
||||
"height": 1024
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
- `file_path`：视频路径（相对于 `DATASET_NAME`，或绝对路径）。
|
||||
- `text`：字幕 / 提示词。
|
||||
- `type`：固定为 `"video"`。
|
||||
- `width` / `height`：建议提供以启用 bucket 训练；若缺失可用 `scripts/process_json_add_width_and_height.py` 补齐。
|
||||
|
||||
### 2.3 相对路径与绝对路径
|
||||
|
||||
```bash
|
||||
# 相对路径
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
|
||||
# 绝对路径（数据在 NAS/OSS 上、跨机共享）
|
||||
export DATASET_NAME=""
|
||||
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 三、训练
|
||||
|
||||
完整顺序为 **Stage-1 暖启 → Stage-2 DMD 蒸馏**。请手动按此顺序运行两个启动脚本。
|
||||
|
||||
### 3.1 下载预训练模型
|
||||
|
||||
```bash
|
||||
# 创建模型目录
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# 下载 Wan2.1 T2V 官方权重
|
||||
modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
|
||||
```
|
||||
|
||||
基座权重会播种 Stage-2 的全部三个网络（稀疏 `generator` 还会额外用 Stage-1 暖启产物初始化）。其目录结构为 `config/wan2.1/wan_civitai.yaml` 读取的 civitai 单文件形态：transformer 位于模型根目录（`transformer_subpath: ./`），与 `Wan2.1_VAE.pth`、UMT5 编码器 `models_t5_umt5-xxl-enc-bf16.pth` 以及 `google/umt5-xxl` 分词器并列。调度器 `shift` 为 `5.0`，由 config 的 `scheduler_kwargs` 提供；蒸馏启动脚本不再覆盖它。
|
||||
|
||||
### 3.2 Stage 1 — 稀疏暖启
|
||||
|
||||
Stage 1 插入 RoLa 稀疏自注意力，并用 step-local flow-matching 目标短暂暖启。这是一个**短**阶段——短暂暖启即足以为蒸馏建立可用先验，因此启动脚本把它限制在 `--max_train_steps=2000`。**不要**用该阶段自身的采样质量来评判它；清晰度是在 Stage 2 建立的。
|
||||
|
||||
用 [scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh](./train_finetune_sparse_wan2.1.sh) 运行：
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.1.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--max_train_steps=2000 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=1e-4 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--random_hw_adapt \
|
||||
--enable_bucket \
|
||||
--uniform_sampling \
|
||||
--trainable_modules "attn_op" \
|
||||
--student_topk=0.10 \
|
||||
--train_mode="normal" \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
或直接：
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh
|
||||
```
|
||||
|
||||
Stage-1 关键参数：
|
||||
|
||||
| 旗标 | 含义 |
|
||||
|------|---------|
|
||||
| `--student_topk` | 稀疏注意力中每个 query 保留的 key block 比例（稀疏度 = `1 - topk`）。`0.10` ≈ 90% 稀疏。**必须与** Stage 2 一致。 |
|
||||
| `--trainable_modules` | `"attn_op"` = Stage A（只训练稀疏分支）；`"."` = Stage B（稀疏网全量微调，用 `--transformer_path` 载入 Stage A）。 |
|
||||
| `--max_train_steps` | 暖启预算；保持简短。 |
|
||||
|
||||
完成后暖启产物为 `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`。
|
||||
|
||||
### 3.3 Stage 2 — DMD 蒸馏
|
||||
|
||||
Stage 2 对着 dense teacher 把稀疏学生蒸馏成少步、CFG-free 的生成器。`--sparse_student` 把 `generator` 变为 `SparkWanTransformer3DModel`，而 teacher 与 critic 保持 dense。`--generator_transformer_path` 载入 Stage-1 暖启 checkpoint。作为单一全区间专家，学生在整个 rf-time 区间 `[0, 1]` 上完成蒸馏——没有分界段、也没有 init-rollout 接力。这是最普通的 DMD 基线：generator 在离散的少步网格上 rollout（`--denoising_step_indices_list 1000 750 500 250` 配 `--flow_euler_rollout`），并通过 `--gen_update_interval` 让 generator/critic 交替更新。在这个离散 DMD 宿主之上，可选的 CrossDistill PCM 多样性项（`--loss_scale_div > 0`，见下方 §3.3.1）额外引入一个 phase-0 边界对齐损失以保持样本多样性。另有两个可选开关（见 §3.3.2）会替换该基线的一部分：`--tdm`（官方 SparkDiffusion TDM 学生损失）与 `--joint_backward_timesteps`（用 rf 节点直接定义少步网格）；但启动脚本中两者都未启用，故默认运行仍是 plain DMD。
|
||||
|
||||
用 [scripts/sparkdiffusion/train_distill_wan2.1.sh](./train_distill_wan2.1.sh) 运行：
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# Stage-1 暖启产物。可用环境变量覆盖；下方 --student_topk 必须与 Stage 1 一致。
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.1.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--random_hw_adapt \
|
||||
--enable_bucket \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--uniform_sampling \
|
||||
--denoising_step_indices_list 1000 750 500 250 \
|
||||
--flow_euler_rollout \
|
||||
--randomize_step_indices \
|
||||
--gen_update_interval 5 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 5.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 5.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--sparse_student \
|
||||
--student_topk=0.10 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
或直接：
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_distill_wan2.1.sh
|
||||
```
|
||||
|
||||
> ⚠️ **蒸馏阶段没有步数上限。** 它会一直运行直到你手动停止（kill 掉 `train_distill_wan2.1.py` 进程）；checkpoint 每 `--checkpointing_steps=200` 步写一次。
|
||||
|
||||
#### 3.3.1 CrossDistill 多样性（PCM）
|
||||
|
||||
设置 `--loss_scale_div > 0` 会在离散 DMD 宿主之上叠加官方 SparkDiffusion 的 **CrossDistill** PCM 多样性项，用于抵消纯分布匹配可能带来的多样性损失。它是一个自包含的加性损失，不改动 DMD rollout：
|
||||
|
||||
- **作用**——一个 phase-0 边界对齐 MSE：学生从高噪 anchor 一步跳到 phase-0 边界（`--div_rf_boundary`，即官方 `backward_timesteps[0]`）的结果，必须与「teacher K 步 Euler rollout（`--dcm_skipping_interval_steps` 个子步）+ 学生自身的 target 跳转」相匹配。它跑在**自己的连续 rf 网格**上（`--dcm_total_steps` 细分、按 `--teacher_timestep_shift` 平移），与离散 DMD 网格解耦。
|
||||
- **anchor**——配对的**真实视频 latent**（忠实官方 `_compute_pcm_div_loss`，其在真实 `x0` 上蒸馏），仅在多样性开启时按需编码。
|
||||
- **合并**——`generator_loss = dmd_loss + loss_scale_div * div_loss`，用**单次** `accelerator.backward`（FSDP 每个 wrapped module 每轮只允许一次 backward）；teacher rollout 与学生 target 在带梯度的学生前向之前于 `no_grad` 下运行。
|
||||
- **scale 与调参**——`--loss_scale_div` 处于**逐元素均值**空间（与 mean 归约的 `dmd_loss` 对齐）；官方在 sum 空间用 `10.0`，故推荐起点为 `5.0`。首轮请观察日志中的 `train_div_loss` 与 `train_dmd_loss` 并调整，使多样性项不压过 DMD；`--div_teacher_guidance` 是多样性 teacher 的 CFG 尺度，独立于 `--real_guidance_scale`。设 `--loss_scale_div 0` 可退回完全一致的 plain-DMD 基线。
|
||||
|
||||
#### 3.3.2 TDM 损失与直接 rf 网格（可选）
|
||||
|
||||
两个可选开关会替换上方 plain-DMD 基线的一部分；**`train_distill_wan2.1.sh` 中两者都未启用**，故默认运行仍是 plain DMD。
|
||||
|
||||
- **`--tdm`**——把 plain-DMD 学生损失替换为 **TDM**（Trajectory Distribution Matching，官方 SparkDiffusion 学生损失）。相对 DMD 有三处差异：(1) rollout 末段不塌缩到 `x0`，而是再走一步 rf-Euler 到 `s_{i-1}`（下一个网格节点，耗尽时为 `0.0`），故锚点为 `G_x_{t_{i-1}}`；(2) `D_time` 被严格 clamp 在 `s_{i-1}` 之上；(3) fake/teacher 的 `x0` 预测经 rf PF-ODE 解析步映回 `s_{i-1}`，surrogate 在 `x_{t_{i-1}}` 域比较。当 `s_{i-1}==0`（网格耗尽）时**精确**退化为 DMD。需 `--flow_euler_rollout`；与 `--dfd` / `--decoupled_dmd` 互斥；与多样性项兼容。
|
||||
- **`--joint_backward_timesteps`**——降序 rf 节点（位于 `[0,1]`），**直接**定义少步网格，替代 `--denoising_step_indices_list`。Wan2.1 是单专家（不做 MoE 分界切分），故整份列表都用上：网格变为 `[1.0] + [k for k in knots if 0<k<1]`，rf 值直接充当 Euler sigma（`sigma == rf == timestep/1000`）。官方 Wan2.1 rCM 节点：`0.933781 0.852895 0.608979`。省略时，网格退回 `--denoising_step_indices_list`。
|
||||
|
||||
### 3.4 通用训练参数
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|-----------|-------------|---------------|
|
||||
| `--config_path` | Wan2.1 config（单专家，无 `boundary`） | `config/wan2.1/wan_civitai.yaml` |
|
||||
| `--pretrained_model_name_or_path` | Wan2.1-T2V-1.3B 基座（transformer + VAE + UMT5） | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
|
||||
| `--train_data_dir` / `--train_data_meta` | 数据集根目录 / metadata JSON | 见 §2 |
|
||||
| `--train_mode` | SparkDiffusion 仅支持 `normal`（T2V） | `normal` |
|
||||
| `--student_topk` | 稀疏注意力保留比例（稀疏度 = 1 − topk）；Stage 1 与 Stage 2 必须一致 | `0.10` |
|
||||
| `--sparse_student` | （Stage 2）把 generator 建成稀疏学生 | - |
|
||||
| `--generator_transformer_path` | （Stage 2）载入 generator 的 Stage-1 暖启 checkpoint | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` |
|
||||
| `--denoising_step_indices_list` | generator rollout 的离散少步网格；validation 复用同一网格（`num_inference_steps` = 其长度） | `1000 750 500 250` |
|
||||
| `--flow_euler_rollout` | 让 generator 在 flow/速度空间用确定性 Euler 步 rollout（LightX2V 式），不做 x0 重新加噪 | - |
|
||||
| `--tdm` | （可选）用 TDM（官方 SparkDiffusion 学生损失）替代 plain DMD；当 `s_{i-1}==0` 时精确退化为 DMD。需 `--flow_euler_rollout`，与 `--dfd` / `--decoupled_dmd` 互斥 | - |
|
||||
| `--joint_backward_timesteps` | （可选）降序 rf 节点，直接定义少步网格，替代 `--denoising_step_indices_list`（单专家：整份列表，不做 MoE 切分） | `0.933781 0.852895 0.608979` |
|
||||
| `--uniform_sampling` | 在训练调度上均匀采样 DMD 噪声档位 | - |
|
||||
| `--randomize_step_indices` | 每步抖动网格索引以增强鲁棒性 | - |
|
||||
| `--gen_update_interval` | 每 N 步更新一次 generator（critic 每步都更新） | `5` |
|
||||
| `--real_guidance_scale` | 冻结 dense teacher 的 CFG 尺度（不覆盖，走脚本默认值） | `6.0` |
|
||||
| `--div_mode` | CrossDistill 多样性模式：`pcm_random`（随机 phase-0 偏移）或 `pcm_fixed` | `pcm_random` |
|
||||
| `--loss_scale_div` | CrossDistill PCM 多样性项权重（逐元素均值空间）；`0` 关闭并退回 plain DMD | `5.0` |
|
||||
| `--dcm_total_steps` | 多样性项连续 rf 网格的细分数（`dw = 1 / N`） | `48` |
|
||||
| `--dcm_skipping_interval_steps` | phase 0 内 teacher Euler 子步数 `K` | `1` |
|
||||
| `--teacher_timestep_shift` | 多样性网格的 rf timestep 平移（与 Wan2.1 scheduler shift 一致） | `5.0` |
|
||||
| `--div_rf_boundary` | phase-0 右边界 `rf_b`（官方 `backward_timesteps[0]`） | `0.933781` |
|
||||
| `--div_teacher_guidance` | 多样性 teacher rollout 的 CFG 尺度（独立于 `--real_guidance_scale`） | `5.0` |
|
||||
| `--learning_rate` / `--learning_rate_critic` | 生成器 / critic 学习率 | `2e-06` / `2e-06` |
|
||||
| `--checkpointing_steps` | 每 N 步保存 checkpoint | `200` |
|
||||
| `--max_grad_norm` | 梯度裁剪阈值 | `0.05` |
|
||||
| `--resume_from_checkpoint` | 续训路径或 `"latest"` | `latest` |
|
||||
|
||||
> 💡 少步网格由 `--denoising_step_indices_list` 设定（是对 `--train_sampling_steps=1000` 调度的索引）；validation 复用同一网格，因此 `num_inference_steps` 等于其长度（此处为 `4`）。
|
||||
|
||||
### 3.5 FSDP 后端与显存说明
|
||||
|
||||
Stage 2 需让**三个** transformer 同驻（稀疏 generator + dense teacher + dense critic）。每个 1.3B 远比 Wan2.2 的 14B 专家轻量，但优化器态仍占主导，因此蒸馏启动脚本保留 FSDP `FULL_SHARD` 与 `--gradient_checkpointing`：
|
||||
|
||||
- `--fsdp_sharding_strategy "FULL_SHARD"`——在多卡间分片 param/grad/optimizer。
|
||||
- `--fsdp_backward_prefetch "BACKWARD_POST"`——降低 backward 期间的中位驻留。
|
||||
- `--gradient_checkpointing`——在 backward 中重算激活。
|
||||
- **OOM 的杠杆是分辨率、而非帧数。** 显存随 `seq_len` 增长；启动脚本以 `--image_sample_size=640 --video_sample_size=640 --token_sample_size=640` 配 `--random_hw_adapt --enable_bucket`、在 `--video_sample_n_frames=81` 下分桶。若 OOM，请降低 sample size（例如 `640 → 480`），而不是减少帧数。
|
||||
- 帧数须满足 `(N − 1) % 4 == 0`。
|
||||
|
||||
> 💡 1.3B 下三个 transformer 通常无需 CPU offload 即可放下；若 OOM，在 `accelerate launch` 行加上 `--fsdp_offload_params true` 把 param/grad/optimizer offload 到 CPU（更慢）。`--low_vram` 只把冻结侧停到 CPU；它是对分片的补充，而非替代。
|
||||
|
||||
### 3.6 其他后端
|
||||
|
||||
Stage-2 启动脚本默认使用 FSDP `FULL_SHARD`。同一个 [`train_distill_wan2.1.py`](./train_distill_wan2.1.py) 也可在 DeepSpeed Zero-2 下运行，或不使用任何分片后端；在 1.3B 规模下这些方案比 Wan2.2 的 14B 专家更可行。请保持脚本参数与 §3.3 完全一致——只改动 `accelerate launch` 头部。
|
||||
|
||||
#### 3.6.1 使用 DeepSpeed-Zero-2 进行训练
|
||||
|
||||
DeepSpeed Zero-2 只分片优化器态与梯度、**不**分片参数。对三个同驻的 1.3B transformer 而言，它在默认分辨率下即可从容放下：
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/sparkdiffusion/train_distill_wan2.1.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--random_hw_adapt \
|
||||
--enable_bucket \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--uniform_sampling \
|
||||
--denoising_step_indices_list 1000 750 500 250 \
|
||||
--flow_euler_rollout \
|
||||
--randomize_step_indices \
|
||||
--gen_update_interval 5 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 5.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 5.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--sparse_student \
|
||||
--student_topk=0.10 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
#### 3.6.2 不使用 DeepSpeed 与 FSDP 训练
|
||||
|
||||
完全去掉分片后端——三个同驻的 1.3B transformer 可放进单张高显存 GPU，不过多卡仍推荐 FSDP。把 `accelerate launch` 头部替换为朴素形式，脚本参数与 §3.6.1 保持一致：
|
||||
|
||||
```bash
|
||||
accelerate launch --mixed_precision="bf16" \
|
||||
scripts/sparkdiffusion/train_distill_wan2.1.py \
|
||||
<与 §3.6.1 相同的 --flag 参数块>
|
||||
```
|
||||
|
||||
### 3.7 多机分布式训练
|
||||
|
||||
启动脚本为单机。多机时在 `accelerate launch` 行追加下列参数（所有机器须通过 NFS/共享存储共享同一份数据集 / 基座权重 / `output_dir`）：
|
||||
|
||||
```bash
|
||||
--main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT \
|
||||
--num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK
|
||||
```
|
||||
|
||||
无 RDMA 时，另设 `NCCL_IB_DISABLE=1` 与 `NCCL_P2P_DISABLE=1`。
|
||||
|
||||
---
|
||||
|
||||
## 四、推理 — 稀疏少步学生
|
||||
|
||||
Stage-2 产物是一个少步、CFG-free 的稀疏学生。作为单专家模型，推理只载入**一个** checkpoint（没有高噪专家）。构建 wan2.1 稀疏类 `SparkWanTransformer3DModel`，设置训练时的 top-k，再载入蒸馏权重——顺序很关键，因为 `set_student_topk` 会（重）建接收训练好 RoLa 参数的 `attn_op` 内核：
|
||||
|
||||
```python
|
||||
from videox_fun.models import SparkWanTransformer3DModel
|
||||
transformer = SparkWanTransformer3DModel.from_pretrained(
|
||||
os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
|
||||
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
||||
low_cpu_mem_usage=True, torch_dtype=weight_dtype,
|
||||
)
|
||||
transformer.set_student_topk(student_topk) # 必须在 load_state_dict 之前
|
||||
# transformer.load_state_dict(load_file(transformer_path), strict=False)
|
||||
```
|
||||
|
||||
Wan2.1 T2V 入口 [examples/wan2.1/predict_t2v.py](../../examples/wan2.1/predict_t2v.py) 已在 `WanPipeline` 下运行少步 CFG-free 采样（`sampler_name="Flow"`、`num_inference_steps=4`、`guidance_scale=1.0`）；把 dense 的 `WanTransformer3DModel` 换成上面的稀疏构建即可。需设置：
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|-----------|-------------|---------------|
|
||||
| `GPU_memory_mode` | 显存模式 | `model_full_load` |
|
||||
| `config_path` | Wan2.1 config | `config/wan2.1/wan_civitai.yaml` |
|
||||
| `model_name` | Wan2.1-T2V-1.3B 基座（transformer / VAE / UMT5） | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
|
||||
| `student_topk` | **必须等于**训练时的 `--student_topk` | `0.10` |
|
||||
| `transformer_path` | 稀疏学生 checkpoint（单专家） | `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
|
||||
| `sampler_name` | `Flow` 与少步 CFG-free 学生匹配 | `Flow` |
|
||||
| `num_inference_steps` | 学生的少步步数 | `4` |
|
||||
| `guidance_scale` | 蒸馏学生为 CFG-free | `1.0` |
|
||||
| `sample_size` / `video_length` / `fps` | `[高, 宽]` / 帧数 / 帧率 | `[480, 832]` / `81` / `16` |
|
||||
|
||||
```bash
|
||||
python examples/wan2.1/predict_t2v.py
|
||||
```
|
||||
|
||||
> 💡 载入顺序为 `from_pretrained → set_student_topk → load_state_dict`；顺序颠倒会把训练好的 RoLa 权重重新清零。当 `attn_op.*` 键载入且 unexpected keys 为 **0** 时，即确认为稀疏 checkpoint。
|
||||
|
||||
> 📌 仓库自带的 [examples/sparkdiffusion/predict_t2v.py](../../examples/sparkdiffusion/predict_t2v.py) 面向 Wan2.2 双专家 MoE 学生（`SparkWan2_2Transformer3DModel` + `Wan2_2Pipeline`）；单专家的 Wan2.1 学生使用 `SparkWanTransformer3DModel` + `WanPipeline`，只载入一个 checkpoint。
|
||||
|
||||
> 📌 论文的 **Stage 3**（FP8/INT8 融合量化推理核）在本仓库中有意未接入。
|
||||
|
||||
---
|
||||
|
||||
## 五、更多资源
|
||||
|
||||
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
|
||||
- **论文**：SparkDiffusion，《Mitigating the High-Sparsity Trap》（arXiv:2609.23153）
|
||||
- **稀疏学生模型**：`videox_fun/models/spark_transformer3d.py`（`SparkWanTransformer3DModel`）
|
||||
- **RoLa 稀疏注意力内核**：`videox_fun/models/attention_utils.py`
|
||||
- **Wan2.2 MoE 对应文档**：[README_TRAIN_DISTILL_wan2.2_zh-CN.md](./README_TRAIN_DISTILL_wan2.2_zh-CN.md)
|
||||
@@ -0,0 +1,495 @@
|
||||
# SparkDiffusion Sparse Distillation Training Guide
|
||||
|
||||
This document provides the complete workflow for **SparkDiffusion** on Wan2.2-T2V-A14B: a two-stage pipeline that first warms up a **sparse (RoLa) student** and then **distills** it into a few-step, CFG-free generator. It covers environment setup, data preparation, the four launchers (Stage-1 warm-up + Stage-2 CrossDistill, per MoE expert), the one-shot orchestrator, and the `predict_t2v.py` inference acceptance test.
|
||||
|
||||
> **What is SparkDiffusion?**
|
||||
>
|
||||
> SparkDiffusion (arXiv:2609.23153, *Mitigating the High-Sparsity Trap*) combines a **compensating sparse attention** (RoLa) with a **CrossDistill** few-step recipe. Wan2.2-T2V-A14B is a two-expert MoE (a high-noise expert and a low-noise expert, split at the config `boundary: 0.875`); each expert is warmed up and distilled **on its own noise range**, never jointly.
|
||||
>
|
||||
> - **Stage 1 — sparse warm-up** ([`train_finetune_sparse_wan2.2.py`](./train_finetune_sparse_wan2.2.py)): insert the RoLa sparse self-attention into the dense backbone and briefly fine-tune it with the plain step-local flow-matching objective, so distillation starts from a usable coarse prior instead of a zero-initialized sparse branch. This stage is intentionally **short**.
|
||||
> - **Stage 2 — CrossDistill** ([`train_distill_wan2.2.py`](./train_distill_wan2.2.py)): distill the sparse student into a few-step generator against the **dense** teacher. It is Distribution Matching Distillation (DMD) with three co-resident transformers — a trainable sparse `generator`, a frozen dense `real_score` teacher, and a trainable dense `fake_score` critic. The **high-noise** expert adds a PCM-style diversity term; the **low-noise** expert is div-free.
|
||||
>
|
||||
> | Stage | High-noise expert | Low-noise expert | Output |
|
||||
> |-------|-------------------|------------------|--------|
|
||||
> | **1 warm-up** | [`train_finetune_sparse_wan2.2_high.sh`](./train_finetune_sparse_wan2.2_high.sh) | [`train_finetune_sparse_wan2.2_low.sh`](./train_finetune_sparse_wan2.2_low.sh) | `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola` |
|
||||
> | **2 distill** | [`train_distill_wan2.2_high.sh`](./train_distill_wan2.2_high.sh) | [`train_distill_wan2.2_low.sh`](./train_distill_wan2.2_low.sh) | `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola_distill` |
|
||||
>
|
||||
> The two experts are distilled separately and both are loaded at inference to form the full few-step MoE student.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
- [1. Environment Setup](#1-environment-setup)
|
||||
- [2. Data Preparation](#2-data-preparation)
|
||||
- [3. Training](#3-training)
|
||||
- [3.1 Download Pretrained Models](#31-download-pretrained-models)
|
||||
- [3.2 Stage 1 — Sparse Warm-up](#32-stage-1--sparse-warm-up)
|
||||
- [3.3 Stage 2 — CrossDistill](#33-stage-2--crossdistill)
|
||||
- [3.4 Common Training Parameters](#34-common-training-parameters)
|
||||
- [3.5 FSDP Backend and Memory Notes](#35-fsdp-backend-and-memory-notes)
|
||||
- [3.6 Other Backends](#36-other-backends)
|
||||
- [3.7 Multi-Node Distributed Training](#37-multi-node-distributed-training)
|
||||
- [4. Inference — Sparse Few-Step Student](#4-inference--sparse-few-step-student)
|
||||
- [5. Additional Resources](#5-additional-resources)
|
||||
|
||||
---
|
||||
|
||||
## 1. Environment Setup
|
||||
|
||||
**Method 1: Using requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**Method 2: Manual Installation**
|
||||
|
||||
```bash
|
||||
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
|
||||
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
|
||||
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
|
||||
pip install yunchang xfuser modelscope openpyxl
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
pip install deepspeed==0.17.0 numpy==1.26.4
|
||||
```
|
||||
|
||||
> 💡 The RoLa sparse attention relies on the in-repo Triton block-sparse kernel (`videox_fun/models/attention_utils.py`), so a working Triton install is required for the sparse student.
|
||||
|
||||
**Method 3: Using Docker**
|
||||
|
||||
```bash
|
||||
# pull image
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
|
||||
# enter image
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Data Preparation
|
||||
|
||||
SparkDiffusion consumes ordinary video-caption data (the same format as the standard Wan2.2 training).
|
||||
|
||||
### 2.1 Quick Test Dataset
|
||||
|
||||
```bash
|
||||
# Download official example dataset
|
||||
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
|
||||
```
|
||||
|
||||
### 2.2 metadata.json Format
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"file_path": "train/video001.mp4",
|
||||
"text": "A beautiful sunset over the ocean, golden hour lighting",
|
||||
"type": "video",
|
||||
"width": 1024,
|
||||
"height": 1024
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
- `file_path`: video path (relative to `DATASET_NAME`, or absolute).
|
||||
- `text`: caption / prompt.
|
||||
- `type`: fixed as `"video"`.
|
||||
- `width` / `height`: recommended for bucket training; add them with `scripts/process_json_add_width_and_height.py` if absent.
|
||||
|
||||
### 2.3 Relative vs Absolute Paths
|
||||
|
||||
```bash
|
||||
# relative paths
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
|
||||
# absolute paths (data on NAS/OSS shared across machines)
|
||||
export DATASET_NAME=""
|
||||
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. Training
|
||||
|
||||
The full order is **Stage-1 high → Stage-2 high → Stage-1 low → Stage-2 low**. Run the four launchers by hand, in that order.
|
||||
|
||||
### 3.1 Download Pretrained Models
|
||||
|
||||
```bash
|
||||
# Create model directory
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# Download Wan2.2 T2V official weights
|
||||
modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-T2V-A14B
|
||||
```
|
||||
|
||||
The base checkpoint seeds all three Stage-2 transformers (the sparse `generator` is additionally initialized from the Stage-1 warm-up product). Its layout is `high_noise_model/`, `low_noise_model/`, the VAE and the UMT5 text encoder; the MoE boundary is read from `config/wan2.2/wan_civitai_t2v.yaml` (`boundary: 0.875`).
|
||||
|
||||
### 3.2 Stage 1 — Sparse Warm-up
|
||||
|
||||
Stage 1 inserts the RoLa sparse self-attention and briefly warms it up with the step-local flow-matching objective. It is a **short** stage — a brief warm-up is enough to establish a usable prior for distillation, so the launchers cap it at `--max_train_steps=2000`. Do **not** judge this stage by its own sample quality; the sharpness is established in Stage 2.
|
||||
|
||||
Run the high-noise expert with [scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh](./train_finetune_sparse_wan2.2_high.sh):
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--max_train_steps=2000 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=1e-4 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--enable_bucket \
|
||||
--uniform_sampling \
|
||||
--trainable_modules "attn_op" \
|
||||
--student_topk=0.05 \
|
||||
--train_mode="normal" \
|
||||
--boundary_type="high" \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
Or simply:
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh
|
||||
```
|
||||
|
||||
The low-noise expert uses the same recipe with `--boundary_type="low"` and its own `--output_dir`:
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.2_low.sh
|
||||
```
|
||||
|
||||
Stage-1 keys:
|
||||
|
||||
| Flag | Meaning |
|
||||
|------|---------|
|
||||
| `--boundary_type` | `high` / `low` — selects the MoE expert to warm up (from the config boundary). |
|
||||
| `--student_topk` | Ratio of key blocks kept per query in the sparse attention (sparsity = `1 - topk`). `0.05` ≈ 95% sparse. **Must match** the value used in Stage 2. |
|
||||
| `--trainable_modules` | `"attn_op"` = Stage A (train only the sparse branch); `"."` = Stage B (full sparse finetune, load Stage A via `--transformer_path`). |
|
||||
| `--max_train_steps` | Warm-up budget; keep it short. |
|
||||
|
||||
On completion the warm-up product is `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`.
|
||||
|
||||
### 3.3 Stage 2 — CrossDistill
|
||||
|
||||
Stage 2 distills the sparse student into a few-step, CFG-free generator against the dense teacher. `--sparse_student` turns the `generator` (and the `init_rollout` model) into `SparkWan2_2Transformer3DModel` while the teacher and critic stay dense. `--generator_transformer_path` loads the Stage-1 warm-up checkpoint. The host is a plain **discrete-grid DMD**: the few-step grid is derived in-script by splitting `--joint_backward_timesteps` at the config `boundary` (`0.875`) and keeping only this expert's segment (`--boundary_type`), then rolled out with rf-native Euler steps (`--flow_euler_rollout`), the generator and critic alternating via `--gen_update_interval`. On top of this host the **high** expert adds an optional CrossDistill PCM diversity term (`--loss_scale_div > 0`, see §3.3.1); the **low** expert is div-free. An opt-in `--tdm` switch (§3.3.2) replaces the plain-DMD student loss with the official SparkDiffusion **TDM** loss; it is not enabled in the published launchers, so the default run stays plain DMD.
|
||||
|
||||
Run the high-noise expert with [scripts/sparkdiffusion/train_distill_wan2.2_high.sh](./train_distill_wan2.2_high.sh):
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# Stage-1 warm-up product.
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
|
||||
--fsdp_offload_params true \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_n_frames=81 \
|
||||
--fix_sample_size 480 832 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--enable_bucket \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--uniform_sampling \
|
||||
--flow_euler_rollout \
|
||||
--gen_update_interval 5 \
|
||||
--joint_backward_timesteps 0.933781 0.875 0.608979 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 10.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 12.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--boundary_type="high" \
|
||||
--real_guidance_scale=5.0 \
|
||||
--sparse_student \
|
||||
--student_topk=0.05 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
Or simply:
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_distill_wan2.2_high.sh
|
||||
```
|
||||
|
||||
The low-noise expert ([scripts/sparkdiffusion/train_distill_wan2.2_low.sh](./train_distill_wan2.2_low.sh)) differs in three ways:
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_distill_wan2.2_low.sh
|
||||
```
|
||||
|
||||
| Difference | High expert | Low expert |
|
||||
|------------|-------------|------------|
|
||||
| `--boundary_type` | `high` | `low` |
|
||||
| Diversity term | `--div_mode="pcm_random" --loss_scale_div=10.0` | div-free (no div flags) |
|
||||
| `--init_rollout_path` | not used | warm-starts the rollout's boundary latent from the trained **high** student (`INIT_ROLLOUT_CKPT`) |
|
||||
|
||||
> ⚠️ The low-noise segment of the few-step trajectory starts at the MoE boundary, not at pure noise. Point `INIT_ROLLOUT_CKPT` at the latest high-distill checkpoint (`output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors`) so the frozen high student rolls the high-noise segment; leaving it empty makes the untrained low student roll that segment itself and the samples degrade.
|
||||
|
||||
> ⚠️ **The distillation phases have no step cap.** They run until you stop them manually (kill the `train_distill_wan2.2.py` processes); checkpoints are written every `--checkpointing_steps=200`.
|
||||
|
||||
#### 3.3.1 CrossDistill Diversity (PCM, high expert only)
|
||||
|
||||
On the **high**-noise expert, setting `--loss_scale_div > 0` layers the official SparkDiffusion **CrossDistill** PCM diversity term on top of the discrete-grid DMD host, to counter the diversity loss that pure distribution matching can cause. It is a self-contained additive loss that leaves the DMD rollout untouched (the **low** expert keeps `--loss_scale_div 0` and is div-free):
|
||||
|
||||
- **What it does** — a phase-0 boundary-alignment MSE: the student's one-step jump from a noisy anchor to the phase-0 boundary (`--div_rf_boundary`, the high segment's first node `0.933781`) must match a teacher K-step Euler rollout (`--dcm_skipping_interval_steps` sub-steps) followed by the student's own target jump. It runs on its **own continuous rf grid** (`--dcm_total_steps` subdivisions shifted by `--teacher_timestep_shift 12.0`, the Wan2.2 scheduler shift), decoupled from the discrete DMD grid.
|
||||
- **Anchor** — the paired **real video latent** (faithful to the official `_compute_pcm_div_loss`, which distills on real `x0`), encoded on demand only when diversity is on.
|
||||
- **Merge** — `generator_loss = dmd_loss + loss_scale_div * div_loss` in a **single** `accelerator.backward` (FSDP allows one backward per wrapped module per iteration); the teacher rollout and the student target run under `no_grad` before the graded student forward.
|
||||
- **Scale & tuning** — `--loss_scale_div` is in **per-element-mean** space (aligned with the mean-reduced `dmd_loss`); the launchers start at `10.0` on the high expert. On the first run watch `train_div_loss` vs `train_dmd_loss` and adjust so diversity does not overwhelm DMD; `--div_teacher_guidance` is the diversity teacher's CFG scale, independent of `--real_guidance_scale`. Set `--loss_scale_div 0` to recover the exact plain-DMD baseline.
|
||||
|
||||
#### 3.3.2 TDM Loss (opt-in)
|
||||
|
||||
An opt-in `--tdm` switch replaces the plain-DMD student loss with **TDM** (Trajectory Distribution Matching, the official SparkDiffusion student loss); **it is not enabled in the published `train_distill_wan2.2_{high,low}.sh`**, so the default run stays plain DMD. Three deltas from DMD: (1) the last rollout segment does not collapse to `x0` but takes one more rf-Euler step to `s_{i-1}` (the next grid knot), so the anchor is `G_x_{t_{i-1}}`; (2) `D_time` is clamped strictly above `s_{i-1}`; (3) the fake/teacher `x0` preds are mapped back to `s_{i-1}` via the rf PF-ODE analytic step and the surrogate is compared in the `x_{t_{i-1}}` domain. When `s_{i-1}==0` (grid exhausted) it reduces **exactly** to DMD. Requires `--flow_euler_rollout`; mutually exclusive with `--dfd` / `--decoupled_dmd`; compatible with the diversity term on the high expert.
|
||||
|
||||
### 3.4 Common Training Parameters
|
||||
|
||||
| Parameter | Description | Example Value |
|
||||
|-----------|-------------|---------------|
|
||||
| `--config_path` | Wan2.2 config; carries the MoE `boundary` | `config/wan2.2/wan_civitai_t2v.yaml` |
|
||||
| `--pretrained_model_name_or_path` | Wan2.2-T2V-A14B base (both experts + VAE + UMT5) | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` |
|
||||
| `--train_data_dir` / `--train_data_meta` | Dataset root / metadata JSON | see §2 |
|
||||
| `--boundary_type` | `high` / `low` MoE expert selector | `high` |
|
||||
| `--train_mode` | SparkDiffusion supports `normal` (T2V) only | `normal` |
|
||||
| `--student_topk` | Sparse attention keep-ratio (sparsity = 1 − topk); Stage 1 and Stage 2 must match | `0.05` |
|
||||
| `--sparse_student` | (Stage 2) build generator + init_rollout as the sparse student | - |
|
||||
| `--generator_transformer_path` | (Stage 2) Stage-1 warm-up checkpoint loaded into the generator | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` |
|
||||
| `--uniform_sampling` | Sample the DMD noise level uniformly over the training schedule | - |
|
||||
| `--flow_euler_rollout` | Roll the generator in flow/velocity space with deterministic rf-native Euler steps (no x0 re-noising) | - |
|
||||
| `--tdm` | (opt-in) Use TDM — the official SparkDiffusion student loss — instead of plain DMD; reduces exactly to DMD when `s_{i-1}==0`. Requires `--flow_euler_rollout`, mutually exclusive with `--dfd` / `--decoupled_dmd` | - |
|
||||
| `--gen_update_interval` | Update the generator once every N steps (the critic updates every step) | `5` |
|
||||
| `--joint_backward_timesteps` | Joint few-step rf grid (descending), split at the config `boundary` into per-expert segments by `--boundary_type`; passed explicitly in the launchers (equals the built-in default), and validation reuses it | `0.933781 0.875 0.608979` |
|
||||
| `--div_mode` | CrossDistill diversity mode: `pcm_random` (random phase-0 offset) or `pcm_fixed` (high expert only) | `pcm_random` |
|
||||
| `--loss_scale_div` | Weight of the CrossDistill PCM diversity term (per-element-mean space); `0` disables it | `10.0` (high) / `0` (low) |
|
||||
| `--dcm_total_steps` | Subdivision of the diversity term's continuous rf grid (`dw = 1 / N`) | `48` |
|
||||
| `--dcm_skipping_interval_steps` | Teacher Euler sub-steps `K` within phase 0 | `1` |
|
||||
| `--teacher_timestep_shift` | rf timestep shift for the diversity grid (matches the Wan2.2 scheduler shift) | `12.0` |
|
||||
| `--div_rf_boundary` | Phase-0 right boundary `rf_b` (the high segment's first node) | `0.933781` |
|
||||
| `--div_teacher_guidance` | CFG scale of the diversity teacher rollout (independent of `--real_guidance_scale`) | `5.0` |
|
||||
| `--real_guidance_scale` | CFG scale of the frozen dense teacher | `5.0` |
|
||||
| `--init_rollout_path` | (low expert) trained high student used to roll the rollout's boundary start latent | high-distill ckpt |
|
||||
| `--learning_rate` / `--learning_rate_critic` | Generator / critic learning rate | `2e-06` / `2e-06` |
|
||||
| `--checkpointing_steps` | Save checkpoint every N steps | `200` |
|
||||
| `--max_grad_norm` | Gradient clipping threshold | `0.05` |
|
||||
| `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` |
|
||||
|
||||
> 💡 The per-expert few-step grid is **auto-derived**: `--joint_backward_timesteps` (the joint rf grid, passed explicitly in the launchers) is split at the config `boundary` by `--boundary_type`, and `--init_rollout_timesteps` (the low expert's frozen high-segment relay) is filled from the knots above the boundary. Leave `--init_rollout_timesteps` at its default; validation reuses the same grid, so `num_inference_steps` = `len(--joint_backward_timesteps) + 1` (`4` here).
|
||||
|
||||
### 3.5 FSDP Backend and Memory Notes
|
||||
|
||||
Stage 2 keeps **three** 14B transformers co-resident (sparse generator + dense teacher + dense critic; the low expert also parks a frozen `init_rollout` high-noise relay, sharded and `no_grad`), all optimizer-state dominated, so a sharding backend is mandatory — FSDP `FULL_SHARD` here. The distill launchers combine these memory levers:
|
||||
|
||||
- `--fsdp_sharding_strategy "FULL_SHARD"` — shard params/grads/optimizer across GPUs.
|
||||
- `--fsdp_backward_prefetch "BACKWARD_POST"` — lower median residency during backward.
|
||||
- `--gradient_checkpointing` — recompute activations in backward.
|
||||
- **Resolution, not frame count, is the OOM lever.** Each step rolls the few-step grid and backprops a grad-carrying 14B forward, so memory scales with `seq_len`; the launchers use `--fix_sample_size 480 832` with `--video_sample_n_frames=81`. If it OOMs, lower `--fix_sample_size` (e.g. to `320 512`) rather than dropping frames.
|
||||
- Frame count must satisfy `(N − 1) % 4 == 0`.
|
||||
|
||||
> 💡 The distill launchers already pass `--fsdp_offload_params true` (mandatory at 14B: it offloads params/grads/optimizer to CPU, at some speed cost); drop it only if you have ample VRAM. `--low_vram` additionally parks the frozen side (`real_score` / `init_rollout` / text encoder) on CPU; it complements, but does not replace, sharding.
|
||||
|
||||
### 3.6 Other Backends
|
||||
|
||||
The Stage-2 launchers default to FSDP `FULL_SHARD`. The same [`train_distill_wan2.2.py`](./train_distill_wan2.2.py) also runs under DeepSpeed Zero-2 or without a sharding backend. All examples below use the **high**-noise expert; for the low expert flip `--boundary_type="low"`, drop the diversity flags (`--div_mode` … `--div_teacher_guidance`), add `--init_rollout_path`, and change `--output_dir` / `SPARSE_STUDENT_CKPT` as in §3.3.
|
||||
|
||||
#### 3.6.1 Training with DeepSpeed Zero-2
|
||||
|
||||
DeepSpeed Zero-2 shards optimizer states and gradients but **not** parameters. Since Stage 2 keeps three 14B transformers co-resident, Zero-2 only fits at low resolution and may OOM where FSDP `FULL_SHARD` (which also shards parameters) does not — prefer FSDP for the 14B experts.
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# Stage-1 warm-up product.
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/sparkdiffusion/train_distill_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_n_frames=81 \
|
||||
--fix_sample_size 480 832 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--enable_bucket \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--uniform_sampling \
|
||||
--flow_euler_rollout \
|
||||
--gen_update_interval 5 \
|
||||
--joint_backward_timesteps 0.933781 0.875 0.608979 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 10.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 12.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--boundary_type="high" \
|
||||
--real_guidance_scale=5.0 \
|
||||
--sparse_student \
|
||||
--student_topk=0.05 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
#### 3.6.2 Training Without DeepSpeed and FSDP
|
||||
|
||||
**This approach is not recommended: with three co-resident 14B transformers and no sharding backend it easily OOMs.** It is provided only as a reference shell.
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# Stage-1 warm-up product.
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" scripts/sparkdiffusion/train_distill_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_n_frames=81 \
|
||||
--fix_sample_size 480 832 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--enable_bucket \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--uniform_sampling \
|
||||
--flow_euler_rollout \
|
||||
--gen_update_interval 5 \
|
||||
--joint_backward_timesteps 0.933781 0.875 0.608979 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 10.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 12.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--boundary_type="high" \
|
||||
--real_guidance_scale=5.0 \
|
||||
--sparse_student \
|
||||
--student_topk=0.05 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
### 3.7 Multi-Node Distributed Training
|
||||
|
||||
The launchers are single-node. For multiple machines, add the following to the `accelerate launch` line (all machines must share the same dataset / base weights / `output_dir` over NFS/shared storage):
|
||||
|
||||
```bash
|
||||
--main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT \
|
||||
--num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK
|
||||
```
|
||||
|
||||
Without RDMA, also set `NCCL_IB_DISABLE=1` and `NCCL_P2P_DISABLE=1`.
|
||||
|
||||
---
|
||||
|
||||
## 4. Inference — Sparse Few-Step Student
|
||||
|
||||
Acceptance is a standalone run of [examples/sparkdiffusion/predict_t2v.py](../../examples/sparkdiffusion/predict_t2v.py), which rebuilds each MoE expert as a sparse `SparkWan2_2Transformer3DModel` and loads the distilled checkpoints. Edit the top of the script:
|
||||
|
||||
| Parameter | Description | Example Value |
|
||||
|-----------|-------------|---------------|
|
||||
| `GPU_memory_mode` | Memory mode | `model_cpu_offload` |
|
||||
| `model_name` | Wan2.2-T2V-A14B base (VAE / UMT5 / schedulers) | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` |
|
||||
| `student_topk` | **Must equal** the training `--student_topk` | `0.05` |
|
||||
| `transformer_path` | **Low**-noise sparse student checkpoint | `output_dir_sparkdiffusion_wan2.2_14b_low_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
|
||||
| `transformer_high_path` | **High**-noise sparse student checkpoint | `output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
|
||||
| `sampler_name` | `Flow` matches the few-step CFG-free student | `Flow` |
|
||||
| `num_inference_steps` | The student's few-step count | `4` |
|
||||
| `guidance_scale` | Distilled student is CFG-free | `1.0` |
|
||||
| `sample_size` / `video_length` / `fps` | `[height, width]` / frames / fps | `[240, 336]` / `81` / `16` |
|
||||
|
||||
```bash
|
||||
python examples/sparkdiffusion/predict_t2v.py
|
||||
```
|
||||
|
||||
> 💡 Per-expert loading order is `from_pretrained → set_student_topk → load_state_dict`; reversing it re-zeroes the trained RoLa weights. An expert whose path is `None` stays dense — for the full SparkDiffusion student, set **both** paths (prefer each expert's Stage-2 distill product).
|
||||
|
||||
> 📌 **Stage 3** of the paper (the FP8/INT8 fused quantized inference kernel) is intentionally not integrated in this repository.
|
||||
|
||||
---
|
||||
|
||||
## 5. Additional Resources
|
||||
|
||||
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
|
||||
- **Paper**: SparkDiffusion, *Mitigating the High-Sparsity Trap* (arXiv:2609.23153)
|
||||
- **Sparse student model**: `videox_fun/models/spark_transformer3d.py` (`SparkWan2_2Transformer3DModel`)
|
||||
- **RoLa sparse attention kernel**: `videox_fun/models/attention_utils.py`
|
||||
@@ -0,0 +1,495 @@
|
||||
# SparkDiffusion 稀疏蒸馏训练指南
|
||||
|
||||
本文档给出 Wan2.2-T2V-A14B 上 **SparkDiffusion** 的完整流程：一条两阶段管线，先暖启一个**稀疏（RoLa）学生**，再把它**蒸馏**成少步、CFG-free 的生成器。内容涵盖环境配置、数据准备、四个启动脚本（Stage-1 暖启 + Stage-2 CrossDistill，按 MoE 专家分别执行）、以及 `predict_t2v.py` 推理验收。
|
||||
|
||||
> **什么是 SparkDiffusion？**
|
||||
>
|
||||
> SparkDiffusion（arXiv:2609.23153，《Mitigating the High-Sparsity Trap》）把**补偿式稀疏注意力**（RoLa）与 **CrossDistill** 少步配方结合起来。Wan2.2-T2V-A14B 是双专家 MoE（一个高噪专家、一个低噪专家，按 config 的 `boundary: 0.875` 切分）；每个专家都在**各自的噪声区间**上单独暖启与蒸馏，绝不联合训练。
|
||||
>
|
||||
> - **Stage 1 — 稀疏暖启**（[`train_finetune_sparse_wan2.2.py`](./train_finetune_sparse_wan2.2.py)）：把 RoLa 稀疏自注意力插入 dense 主干，用标准的 step-local flow-matching 目标短暂微调，使蒸馏从一个可用的粗先验出发，而不是从零初始化的稀疏分支出发。该阶段刻意保持**短**。
|
||||
> - **Stage 2 — CrossDistill**（[`train_distill_wan2.2.py`](./train_distill_wan2.2.py)）：对着 **dense** teacher 把稀疏学生蒸馏成少步生成器。它是分布匹配蒸馏（DMD），含三个同驻网络——可训练的稀疏 `generator`、冻结的 dense `real_score` teacher、可训练的 dense `fake_score` critic。**高噪**专家额外加一项 PCM 式多样性损失；**低噪**专家则不含多样性项。
|
||||
>
|
||||
> | 阶段 | 高噪专家 | 低噪专家 | 产物 |
|
||||
> |-------|-------------------|------------------|--------|
|
||||
> | **1 暖启** | [`train_finetune_sparse_wan2.2_high.sh`](./train_finetune_sparse_wan2.2_high.sh) | [`train_finetune_sparse_wan2.2_low.sh`](./train_finetune_sparse_wan2.2_low.sh) | `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola` |
|
||||
> | **2 蒸馏** | [`train_distill_wan2.2_high.sh`](./train_distill_wan2.2_high.sh) | [`train_distill_wan2.2_low.sh`](./train_distill_wan2.2_low.sh) | `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola_distill` |
|
||||
>
|
||||
> 两个专家分别蒸馏，推理时同时载入，共同构成完整的少步 MoE 学生。
|
||||
|
||||
---
|
||||
|
||||
## 目录
|
||||
- [一、环境配置](#一环境配置)
|
||||
- [二、数据准备](#二数据准备)
|
||||
- [三、训练](#三训练)
|
||||
- [3.1 下载预训练模型](#31-下载预训练模型)
|
||||
- [3.2 Stage 1 — 稀疏暖启](#32-stage-1--稀疏暖启)
|
||||
- [3.3 Stage 2 — CrossDistill](#33-stage-2--crossdistill)
|
||||
- [3.4 通用训练参数](#34-通用训练参数)
|
||||
- [3.5 FSDP 后端与显存说明](#35-fsdp-后端与显存说明)
|
||||
- [3.6 其他后端](#36-其他后端)
|
||||
- [3.7 多机分布式训练](#37-多机分布式训练)
|
||||
- [四、推理 — 稀疏少步学生](#四推理--稀疏少步学生)
|
||||
- [五、更多资源](#五更多资源)
|
||||
|
||||
---
|
||||
|
||||
## 一、环境配置
|
||||
|
||||
**方式一：使用 requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**方式二：手动安装**
|
||||
|
||||
```bash
|
||||
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
|
||||
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
|
||||
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
|
||||
pip install yunchang xfuser modelscope openpyxl
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
pip install deepspeed==0.17.0 numpy==1.26.4
|
||||
```
|
||||
|
||||
> 💡 RoLa 稀疏注意力依赖仓库内的 Triton block-sparse 内核（`videox_fun/models/attention_utils.py`），因此稀疏学生需要一个可用的 Triton 安装。
|
||||
|
||||
**方式三：使用 Docker**
|
||||
|
||||
```bash
|
||||
# 拉取镜像
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
|
||||
# 进入镜像
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 二、数据准备
|
||||
|
||||
SparkDiffusion 使用普通的「视频-字幕」数据（与标准 Wan2.2 训练格式一致）。
|
||||
|
||||
### 2.1 快速测试数据集
|
||||
|
||||
```bash
|
||||
# 下载官方示例数据集
|
||||
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
|
||||
```
|
||||
|
||||
### 2.2 metadata.json 格式
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"file_path": "train/video001.mp4",
|
||||
"text": "A beautiful sunset over the ocean, golden hour lighting",
|
||||
"type": "video",
|
||||
"width": 1024,
|
||||
"height": 1024
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
- `file_path`：视频路径（相对于 `DATASET_NAME`，或绝对路径）。
|
||||
- `text`：字幕 / 提示词。
|
||||
- `type`：固定为 `"video"`。
|
||||
- `width` / `height`：建议提供以启用 bucket 训练；若缺失可用 `scripts/process_json_add_width_and_height.py` 补齐。
|
||||
|
||||
### 2.3 相对路径与绝对路径
|
||||
|
||||
```bash
|
||||
# 相对路径
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
|
||||
# 绝对路径（数据在 NAS/OSS 上、跨机共享）
|
||||
export DATASET_NAME=""
|
||||
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 三、训练
|
||||
|
||||
完整顺序为 **Stage-1 high → Stage-2 high → Stage-1 low → Stage-2 low**。请手动按此顺序运行四个启动脚本。
|
||||
|
||||
### 3.1 下载预训练模型
|
||||
|
||||
```bash
|
||||
# 创建模型目录
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# 下载 Wan2.2 T2V 官方权重
|
||||
modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-T2V-A14B
|
||||
```
|
||||
|
||||
基座权重会播种 Stage-2 的全部三个网络（稀疏 `generator` 还会额外用 Stage-1 暖启产物初始化）。其目录结构为 `high_noise_model/`、`low_noise_model/`、VAE 以及 UMT5 文本编码器；MoE 分界从 `config/wan2.2/wan_civitai_t2v.yaml`（`boundary: 0.875`）读取。
|
||||
|
||||
### 3.2 Stage 1 — 稀疏暖启
|
||||
|
||||
Stage 1 插入 RoLa 稀疏自注意力，并用 step-local flow-matching 目标短暂暖启。这是一个**短**阶段——短暂暖启即足以为蒸馏建立可用先验，因此启动脚本把它限制在 `--max_train_steps=2000`。**不要**用该阶段自身的采样质量来评判它；清晰度是在 Stage 2 建立的。
|
||||
|
||||
用 [scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh](./train_finetune_sparse_wan2.2_high.sh) 运行高噪专家：
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--max_train_steps=2000 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=1e-4 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--enable_bucket \
|
||||
--uniform_sampling \
|
||||
--trainable_modules "attn_op" \
|
||||
--student_topk=0.05 \
|
||||
--train_mode="normal" \
|
||||
--boundary_type="high" \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
或直接：
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh
|
||||
```
|
||||
|
||||
低噪专家用同样的配方，只把 `--boundary_type` 改为 `"low"` 并指定各自的 `--output_dir`：
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_finetune_sparse_wan2.2_low.sh
|
||||
```
|
||||
|
||||
Stage-1 关键参数：
|
||||
|
||||
| 旗标 | 含义 |
|
||||
|------|---------|
|
||||
| `--boundary_type` | `high` / `low`——选择要暖启的 MoE 专家（由 config 分界决定）。 |
|
||||
| `--student_topk` | 稀疏注意力中每个 query 保留的 key block 比例（稀疏度 = `1 - topk`）。`0.05` ≈ 95% 稀疏。**必须与** Stage 2 一致。 |
|
||||
| `--trainable_modules` | `"attn_op"` = Stage A（只训练稀疏分支）；`"."` = Stage B（稀疏网全量微调，用 `--transformer_path` 载入 Stage A）。 |
|
||||
| `--max_train_steps` | 暖启预算；保持简短。 |
|
||||
|
||||
完成后暖启产物为 `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`。
|
||||
|
||||
### 3.3 Stage 2 — CrossDistill
|
||||
|
||||
Stage 2 对着 dense teacher 把稀疏学生蒸馏成少步、CFG-free 的生成器。`--sparse_student` 把 `generator`（以及 `init_rollout` 模型）变为 `SparkWan2_2Transformer3DModel`，而 teacher 与 critic 保持 dense。`--generator_transformer_path` 载入 Stage-1 暖启 checkpoint。宿主是一个朴素的**离散网格 DMD**：少步网格在脚本内由 `--joint_backward_timesteps` 按 config 的 `boundary`（`0.875`）切分、只保留本专家的段（`--boundary_type`）而派生，随后用 rf-native 的 Euler 步进滚动（`--flow_euler_rollout`），generator 与 critic 通过 `--gen_update_interval` 交替更新。在此宿主之上，**高噪**专家额外叠加一个可选的 CrossDistill PCM 多样性项（`--loss_scale_div > 0`，见 §3.3.1）；**低噪**专家则不含多样性项。一个可选的 `--tdm` 开关（§3.3.2）把 plain-DMD 学生损失替换为官方 SparkDiffusion **TDM** 损失；它未在发布态启动脚本中启用，故默认运行仍是 plain DMD。
|
||||
|
||||
用 [scripts/sparkdiffusion/train_distill_wan2.2_high.sh](./train_distill_wan2.2_high.sh) 运行高噪专家：
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# Stage-1 暖启产物。
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
|
||||
--fsdp_offload_params true \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_n_frames=81 \
|
||||
--fix_sample_size 480 832 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--enable_bucket \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--uniform_sampling \
|
||||
--flow_euler_rollout \
|
||||
--gen_update_interval 5 \
|
||||
--joint_backward_timesteps 0.933781 0.875 0.608979 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 10.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 12.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--boundary_type="high" \
|
||||
--real_guidance_scale=5.0 \
|
||||
--sparse_student \
|
||||
--student_topk=0.05 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
或直接：
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_distill_wan2.2_high.sh
|
||||
```
|
||||
|
||||
低噪专家（[scripts/sparkdiffusion/train_distill_wan2.2_low.sh](./train_distill_wan2.2_low.sh)）有三处不同：
|
||||
|
||||
```bash
|
||||
bash scripts/sparkdiffusion/train_distill_wan2.2_low.sh
|
||||
```
|
||||
|
||||
| 差异项 | 高噪专家 | 低噪专家 |
|
||||
|------------|-------------|------------|
|
||||
| `--boundary_type` | `high` | `low` |
|
||||
| 多样性项 | `--div_mode="pcm_random" --loss_scale_div=10.0` | 无多样性项（不加 div 旗标） |
|
||||
| `--init_rollout_path` | 不使用 | 用训练好的**高噪**学生暖启 rollout 的分界起点 latent（`INIT_ROLLOUT_CKPT`） |
|
||||
|
||||
> ⚠️ 少步轨迹的低噪段起点在 MoE 分界处、而非纯噪声。请把 `INIT_ROLLOUT_CKPT` 指向最新的高噪蒸馏 checkpoint（`output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors`），让冻结的高噪学生去滚动高噪段；若留空，则由未训练的低噪学生自己滚动该段，样本会退化。
|
||||
|
||||
> ⚠️ **蒸馏阶段没有步数上限。** 它们会一直运行直到你手动停止（kill 掉 `train_distill_wan2.2.py` 进程）；checkpoint 每 `--checkpointing_steps=200` 步写一次。
|
||||
|
||||
#### 3.3.1 CrossDistill 多样性项（PCM，仅高噪专家）
|
||||
|
||||
在**高噪**专家上，设置 `--loss_scale_div > 0` 会在离散网格 DMD 宿主之上叠加官方 SparkDiffusion 的 **CrossDistill** PCM 多样性项，用以抵消纯分布匹配可能带来的多样性损失。它是一个自包含的加性损失，不改动 DMD rollout（**低噪**专家保持 `--loss_scale_div 0`，不含多样性项）：
|
||||
|
||||
- **它做什么**——一个 phase-0 分界对齐 MSE：学生从带噪锚点一步跳到 phase-0 分界（`--div_rf_boundary`，即高噪段的首个节点 `0.933781`），必须与 teacher 的 K 步 Euler rollout（`--dcm_skipping_interval_steps` 个子步）再接学生自己的目标跳变相匹配。它运行在**自己独立的连续 rf 网格**上（`--dcm_total_steps` 等分、并按 `--teacher_timestep_shift 12.0`——Wan2.2 调度器的 shift——偏移），与离散的 DMD 网格解耦。
|
||||
- **锚点**——配对的**真实视频 latent**（忠于官方 `_compute_pcm_div_loss`，其在真实 `x0` 上蒸馏），仅在开启多样性项时按需编码。
|
||||
- **合并**——`generator_loss = dmd_loss + loss_scale_div * div_loss`，在**单次** `accelerator.backward` 中完成（FSDP 每个被包裹的模块每轮迭代只允许一次 backward）；teacher rollout 与学生目标在带梯度的学生前向之前于 `no_grad` 下运行。
|
||||
- **尺度与调参**——`--loss_scale_div` 处于**逐元素均值**空间（与均值归约的 `dmd_loss` 对齐）；启动脚本在高噪专家上从 `10.0` 起步。首次运行时观察 `train_div_loss` 与 `train_dmd_loss`，调整到多样性不会压过 DMD；`--div_teacher_guidance` 是多样性 teacher 的 CFG 尺度，独立于 `--real_guidance_scale`。设 `--loss_scale_div 0` 即可恢复到完全一致的朴素 DMD 基线。
|
||||
|
||||
#### 3.3.2 TDM 损失（可选）
|
||||
|
||||
一个可选的 `--tdm` 开关把 plain-DMD 学生损失替换为 **TDM**（Trajectory Distribution Matching，官方 SparkDiffusion 学生损失）；**发布态的 `train_distill_wan2.2_{high,low}.sh` 中未启用它**，故默认运行仍是 plain DMD。相对 DMD 有三处差异：(1) rollout 末段不塌缩到 `x0`，而是再走一步 rf-Euler 到 `s_{i-1}`（下一个网格节点），故锚点为 `G_x_{t_{i-1}}`；(2) `D_time` 被严格 clamp 在 `s_{i-1}` 之上；(3) fake/teacher 的 `x0` 预测经 rf PF-ODE 解析步映回 `s_{i-1}`，surrogate 在 `x_{t_{i-1}}` 域比较。当 `s_{i-1}==0`（网格耗尽）时**精确**退化为 DMD。需 `--flow_euler_rollout`；与 `--dfd` / `--decoupled_dmd` 互斥；在高噪专家上与多样性项兼容。
|
||||
|
||||
### 3.4 通用训练参数
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|-----------|-------------|---------------|
|
||||
| `--config_path` | Wan2.2 config；携带 MoE `boundary` | `config/wan2.2/wan_civitai_t2v.yaml` |
|
||||
| `--pretrained_model_name_or_path` | Wan2.2-T2V-A14B 基座（两个专家 + VAE + UMT5） | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` |
|
||||
| `--train_data_dir` / `--train_data_meta` | 数据集根目录 / metadata JSON | 见 §2 |
|
||||
| `--boundary_type` | `high` / `low` MoE 专家选择器 | `high` |
|
||||
| `--train_mode` | SparkDiffusion 仅支持 `normal`（T2V） | `normal` |
|
||||
| `--student_topk` | 稀疏注意力保留比例（稀疏度 = 1 − topk）；Stage 1 与 Stage 2 必须一致 | `0.05` |
|
||||
| `--sparse_student` | （Stage 2）把 generator + init_rollout 建成稀疏学生 | - |
|
||||
| `--generator_transformer_path` | （Stage 2）载入 generator 的 Stage-1 暖启 checkpoint | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` |
|
||||
| `--uniform_sampling` | 在训练调度上均匀采样 DMD 的噪声水平 | - |
|
||||
| `--flow_euler_rollout` | 在 flow/velocity 空间用确定性的 rf-native Euler 步进滚动 generator（不做 x0 再加噪） | - |
|
||||
| `--tdm` | （可选）用 TDM（官方 SparkDiffusion 学生损失）替代 plain DMD；当 `s_{i-1}==0` 时精确退化为 DMD。需 `--flow_euler_rollout`，与 `--dfd` / `--decoupled_dmd` 互斥 | - |
|
||||
| `--gen_update_interval` | 每 N 步更新一次 generator（critic 每步都更新） | `5` |
|
||||
| `--joint_backward_timesteps` | 联合少步 rf 网格（降序），按 config 的 `boundary` 由 `--boundary_type` 切分成各专家的段；在启动脚本中显式传入（与内置默认值一致），validation 复用它 | `0.933781 0.875 0.608979` |
|
||||
| `--div_mode` | CrossDistill 多样性模式：`pcm_random`（随机 phase-0 偏移）或 `pcm_fixed`（仅高噪专家） | `pcm_random` |
|
||||
| `--loss_scale_div` | CrossDistill PCM 多样性项的权重（逐元素均值空间）；`0` 表示关闭 | `10.0`（高噪）/ `0`（低噪） |
|
||||
| `--dcm_total_steps` | 多样性项连续 rf 网格的等分数（`dw = 1 / N`） | `48` |
|
||||
| `--dcm_skipping_interval_steps` | phase 0 内 teacher 的 Euler 子步数 `K` | `1` |
|
||||
| `--teacher_timestep_shift` | 多样性网格的 rf 时间步 shift（与 Wan2.2 调度器的 shift 一致） | `12.0` |
|
||||
| `--div_rf_boundary` | phase-0 右边界 `rf_b`（高噪段的首个节点） | `0.933781` |
|
||||
| `--div_teacher_guidance` | 多样性 teacher rollout 的 CFG 尺度（独立于 `--real_guidance_scale`） | `5.0` |
|
||||
| `--real_guidance_scale` | 冻结 dense teacher 的 CFG 尺度 | `5.0` |
|
||||
| `--init_rollout_path` | （低噪专家）用于滚动 rollout 分界起点 latent 的、已训练好的高噪学生 | 高噪蒸馏 ckpt |
|
||||
| `--learning_rate` / `--learning_rate_critic` | 生成器 / critic 学习率 | `2e-06` / `2e-06` |
|
||||
| `--checkpointing_steps` | 每 N 步保存 checkpoint | `200` |
|
||||
| `--max_grad_norm` | 梯度裁剪阈值 | `0.05` |
|
||||
| `--resume_from_checkpoint` | 续训路径或 `"latest"` | `latest` |
|
||||
|
||||
> 💡 各专家的少步网格是**自动推导**的：`--joint_backward_timesteps`（联合 rf 网格，在启动脚本中显式传入）按 config 的 `boundary` 由 `--boundary_type` 切分，`--init_rollout_timesteps`（低噪专家冻结的高噪段中继）由分界以上的节点填充。`--init_rollout_timesteps` 保持默认值即可；validation 复用同一网格，故 `num_inference_steps` = `len(--joint_backward_timesteps) + 1`（此处为 `4`）。
|
||||
|
||||
### 3.5 FSDP Backend 与显存说明
|
||||
|
||||
Stage 2 需让**三个** 14B 网络同驻（稀疏 generator + dense teacher + dense critic；低噪专家还会停放一个冻结的 `init_rollout` 高噪中继，已分片且 `no_grad`），且都以优化器态为主，因此必须使用分片后端——此处为 FSDP `FULL_SHARD`。蒸馏启动脚本组合了下列显存杠杆：
|
||||
|
||||
- `--fsdp_sharding_strategy "FULL_SHARD"`——在多卡间分片 param/grad/optimizer。
|
||||
- `--fsdp_backward_prefetch "BACKWARD_POST"`——降低 backward 期间的中位驻留。
|
||||
- `--gradient_checkpointing`——在 backward 中重算激活。
|
||||
- **OOM 的杠杆是分辨率、而非帧数。** 每一步都会滚动少步网格并对一段带梯度的 14B 前向做反传，因此显存随 `seq_len` 增长；启动脚本使用 `--fix_sample_size 480 832` 配 `--video_sample_n_frames=81`。若 OOM，请降低 `--fix_sample_size`（例如降到 `320 512`），而不是减少帧数。
|
||||
- 帧数须满足 `(N − 1) % 4 == 0`。
|
||||
|
||||
> 💡 蒸馏启动脚本已经传入 `--fsdp_offload_params true`（14B 下为必需：它把 param/grad/optimizer offload 到 CPU，会牺牲一些速度）；仅在显存充裕时才可去掉。`--low_vram` 还会额外把冻结侧（`real_score` / `init_rollout` / 文本编码器）停到 CPU；它是对分片的补充，而非替代。
|
||||
|
||||
### 3.6 其他后端
|
||||
|
||||
Stage-2 启动脚本默认使用 FSDP `FULL_SHARD`。同一个 [`train_distill_wan2.2.py`](./train_distill_wan2.2.py) 也可在 DeepSpeed Zero-2 下运行，或不使用任何分片后端。下面所有示例都以**高噪**专家为例；低噪专家请参照 §3.3 把 `--boundary_type` 改为 `"low"`、去掉多样性旗标（`--div_mode` …… `--div_teacher_guidance`）、加上 `--init_rollout_path`，并修改 `--output_dir` / `SPARSE_STUDENT_CKPT`。
|
||||
|
||||
#### 3.6.1 使用 DeepSpeed-Zero-2 进行训练
|
||||
|
||||
DeepSpeed Zero-2 只分片优化器态与梯度、**不**分片参数。由于 Stage 2 需让三个 14B 网络同驻，Zero-2 仅在低分辨率下放得下，且可能在 FSDP `FULL_SHARD`（它还会分片参数）不会 OOM 的地方 OOM——14B 专家请优先选用 FSDP。
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# Stage-1 暖启产物。
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/sparkdiffusion/train_distill_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_n_frames=81 \
|
||||
--fix_sample_size 480 832 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--enable_bucket \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--uniform_sampling \
|
||||
--flow_euler_rollout \
|
||||
--gen_update_interval 5 \
|
||||
--joint_backward_timesteps 0.933781 0.875 0.608979 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 10.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 12.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--boundary_type="high" \
|
||||
--real_guidance_scale=5.0 \
|
||||
--sparse_student \
|
||||
--student_topk=0.05 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
#### 3.6.2 不使用 DeepSpeed 与 FSDP 训练
|
||||
|
||||
**该方案并不被推荐：三个 14B 网络同驻且没有分片后端，很容易造成显存不足**。这里仅提供训练 Shell 用于参考。
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# Stage-1 暖启产物。
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" scripts/sparkdiffusion/train_distill_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_n_frames=81 \
|
||||
--fix_sample_size 480 832 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--enable_bucket \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--uniform_sampling \
|
||||
--flow_euler_rollout \
|
||||
--gen_update_interval 5 \
|
||||
--joint_backward_timesteps 0.933781 0.875 0.608979 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 10.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 12.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--boundary_type="high" \
|
||||
--real_guidance_scale=5.0 \
|
||||
--sparse_student \
|
||||
--student_topk=0.05 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
### 3.7 多机分布式训练
|
||||
|
||||
启动脚本为单机。多机时在 `accelerate launch` 行追加下列参数（所有机器须通过 NFS/共享存储共享同一份数据集 / 基座权重 / `output_dir`）：
|
||||
|
||||
```bash
|
||||
--main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT \
|
||||
--num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK
|
||||
```
|
||||
|
||||
无 RDMA 时，另设 `NCCL_IB_DISABLE=1` 与 `NCCL_P2P_DISABLE=1`。
|
||||
|
||||
---
|
||||
|
||||
## 四、推理 — 稀疏少步学生
|
||||
|
||||
验收方式为独立运行 [examples/sparkdiffusion/predict_t2v.py](../../examples/sparkdiffusion/predict_t2v.py)，它会把每个 MoE 专家重建为稀疏 `SparkWan2_2Transformer3DModel` 并载入蒸馏 checkpoint。编辑脚本顶部：
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|-----------|-------------|---------------|
|
||||
| `GPU_memory_mode` | 显存模式 | `model_cpu_offload` |
|
||||
| `model_name` | Wan2.2-T2V-A14B 基座（VAE / UMT5 / 调度器） | `models/Diffusion_Transformer/Wan2.2-T2V-A14B` |
|
||||
| `student_topk` | **必须等于**训练时的 `--student_topk` | `0.05` |
|
||||
| `transformer_path` | **低噪**稀疏学生 checkpoint | `output_dir_sparkdiffusion_wan2.2_14b_low_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
|
||||
| `transformer_high_path` | **高噪**稀疏学生 checkpoint | `output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors` |
|
||||
| `sampler_name` | `Flow` 与少步 CFG-free 学生匹配 | `Flow` |
|
||||
| `num_inference_steps` | 学生的少步步数 | `4` |
|
||||
| `guidance_scale` | 蒸馏学生为 CFG-free | `1.0` |
|
||||
| `sample_size` / `video_length` / `fps` | `[高, 宽]` / 帧数 / 帧率 | `[240, 336]` / `81` / `16` |
|
||||
|
||||
```bash
|
||||
python examples/sparkdiffusion/predict_t2v.py
|
||||
```
|
||||
|
||||
> 💡 每个专家的载入顺序为 `from_pretrained → set_student_topk → load_state_dict`；顺序颠倒会把训练好的 RoLa 权重重新清零。路径为 `None` 的专家保持 dense——要得到完整的 SparkDiffusion 学生，请设置**两个**路径（优先用各专家的 Stage-2 蒸馏产物）。
|
||||
|
||||
> 📌 论文的 **Stage 3**（FP8/INT8 融合量化推理核）在本仓库中有意未接入。
|
||||
|
||||
---
|
||||
|
||||
## 五、更多资源
|
||||
|
||||
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
|
||||
- **论文**：SparkDiffusion，《Mitigating the High-Sparsity Trap》（arXiv:2609.23153）
|
||||
- **稀疏学生模型**：`videox_fun/models/spark_transformer3d.py`（`SparkWan2_2Transformer3DModel`）
|
||||
- **RoLa 稀疏注意力内核**：`videox_fun/models/attention_utils.py`
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,77 @@
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# Stage-1 RoLa sparse warm-up product (train_finetune_sparse_wan2.1.sh -> _1.3b_sparse_rola/checkpoint-400).
|
||||
# --student_topk below MUST match the Stage-1 value (0.10, i.e. sparsity 0.9).
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
# SparkDiffusion Stage-2 (CrossDistill) for the Wan2.1 single-expert T2V model: distill the sparse student into a
|
||||
# few-step, CFG-free generator against the dense teacher. --sparse_student turns the generator into
|
||||
# SparkWanTransformer3DModel while teacher + critic stay dense; --generator_transformer_path loads the Stage-1
|
||||
# sparse warm-up ckpt (no prefix stripping, unexpected keys must be 0). Wan2.1-T2V-1.3B is a single full-range
|
||||
# expert, so the student distills over the whole rf-time range [0, 1] (no MoE boundary, no init-rollout relay).
|
||||
# The host is a plain discrete-grid DMD (--denoising_step_indices_list 1000 750 500 250 + --flow_euler_rollout, mean
|
||||
# reduction). CrossDistill layers the official PCM diversity term on top as a self-contained additive loss
|
||||
# (--div_mode pcm_random --loss_scale_div 5.0): a phase-0 boundary-alignment MSE anchored on the paired real video
|
||||
# latent, rolled on its own continuous rf grid (--dcm_total_steps 48 --dcm_skipping_interval_steps 1,
|
||||
# --teacher_timestep_shift 5.0 = the Wan2.1 scheduler shift, --div_rf_boundary 0.933781 = phase-0 right edge,
|
||||
# --div_teacher_guidance 5.0). --loss_scale_div is in per-element-mean space (the official uses 10.0 in sum space);
|
||||
# set it to 0 to recover the plain-DMD baseline. Memory: --gradient_checkpointing keeps the three nets
|
||||
# (student/teacher/critic) off the activation OOM path (1.3B single expert needs no CPU offload; add
|
||||
# --fsdp_offload_params true only if it OOMs). No --max_train_steps: runs until you kill the procs.
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.1.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--video_repeat=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--max_grad_norm=0.05 \
|
||||
--random_hw_adapt \
|
||||
--training_with_video_token_length \
|
||||
--enable_bucket \
|
||||
--uniform_sampling \
|
||||
--denoising_step_indices_list 1000 750 500 250 \
|
||||
--flow_euler_rollout \
|
||||
--randomize_step_indices \
|
||||
--gen_update_interval 5 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 5.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 5.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--sparse_student \
|
||||
--student_topk=0.10 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,80 @@
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
# SparkDiffusion Stage-2 (CrossDistill) for the HIGH-noise expert of Wan2.2-T2V-A14B (an MoE T2V model):
|
||||
# distill the sparse student into a few-step, CFG-free generator against the dense teacher. --sparse_student
|
||||
# turns the generator (and the init-rollout model) into SparkWan2_2Transformer3DModel while teacher + critic
|
||||
# stay dense; --generator_transformer_path loads the Stage-1 sparse warm-up ckpt (no prefix stripping,
|
||||
# unexpected keys must be 0). --boundary_type="high" restricts the student to the high-noise rf segment: the
|
||||
# few-step grid is derived IN-SCRIPT by splitting --joint_backward_timesteps at the config boundary (0.875) and
|
||||
# keeping the knots >= boundary, so the high expert rolls [1.0 -> 0.933781 -> 0.875] with rf-native Euler steps
|
||||
# (--flow_euler_rollout) -- no contiguous backward simulation. The host is a plain discrete-grid DMD
|
||||
# (--uniform_sampling over that grid, mean reduction, generator/critic alternating via --gen_update_interval 5).
|
||||
# CrossDistill layers the official PCM diversity term on the HIGH expert only (--div_mode pcm_random
|
||||
# --loss_scale_div 10.0): a phase-0 boundary-alignment MSE anchored on the paired real-video latent, rolled on
|
||||
# its own continuous rf grid (--dcm_total_steps 48 --dcm_skipping_interval_steps 1, --teacher_timestep_shift 12.0
|
||||
# = the Wan2.2 scheduler shift, --div_rf_boundary 0.933781 = the high-segment first node, --div_teacher_guidance
|
||||
# 5.0). Set --loss_scale_div 0 to recover the plain-DMD baseline. Memory: --fsdp_offload_params true +
|
||||
# --gradient_checkpointing are mandatory (three 14B nets are optimizer-state dominated); resolution
|
||||
# (--fix_sample_size), not frame count, is the OOM lever. No --max_train_steps: runs until you kill the procs.
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
|
||||
--fsdp_offload_params true \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=1 \
|
||||
--video_sample_n_frames=81 \
|
||||
--fix_sample_size 480 832 \
|
||||
--train_batch_size=1 \
|
||||
--video_repeat=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--max_grad_norm=0.05 \
|
||||
--random_hw_adapt \
|
||||
--training_with_video_token_length \
|
||||
--enable_bucket \
|
||||
--uniform_sampling \
|
||||
--flow_euler_rollout \
|
||||
--gen_update_interval 5 \
|
||||
--joint_backward_timesteps 0.933781 0.875 0.608979 \
|
||||
--div_mode pcm_random \
|
||||
--loss_scale_div 10.0 \
|
||||
--dcm_total_steps 48 \
|
||||
--dcm_skipping_interval_steps 1 \
|
||||
--teacher_timestep_shift 12.0 \
|
||||
--div_rf_boundary 0.933781 \
|
||||
--div_teacher_guidance 5.0 \
|
||||
--boundary_type="high" \
|
||||
--real_guidance_scale=5.0 \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--sparse_student \
|
||||
--student_topk=0.05 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
@@ -0,0 +1,77 @@
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# Stage-1 RoLa sparse warm-up product for the LOW expert (train_finetune_sparse_wan2.2_low.sh ->
|
||||
# _low_sparse_rola/checkpoint-400).
|
||||
export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_low_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors"
|
||||
# Warm-start of the low expert's rollout from the trained HIGH distilled student
|
||||
# (train_distill_wan2.2_high.sh -> _high_sparse_rola_distill/checkpoint-<step>/diffusion_pytorch_model.safetensors).
|
||||
# Empty = cold start from pure noise; set it to the latest HIGH distill checkpoint for a proper relay.
|
||||
# --init_rollout_path is fault-tolerant (auto-finds the safetensors and strips generator/transformer/model
|
||||
# prefixes, so a distill checkpoint dir also works).
|
||||
export INIT_ROLLOUT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola_distill/checkpoint-4000/diffusion_pytorch_model.safetensors"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
# SparkDiffusion Stage-2 (CrossDistill) for the LOW-noise expert of Wan2.2-T2V-A14B: the same sparse-student
|
||||
# discrete-grid DMD as train_distill_wan2.2_high.sh, with two low-specific differences. (1) --boundary_type="low"
|
||||
# keeps only the --joint_backward_timesteps knots BELOW the config boundary (0.875), so the trainable low student
|
||||
# rolls [0.875 -> 0.608979] with rf-native Euler steps (--flow_euler_rollout). Because that segment starts at the
|
||||
# MoE boundary rather than pure noise, --init_rollout_path warm-starts it: a FROZEN copy of the trained HIGH
|
||||
# student (INIT_ROLLOUT_CKPT) first rolls the high segment [1.0 -> 0.933781 -> 0.875] to produce the boundary
|
||||
# latent, then the low student takes over. Leaving INIT_ROLLOUT_CKPT empty = cold start (the untrained low student
|
||||
# rolls the high segment itself and the samples degrade). (2) The low expert is div-free (no --div_mode /
|
||||
# --loss_scale_div): the PCM diversity term is gated on --loss_scale_div > 0 and lives on the HIGH expert only.
|
||||
# Memory levers are identical to train_distill_wan2.2_high.sh (--fsdp_offload_params true + --gradient_checkpointing;
|
||||
# lower --fix_sample_size, not frame count, on OOM). No --max_train_steps: runs until you kill the procs.
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_POST" \
|
||||
--fsdp_offload_params true \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_distill_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=1 \
|
||||
--video_sample_n_frames=81 \
|
||||
--fix_sample_size 480 832 \
|
||||
--train_batch_size=1 \
|
||||
--video_repeat=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=2e-06 \
|
||||
--learning_rate_critic=2e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_low_sparse_rola_distill" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--max_grad_norm=0.05 \
|
||||
--random_hw_adapt \
|
||||
--training_with_video_token_length \
|
||||
--enable_bucket \
|
||||
--uniform_sampling \
|
||||
--flow_euler_rollout \
|
||||
--gen_update_interval 5 \
|
||||
--joint_backward_timesteps 0.933781 0.875 0.608979 \
|
||||
--boundary_type="low" \
|
||||
--real_guidance_scale=5.0 \
|
||||
--init_rollout_path=$INIT_ROLLOUT_CKPT \
|
||||
--trainable_modules "." \
|
||||
--train_mode="normal" \
|
||||
--sparse_student \
|
||||
--student_topk=0.05 \
|
||||
--generator_transformer_path=$SPARSE_STUDENT_CKPT \
|
||||
--resume_from_checkpoint="latest"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,55 @@
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
# SparkDiffusion Stage-1 (sparse warm-up) for the Wan2.1 single-expert T2V model: insert the compensating
|
||||
# sparse self-attention (RoLa: blocks.*.self_attn.attn_op.{proj_q,proj_k,gate_proj,gate_bias}) into the dense
|
||||
# Wan2.1 backbone and briefly warm it up with the plain step-local flow-matching objective, so Stage-2
|
||||
# distillation starts from a usable coarse prior instead of a zero-initialized sparse branch. Wan2.1-T2V-1.3B is
|
||||
# a single full-range expert (no MoE boundary), so the student warms up over the whole rf-time range [0, 1].
|
||||
# --student_topk=0.10 (~90% sparsity) must match the value used later in train_distill_wan2.1.sh. Stage A trains
|
||||
# only "attn_op"; for Stage B (full sparse finetune) set --trainable_modules "." plus --transformer_path=<stage-A
|
||||
# ckpt> and a new --output_dir.
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.1.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--video_repeat=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--max_train_steps=2000 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=1e-4 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=1e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--max_grad_norm=1.0 \
|
||||
--random_hw_adapt \
|
||||
--training_with_video_token_length \
|
||||
--enable_bucket \
|
||||
--uniform_sampling \
|
||||
--trainable_modules "attn_op" \
|
||||
--student_topk=0.10 \
|
||||
--train_mode="normal" \
|
||||
--resume_from_checkpoint="latest"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,57 @@
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
# SparkDiffusion Stage-1 (sparse warm-up) for the HIGH-noise expert: insert the compensating sparse
|
||||
# self-attention (RoLa: blocks.*.self_attn.attn_op.{proj_q,proj_k,gate_proj,gate_bias}) into the dense
|
||||
# Wan2.2 backbone and briefly warm it up with the plain step-local flow-matching objective, so Stage-2
|
||||
# distillation starts from a usable coarse prior instead of a zero-initialized sparse branch.
|
||||
# --boundary_type="high" picks the high-noise expert from the config boundary (config/wan2.2/wan_civitai_t2v.yaml,
|
||||
# boundary:0.875); the low expert is warmed up separately by train_finetune_sparse_wan2.2_low.sh (Sec 4.3: each MoE
|
||||
# expert group is warmed up and distilled on its own noise range). --student_topk=0.05 (~95% sparsity) must
|
||||
# match the value used later in train_distill_wan2.2_high.sh. Stage A trains only "attn_op"; for Stage B (full sparse
|
||||
# finetune) set --trainable_modules "." plus --transformer_path=<stage-A ckpt> and a new --output_dir.
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--video_repeat=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--max_train_steps=2000 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=1e-4 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=1e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--max_grad_norm=1.0 \
|
||||
--random_hw_adapt \
|
||||
--training_with_video_token_length \
|
||||
--enable_bucket \
|
||||
--uniform_sampling \
|
||||
--trainable_modules "attn_op" \
|
||||
--student_topk=0.05 \
|
||||
--train_mode="normal" \
|
||||
--boundary_type="high" \
|
||||
--resume_from_checkpoint="latest"
|
||||
@@ -0,0 +1,53 @@
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
NCCL_DEBUG=INFO
|
||||
|
||||
# SparkDiffusion Stage-1 (sparse warm-up) for the LOW-noise expert: same recipe as
|
||||
# train_finetune_sparse_wan2.2_high.sh, only --boundary_type flips to "low" so the low-noise expert is picked from
|
||||
# the config boundary (config/wan2.2/wan_civitai_t2v.yaml, boundary:0.875) and warmed up on its own
|
||||
# noise range (Sec 4.3: each MoE expert group is warmed up and distilled separately). --student_topk=0.05
|
||||
# must match train_distill_wan2.2_low.sh. Its checkpoint feeds train_distill_wan2.2_low.sh as the low sparse student.
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False scripts/sparkdiffusion/train_finetune_sparse_wan2.2.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--video_repeat=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--max_train_steps=2000 \
|
||||
--checkpointing_steps=200 \
|
||||
--learning_rate=1e-4 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_sparkdiffusion_wan2.2_14b_low_sparse_rola" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=1e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--max_grad_norm=1.0 \
|
||||
--random_hw_adapt \
|
||||
--training_with_video_token_length \
|
||||
--enable_bucket \
|
||||
--uniform_sampling \
|
||||
--trainable_modules "attn_op" \
|
||||
--student_topk=0.05 \
|
||||
--train_mode="normal" \
|
||||
--boundary_type="low" \
|
||||
--resume_from_checkpoint="latest"
|
||||
@@ -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 |
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -12,11 +12,9 @@ This document provides the complete workflow for **Flex-Forcing Stage 2 — Dist
|
||||
> 2. **§3.2 Pyramid timestep chunking** — after finishing step `t+1`, the partition for step `t` is *refined* by inserting boundaries while preserving the existing ones. The step-`t` result is buffered over the whole original chunk, then each sub-chunk resumes autoregressively once the KV it needs is available. Coarse planning first, fine detail later.
|
||||
> 3. **§3.3 K-Projection** — clean-context keys live in a different space from noisy queries at timestep `t`, so an identity-initialised, timestep-conditioned projection maps the cache into the current noisy latent space. It is applied on the fly, never mutates the cache, and trains together with the generator.
|
||||
>
|
||||
> Training draws a **fresh random partition every iteration** (chunk sizes 2–10), which is what lets one set of weights cover the whole causal↔bidirectional spectrum. §4.2 of the paper then reuses the same machinery for **any-order editing**.
|
||||
>
|
||||
> The Flex-Forcing pipeline has two stages:
|
||||
>
|
||||
> 1. **Stage 1 — ODE Regression** (`train_ode.py`): regress the teacher's ODE trajectory so the model becomes a competent few-step generator. This is where the random partitions are first seen.
|
||||
> 1. **Stage 1 — ODE Regression** (`train_ode.py`): regress the teacher's ODE trajectory so the model becomes a competent few-step generator.
|
||||
> 2. **Stage 2 — DMD Distillation** (`train_distill.py`, this README): compress to a **4-step** generator with a **Wan2.1-T2V-14B** real-score teacher, and add the pyramid timestep schedule.
|
||||
>
|
||||
> This README covers **Stage 2 only**. See [README_TRAIN_ODE.md](./README_TRAIN_ODE.md) for Stage 1.
|
||||
@@ -210,7 +208,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
|
||||
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
|
||||
```
|
||||
|
||||
The launcher pins: 81 pixel frames = **21 latent frames** (a 5 s clip) at **480×832**, a **4-step** model (`--denoising_step_indices_list 1000 750 500 250`, jittered per iteration by `--randomize_step_indices`), `--flex_pyramid_levels=4` - which over 21 latent frames reproduces the paper's ladder `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` - partitions drawn in **2..10** (plus, on half the iterations, the whole-clip coarse layout of §3.2 - see below), `--flow_euler_rollout`, `diag_rank1` K-Projection, batch 8 (8 GPUs × `--gradient_accumulation_steps=1`).
|
||||
The launcher pins: 81 pixel frames = **21 latent frames** (a 5 s clip) at **480×832**, a **4-step** model (`--denoising_step_indices_list 1000 750 500 250`, jittered per iteration by `--randomize_step_indices`), `--flex_pyramid_levels=4` - which over 21 latent frames reproduces the paper's ladder `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` - the *random* arm's partitions drawn in **2..10** (that arm is only 10% of iterations; another 60% start from the whole-clip coarse layout of §3.2 - see below), `--flow_euler_rollout`, `diag_rank1` K-Projection, batch 8 (8 GPUs × `--gradient_accumulation_steps=1`).
|
||||
|
||||
Output: `output_dir_wan2.1_flex_forcing_distill/`.
|
||||
|
||||
@@ -243,7 +241,7 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
|
||||
| `--denoising_step_indices_list` | Denoising step indices (DMD core param). Mapped through `timesteps[train_sampling_steps - idx]`, so `1000 750 500 250` becomes t = `[1000, 937.5, 833.33, 625]` at `shift=5.0` - bit-identical to what `stochastic_sampling_timesteps(4, shift)` feeds the pipeline. | `1000 750 500 250` |
|
||||
| `--randomize_step_indices` | Jitter every index but the first symmetrically within `--index_jitter_ratio` of its neighbouring gap each iteration. The expected schedule is unchanged and monotonicity is enforced by construction, so inference's fixed t stays inside every step's trained envelope. | off |
|
||||
| `--index_jitter_ratio` | Jitter budget as a fraction of the neighbouring gap. At `1000 750 500 250` with 0.3, 20 000 draws give idx `[675,825] / [425,575] / [175,325]`, i.e. t `[912.2,959.3] / [787.0,871.2] / [514.7,706.5]`; the inference points 937.5 / 833.33 / 625 all fall inside. | 0.3 |
|
||||
| `--flow_euler_rollout` | Advance the generator's self-rollout with a deterministic Euler ODE step in flow space (`x_next = x_t + (sigma_next - sigma_t) * v`, fp32) instead of converting to x0 and re-noising with fresh noise. Algebraically this *is* re-noising with the same implied noise - measured `8.9e-16` apart in fp64, versus up to `6.4` for a fresh draw - so it removes the rollout's stochasticity and matches normal flow-matching inference. The final step still converts to x0, which is what DMD is defined on. The critic re-rolls the generator and follows the same switch. | off |
|
||||
| `--flow_euler_rollout` | Advance the generator's self-rollout with a deterministic Euler ODE step in flow space (`x_next = x_t + (sigma_next - sigma_t) * v`, fp32) instead of converting to x0 and re-noising with fresh noise. Algebraically this *is* re-noising with the same implied noise - identical to it in exact arithmetic, whereas drawing fresh noise is an independent sample - so it removes the rollout's stochasticity and matches normal flow-matching inference. The final step still converts to x0, which is what DMD is defined on. The critic re-rolls the generator and follows the same switch. | off |
|
||||
| `--output_dir` | Output directory | `output_dir_wan2.1_flex_forcing_distill` |
|
||||
| `--gradient_checkpointing` | Trade compute for memory | - |
|
||||
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
|
||||
@@ -256,24 +254,14 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
|
||||
| Parameter | Paper | Description | Default |
|
||||
|-----------|-------|-------------|---------|
|
||||
| `--flex_forcing` | §3.1 | Instantiate `WanTransformer3DModel_FlexForcing` and validate with `WanFlexForcingPipeline`. Off = the inherited path, unchanged. | off |
|
||||
| `--flex_chunk_min` | §3.1 | Smallest chunk drawn per iteration. `1` is legal but adds no coverage under a pyramid - refining a 2-frame chunk already yields 1-frame leaves - and it would let single-frame chunks appear mid-clip. | 2 |
|
||||
| `--flex_chunk_max` | §3.1 | Largest chunk drawn per iteration. Set equal to `--flex_chunk_min` to pin one fixed layout. **Keep it below the latent frame count**: the whole-clip layout is already covered by the mixture below, so raising this to the frame count only spends draws on a layout you already have and thins out §3.1 (at 21 latent frames, `max=21` leaves 3649 distinct partitions versus 4882 at `max=10`). | 10 |
|
||||
| `--flex_pyramid_levels` | §3.2 | `1` = one partition per iteration. `>1` = nest it into a coarse-to-fine ladder, one level per denoising step; both training paths roll it out. **Stage 2 only.** | 1 |
|
||||
| `--flex_chunk_min` | §5 | Smallest chunk drawn by the *random* arm (10% of iterations, regardless of `--flex_pyramid_levels`); the other three arms pin their own layouts and split the remaining 90% equally. `1` is legal but adds no coverage under a pyramid - refining a 2-frame chunk already yields 1-frame leaves - and it would let single-frame chunks appear mid-clip. | 2 |
|
||||
| `--flex_chunk_max` | §5 | Largest chunk drawn by the *random* arm. Set equal to `--flex_chunk_min` to pin that arm to one fixed layout. **Keep it below the latent frame count**: the whole-clip layout is its own arm, so raising this to the frame count only spends random-arm draws on a layout you already have and thins out every other one (exhaustive enumeration at 21 latent frames: 6441 distinct partitions at `max=10`, 6765 at `max=21` - but `max=21` makes `[21]` alone 1/19 of that arm and cuts `P([3]*7)` from `1/183708` to `1/1106560`). | 10 |
|
||||
| `--flex_pyramid_levels` | §3.2 | `1` = one partition per iteration. `>1` = nest it into a coarse-to-fine ladder, one level per denoising step; both training paths roll it out. The ladder itself is **Stage 2 only** (Stage 1 regresses `x0` in one forward per chunk, so it has no ladder to build), but Stage 1 **must still carry the same value**: it decides whether the shared level-0 mixture's coarse arm exists. | 1 |
|
||||
| `--flex_min_num_frame_per_block` | §3.2 | Block size the refinement stops at; `1` is fully causal at the leaves. | 1 |
|
||||
| `--flex_self_generated_context` | §3.3 | Feeds the K-Projection the model's own previous-step `x0` prediction as context, so it enters the autograd graph under prompt-only training. Costs a doubled attention sequence; from denoising step 1 on. Inert under `--use_kv_cache_training` - that path passes no `flex_state`, so the K-Projection is never invoked and stays at its identity init. | off |
|
||||
| `--num_frame_per_block` | — | The uniform fallback for the inherited (non-Flex) path; with Flex on it is also one of the three level-0 bands (`UNIFORM_BLOCK_PROB = 0.1` of iterations pin level 0 to its uniform partition), and it is what `log_validation` samples with at `--flex_pyramid_levels=1` - see the two bullets below. | 3 |
|
||||
| `--flex_self_generated_context` | §3.3 | Feeds the K-Projection the model's own previous-step `x0` prediction as context, so it enters the autograd graph under prompt-only training. Costs a doubled attention sequence; from denoising step 1 on. Inert under `--use_kv_cache_training`, but **not because the K-Projection is idle** - that path already commits its own clean context into the KV cache and needs no extra `clean_x` half; Π itself **is** invoked there. | off |
|
||||
| `--num_frame_per_block` | — | The uniform fallback for the inherited (non-Flex) path; with Flex on it is also reached inside level 0's four-arm mixture (the uniform arm is 30% of iterations and one of its three equal sub-arms pins level 0 to its uniform partition, so `[nfpb x k]` alone is 10%), and it is what `log_validation` samples with at `--flex_pyramid_levels=1`. | 3 |
|
||||
| `--independent_first_frame` | — | Prepend a length-1 chunk (`[1, N, N, ...]`); honoured by the partition sampler. | off |
|
||||
|
||||
Seven interactions are worth stating explicitly:
|
||||
|
||||
- **Level 0 is a three-band mixture decided by one draw.** `denoise_mode="pyramid"` at inference builds its level 0 as the whole clip in one chunk (fully bidirectional planning), so a partition drawn only from 2..10 would leave that step untrained - and the uniform `--num_frame_per_block` layout is just as unreachable by chance alone (at 21 latent frames `[3]*7` has probability `1/183708`, i.e. 0.11 expected hits in 20 000 iterations). `train_distill.py` therefore splits level 0 with a single uniform draw into three bands, neither constant a CLI flag: `COARSE_GLOBAL_PROB = 0.5` for the whole clip as one chunk, `UNIFORM_BLOCK_PROB = 0.1` pinned to the uniform `--num_frame_per_block` partition, and the remaining `0.4` for §3.1's random draw; all three refine into the same ladder. Measured over 200 000 iterations at 21 latent frames with `--flex_pyramid_levels=4`: whole clip `49.999%`, `[3]*7` `10.034%`, random `39.968%`, **4078** distinct level-0 layouts, and the inference ladder `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` at `49.999%`. At `--flex_pyramid_levels=1` the whole-clip band is empty and the split becomes `10% / 90%`, which keeps that mode as the paper describes it. Peak memory is set by the densest case (level 0 = whole clip), so budget for it.
|
||||
- The ladder is **clamped to the number of trained denoising steps**. A 4-level ladder under a 2-step schedule can never reach its last two levels, so `sample_flex_partitions` drops the unreachable tail rather than silently reporting it as trained. Keep `len(--denoising_step_indices_list) >= --flex_pyramid_levels`; the launcher's 4 indices and `--flex_pyramid_levels=4` are matched on purpose. The converse also wastes compute: more steps than levels leaves `install_flex_partition` clamped on its finest level for the extra steps.
|
||||
- **`--flow_euler_rollout` keeps the prediction in flow space between steps, so §3.3's context needs its own x0.** `--flex_self_generated_context` feeds the K-Projection the previous step's *clean* prediction; under an Euler rollout that variable is a velocity, so both rollout loops convert a detached copy for the context while the rollout state itself advances by Euler. The generator and the critic convert identically, otherwise the critic would score a rollout that is not the one being trained.
|
||||
- **`--randomize_step_indices` does not break the pyramid.** The jitter moves the *timesteps*, not the partitions, and `sample_flex_partitions` keys off `len(denoising_step_list)`, which the jitter preserves. Every level of the ladder still lines up with its denoising step.
|
||||
- **`log_validation` renders the layout training renders.** With `--flex_pyramid_levels > 1` it passes `chunk_spec=None` + `denoise_mode="pyramid"`, so its ladder is identical to `predict_t2v.py`'s and measured in `49.999%` of iterations. It used to pin `chunk_spec` to `--num_frame_per_block`, which nested `[[3]*7, [2,1]*7, [1]*21]` - a **3-level** ladder under a 4-step schedule (refinement hits the all-ones fixed point) that the trainer drew **0 times in 20 000** iterations, so those clips said nothing about the model being trained. At `--flex_pyramid_levels=1` the conclusion reverses: the whole-clip band is empty there and `[21]` is never drawn, so validation uses the uniform `--num_frame_per_block` partition - the one band that is both fixed across checkpoints and genuinely trained (`10.066%` of iterations, against ~`1.3%` for the most common random partition). On that 3-level ladder, now covered `10.034%` of the time: step 4 reuses the finest level `[1]*21` through `install_flex_partition`'s index clamp (`min(step_index, len(partitions) - 1)`), so nothing runs off the end.
|
||||
- §3.3 has **no launcher flag**: the config yaml chooses the variant through `transformer_additional_kwargs.flex_kproj_mode` (default `diag_rank1`; set it to `none` there to ablate Pi, `diag` for the cheaper head-only map). Its parameters are named `flex_kproj.proj_out.*` on purpose - that name is what lets `initialize_missing_parameters()` zero-fill them when a checkpoint predates §3.3, which is how identity init survives `low_cpu_mem_usage=True` - and they otherwise train like any other layer, following `--learning_rate`. The trainer forces them `requires_grad` whenever the built model carries them, so a narrowed `--trainable_modules` cannot silently freeze Pi at its identity. Under `--use_kv_cache_training` it is never invoked and stays at its identity init as well.
|
||||
- Every rank must train the same layout, otherwise the FlexAttention mask and the `num_frame_per_block` derived from it disagree across the SP/FSDP group. The trainer therefore passes the drawn partition through `broadcast_chunk_sizes` before using it, and `randomize_denoising_step_indices` broadcasts its jittered indices from rank 0 for the same reason.
|
||||
|
||||
### 3.4 Training Validation
|
||||
|
||||
You can configure validation parameters to periodically generate test videos during training, allowing you to monitor training progress and model quality.
|
||||
@@ -757,7 +745,7 @@ num_frame_per_block = 7 # clean-context commit granularity; None = one
|
||||
|
||||
This `num_frame_per_block` is the same knob as Self-Forcing's: the clean context is committed in uniform blocks of that many frames, in temporal order. It differs in exactly two ways - the trailing chunk absorbs the remainder so divisibility is not required (the inherited rollout asserts it), and the transformer's block width becomes `max(it, edit span width)`.
|
||||
|
||||
Any-order editing needs the full KV cache: both entries set `local_attn_size = -1`. A pyramid needs the same, plus `stochastic_sampling = True` — the buffered step between two levels *is* the schedule's own re-noising. The pipeline raises rather than silently degrading if either is violated.
|
||||
Any-order editing needs the full KV cache: both entries set `local_attn_size = -1`. Any per-step ladder (`denoise_mode="pyramid"` or `"full_then_blocks"`) needs the same, plus `stochastic_sampling = True` — the buffered step between two levels *is* the schedule's own re-noising. The pipeline raises rather than silently degrading if either is violated.
|
||||
|
||||
---
|
||||
|
||||
@@ -765,7 +753,7 @@ Any-order editing needs the full KV cache: both entries set `local_attn_size = -
|
||||
|
||||
- **Flex-Forcing paper**: https://arxiv.org/abs/2607.03509
|
||||
- **Stage 1**: [README_TRAIN_ODE.md](./README_TRAIN_ODE.md)
|
||||
- **Shared chunking helpers**: `videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`, `sample_flexible_chunks`, `refine_partition`, `build_pyramid_partitions`, `validate_nested_partitions`, `broadcast_chunk_sizes`
|
||||
- **Shared chunking helpers**: `videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`, `sample_flexible_chunks`, `sample_flex_level0` (the level-0 four-arm mixture, drawn by **both** trainers), `refine_partition`, `build_pyramid_partitions`, `validate_nested_partitions`, `broadcast_chunk_sizes`
|
||||
- **Model / pipeline**: `videox_fun/models/wan_transformer3d_flex_forcing.py`, `videox_fun/pipeline/pipeline_wan_flex_forcing.py`
|
||||
- **Upstream stages**: [README_TRAIN_ODE.md](../wan2.1_self_forcing/README_TRAIN_ODE.md), [README_TRAIN_DISTILL.md](../wan2.1_self_forcing/README_TRAIN_DISTILL.md)
|
||||
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
|
||||
|
||||
@@ -12,11 +12,9 @@
|
||||
> 2. **§3.2 金字塔时间步分块** —— 完成步 `t+1` 后，步 `t` 的划分是在**保留已有边界的前提下插入新边界**细化的。步 `t` 的结果先在整个原 chunk 上缓冲，随后每个子 chunk 在其所需 KV 可用后逐个自回归续写。先粗规划、后细刻画。
|
||||
> 3. **§3.3 K-Projection** —— 干净上下文的 key 与时间步 `t` 的含噪 query 处在不同的空间，因此用一个恒等初始化、以时间步为条件的投影把 cache 映射进当前含噪 latent 空间。它即时施加、从不改写 cache，并与生成器一起训练。
|
||||
>
|
||||
> 训练时每次迭代都抽取**全新的随机划分**（chunk 大小 2–10），这正是让同一套权重覆盖整个"因果↔双向"谱系的原因。论文 §4.2 随后把同一套机制复用为**任意序编辑**。
|
||||
>
|
||||
> Flex-Forcing 流程分两个阶段：
|
||||
>
|
||||
> 1. **阶段 1——ODE 回归**（`train_ode.py`）：对教师模型的 ODE 轨迹做回归，让模型先成为一个合格的小步数生成器。随机分块在这一阶段就首次出现。
|
||||
> 1. **阶段 1——ODE 回归**（`train_ode.py`）：对教师模型的 ODE 轨迹做回归，让模型先成为一个合格的小步数生成器。
|
||||
> 2. **阶段 2——DMD 蒸馏**（`train_distill.py`，即本 README）：用 **Wan2.1-T2V-14B** 真实分数教师蒸馏出 **4 步**生成器，并加入金字塔时间步调度。
|
||||
>
|
||||
> 本 README 只覆盖**阶段 2**。阶段 1 见 [README_TRAIN_ODE_zh-CN.md](./README_TRAIN_ODE_zh-CN.md)。
|
||||
@@ -210,7 +208,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
|
||||
--ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
|
||||
```
|
||||
|
||||
launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒片段）、分辨率 **480×832**、**4 步**模型（`--denoising_step_indices_list 1000 750 500 250`，并由 `--randomize_step_indices` 每次迭代抖动）、`--flex_pyramid_levels=4` —— 在 21 个 latent 帧上正好复现论文的阶梯 `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` —— 划分抽取范围 **2..10**（另有半数的迭代以 §3.2 的整段粗粒度布局起步，见下）、`--flow_euler_rollout`、`diag_rank1` K-Projection、batch 8（8 卡 × `--gradient_accumulation_steps=1`）。
|
||||
launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒片段）、分辨率 **480×832**、**4 步**模型（`--denoising_step_indices_list 1000 750 500 250`，并由 `--randomize_step_indices` 每次迭代抖动）、`--flex_pyramid_levels=4` —— 在 21 个 latent 帧上正好复现论文的阶梯 `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` —— 随机档的划分抽取范围 **2..10**（该档只占 10% 的迭代；另有 60% 以 §3.2 的整段粗粒度布局起步，见下）、`--flow_euler_rollout`、`diag_rank1` K-Projection、batch 8（8 卡 × `--gradient_accumulation_steps=1`）。
|
||||
|
||||
输出目录：`output_dir_wan2.1_flex_forcing_distill/`。
|
||||
|
||||
@@ -243,7 +241,7 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
|
||||
| `--denoising_step_indices_list` | 去噪步索引（DMD 核心参数）。经 `timesteps[train_sampling_steps - idx]` 映射，所以 `shift=5.0` 下 `1000 750 500 250` 变成 t = `[1000, 937.5, 833.33, 625]` —— 与 `stochastic_sampling_timesteps(4, shift)` 喂给 pipeline 的值逐位一致。 | `1000 750 500 250` |
|
||||
| `--randomize_step_indices` | 每次迭代对除首个索引以外的所有索引做对称抖动，幅度为相邻间隔的 `--index_jitter_ratio`。期望调度不变，单调性由构造保证，所以推理用的固定 t 始终落在每一步的训练包络内。 | 关 |
|
||||
| `--index_jitter_ratio` | 抖动预算（相邻间隔的比例）。在 `1000 750 500 250` 且取 0.3 时，20000 次抽取给出 idx `[675,825] / [425,575] / [175,325]`，即 t `[912.2,959.3] / [787.0,871.2] / [514.7,706.5]`；推理点 937.5 / 833.33 / 625 全部落在其中。 | 0.3 |
|
||||
| `--flow_euler_rollout` | 生成器自滚用 flow 空间里的确定性 Euler ODE 步推进（`x_next = x_t + (sigma_next - sigma_t) * v`，fp32），而不是转成 x0 再用新噪声重加噪。代数上这**就等于**用同一个隐含噪声重加噪 —— fp64 下实测相差 `8.9e-16`，而重采新噪声最大差 `6.4` —— 因此它消掉了 rollout 的随机性，与正常的 flow-matching 推理一致。最后一步仍然转 x0，因为 DMD 目标定义在 x0 上。critic 重滚生成器，跟随同一开关。 | 关 |
|
||||
| `--flow_euler_rollout` | 生成器自滚用 flow 空间里的确定性 Euler ODE 步推进（`x_next = x_t + (sigma_next - sigma_t) * v`，fp32），而不是转成 x0 再用新噪声重加噪。代数上这**就等于**用同一个隐含噪声重加噪 —— 在精确算术下与之完全相同，而重采新噪声是一次独立抽样 —— 因此它消掉了 rollout 的随机性，与正常的 flow-matching 推理一致。最后一步仍然转 x0，因为 DMD 目标定义在 x0 上。critic 重滚生成器，跟随同一开关。 | 关 |
|
||||
| `--output_dir` | 输出目录 | `output_dir_wan2.1_flex_forcing_distill` |
|
||||
| `--gradient_checkpointing` | 用算力换显存 | - |
|
||||
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
|
||||
@@ -256,24 +254,14 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
|
||||
| 参数 | 论文 | 说明 | 默认值 |
|
||||
|-----------|-------|-------------|---------|
|
||||
| `--flex_forcing` | §3.1 | 实例化 `WanTransformer3DModel_FlexForcing` 并用 `WanFlexForcingPipeline` 校验。关闭 = 走继承的原路径，行为不变。 | 关 |
|
||||
| `--flex_chunk_min` | §3.1 | 每次迭代抽取的最小 chunk 大小。金字塔下 `1` 合法但不增加覆盖：2 帧 chunk 细化后本来就会产生 1 帧叶子，而它会让单帧 chunk 出现在片段中间。 | 2 |
|
||||
| `--flex_chunk_max` | §3.1 | 每次迭代抽取的最大 chunk 大小。设为与 `--flex_chunk_min` 相等即固定一种布局。**必须小于 latent 帧数**：“整段一个 chunk”这个布局已经由下面的混合机制覆盖，把这个值抬到帧数只是把抽取次数花在一个你已经有的布局上，反而稀释 §3.1（21 个 latent 帧时，`max=21` 只剩 3649 种不同划分，而 `max=10` 有 4882 种）。 | 10 |
|
||||
| `--flex_pyramid_levels` | §3.2 | `1` = 每次迭代一个划分。`>1` = 把它嵌套成由粗到细的阶梯，每个去噪步一层；两条训练路径都会展开。**仅阶段 2 有。** | 1 |
|
||||
| `--flex_chunk_min` | §5 | **random 档**（占 10% 迭代，与 `--flex_pyramid_levels` 无关）抽取的最小 chunk 大小；其余三档各自钉定自己的布局、平分剩下的 90%。金字塔下 `1` 合法但不增加覆盖：2 帧 chunk 细化后本来就会产生 1 帧叶子，而它会让单帧 chunk 出现在片段中间。 | 2 |
|
||||
| `--flex_chunk_max` | §5 | **random 档**抽取的最大 chunk 大小。设为与 `--flex_chunk_min` 相等即把该档钉成一种固定布局。**必须小于 latent 帧数**：“整段一个 chunk”本身就是独立的一档，把这个值抬到帧数只是把 random 档的抽取次数花在一个你已经有的布局上，反而稀释其他所有布局（21 个 latent 帧下穷举计数：`max=10` 有 6441 种不同划分、`max=21` 有 6765 种，但 `max=21` 会让单是 `[21]` 就占该档 1/19 的抽取，并把 `P([3]*7)` 从 1/183708 压到 1/1106560）。 | 10 |
|
||||
| `--flex_pyramid_levels` | §3.2 | `1` = 每次迭代一个划分。`>1` = 把它嵌套成由粗到细的阶梯，每个去噪步一层；两条训练路径都会展开。阶梯本身**仅阶段 2 有**（阶段 1 每 chunk 一次前向回归 x0，没有阶梯可搭），但阶段 1 **仍然要带同一个值**：它决定共享的 level-0 混合里 coarse 档是否存在。 | 1 |
|
||||
| `--flex_min_num_frame_per_block` | §3.2 | 细化终止的块大小（每块二分直到不超过它）；`1` 表示叶子级完全因果。 | 1 |
|
||||
| `--flex_self_generated_context` | §3.3 | 把模型自己上一步的 `x0` 预测喂给 K-Projection 当上下文，使其在纯 prompt 训练下也能进入 autograd 图。代价是注意力序列翻倍；从去噪步 1 起生效。在 `--use_kv_cache_training` 下不生效 —— 该路径不传 `flex_state`，K-Projection 不会被调用，停在恒等初始化。 | 关 |
|
||||
| `--num_frame_per_block` | — | 继承路径（未开 Flex）的均匀 fallback；开了 Flex 后它还是 level 0 三档混合中的一档（`UNIFORM_BLOCK_PROB = 0.1` 的迭代把 level 0 钉成它的均匀划分），且 `--flex_pyramid_levels=1` 时 `log_validation` 就用它采样 —— 见下面两条 bullet。 | 3 |
|
||||
| `--flex_self_generated_context` | §3.3 | 把模型自己上一步的 `x0` 预测作为上下文喂给 K-Projection，因此在纯 prompt 训练下它会进入自动求导图。代价是注意力序列翻倍；从去噪第 1 步起生效。在 `--use_kv_cache_training` 下它不起作用，但**并不是因为 K-Projection 空闲** —— 那条路径已经把自己的干净上下文写进 KV cache，不需要额外的 `clean_x` 一半；Π 本身在那里**是**被调用的。 | 关 |
|
||||
| `--num_frame_per_block` | — | 继承路径（未开 Flex）的均匀 fallback；开了 Flex 后它还是 level 0 四档混合中 uniform 那一档（该档占 30% 的迭代、内部再三等分，其一把 level 0 钉成它的均匀划分，故 `[nfpb x k]` 单独占 10%），且 `--flex_pyramid_levels=1` 时 `log_validation` 就用它采样。 | 3 |
|
||||
| `--independent_first_frame` | — | 在划分前插入一个长度为 1 的 chunk（`[1, N, N, ...]`）；划分采样器会遵守它。 | 关 |
|
||||
|
||||
有七个交互值得明说：
|
||||
|
||||
- **level 0 是三档混合，一次抽样定档。** 推理时 `denoise_mode="pyramid"` 的 level 0 是整段一个 chunk（全双向的全局规划），只从 2..10 抽划分会让那一步从未被训练；而 `--num_frame_per_block` 那个均匀布局光靠随机抽样又几乎碰不到（21 个 latent 帧下 `[3]*7` 的概率是 `1/183708`，20000 次迭代期望 0.11 次）。所以 `train_distill.py` 用一个均匀抽样把 level 0 分成三档，两个常量都故意不做成 CLI 参数：`COARSE_GLOBAL_PROB = 0.5` 整段一个 chunk、`UNIFORM_BLOCK_PROB = 0.1` 钉成 `--num_frame_per_block` 的均匀划分、剩下 `0.4` 是 §3.1 的随机划分；三档都细化成同一条阶梯。实测 200000 次（21 个 latent 帧、`--flex_pyramid_levels=4`）：整段 `49.999%`、`[3]*7` `10.034%`、随机 `39.968%`，共 **4078** 种不同的 level-0 布局，推理那条阶梯 `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` 占 `49.999%`。`--flex_pyramid_levels=1` 时整段那一档为空，比例变成 `10% / 90%`，与论文描述一致。峰值显存由最密的那种情形（level 0 = 整段）决定，请据此预留。
|
||||
- 阶梯会被**截断到训练的去噪步数**。2 步调度下配 4 级阶梯永远到不了最后两级，所以 `sample_flex_partitions` 会丢弃到不了的那部分，而不是默默把它们当作已训练。请保持 `len(--denoising_step_indices_list) >= --flex_pyramid_levels`；launcher 里的 4 个索引与 `--flex_pyramid_levels=4` 就是故意对齐的。反过来也浪费算力：步数多于级数时，多出来的那些步里 `install_flex_partition` 会被钉在最细的一级上。
|
||||
- **`--flow_euler_rollout` 会把预测留在 flow 空间，所以 §3.3 的上下文需要自己那份 x0。** `--flex_self_generated_context` 喂给 K-Projection 的是上一步的**干净**预测；在 Euler rollout 下那个变量是速度，所以两个 rollout 循环都会另外转换一份 detached 副本给上下文用，而 rollout 状态本身仍按 Euler 推进。生成器与 critic 的转换必须完全一致，否则 critic 评的就不是正在被训的那个 rollout。
|
||||
- **`--randomize_step_indices` 不会破坏金字塔。** 抖动动的是**时间步**而不是划分，而 `sample_flex_partitions` 只看 `len(denoising_step_list)`，抖动不会改变它。阶梯的每一级仍然对得上自己的去噪步。
|
||||
- **`log_validation` 渲染的就是训练渲染的那个布局。** `--flex_pyramid_levels > 1` 时它传 `chunk_spec=None` + `denoise_mode="pyramid"`，阶梯与 `predict_t2v.py` 的完全相同，实测占 `49.999%` 的迭代。它原本把 `chunk_spec` 钉成 `--num_frame_per_block`，那会嵌套出 `[[3]*7, [2,1]*7, [1]*21]` —— 在 4 步调度下只有 **3 级**（细化到全 1 就撞上固定点），而当时那条阶梯 trainer 在 20000 次迭代里抽出 **0 次**，所以那些校验视频根本说明不了正在训练的模型。`--flex_pyramid_levels=1` 时结论反过来：整段那一档为空、`[21]` 一次也抽不到，所以校验用 `--num_frame_per_block` 的均匀划分 —— 它是唯一既跨 checkpoint 固定、又真被训练的一档（`10.066%` 的迭代，而最高频的随机划分只有约 `1.3%`）。顺带说明那条 3 级阶梯：它现在有 `10.034%` 的覆盖率，第 4 步由 `install_flex_partition` 的索引 clamp（`min(step_index, len(partitions) - 1)`）复用最细一级 `[1]*21`，不会越界。
|
||||
- §3.3 **没有 launcher 参数**：变体由 config yaml 里的 `transformer_additional_kwargs.flex_kproj_mode` 决定（默认 `diag_rank1`；要消融就在那里改成 `none`，想更省就改 `diag`）。它的参数故意取名 `flex_kproj.proj_out.*` —— 正是这个名字让 `initialize_missing_parameters()` 在 checkpoint 早于 §3.3 时将其置零，从而在 `low_cpu_mem_usage=True` 下保住恒等初始化；除此之外它和其他层一样训，跟 `--learning_rate`。只要构建出的模型带 Pi，trainer 就强制其 `requires_grad`，所以收窄 `--trainable_modules` 也不会把 Pi 悄悄冻在恒等点。在 `--use_kv_cache_training` 下它同样不会被调用，始终停在恒等初始化。
|
||||
- 每张卡必须训练同一种布局，否则 FlexAttention 掩码与其派生的 `num_frame_per_block` 在 SP/FSDP 组之间会不一致。trainer 在使用抽取的划分前会先经 `broadcast_chunk_sizes` 广播对齐；`randomize_denoising_step_indices` 同样从 rank 0 广播它抖动后的索引，道理一样。
|
||||
|
||||
---
|
||||
|
||||
### 3.4 训练验证
|
||||
@@ -746,15 +734,11 @@ num_frame_per_block = 7 # 干净上下文的提交粒度；None = 整
|
||||
|
||||
这里的 `num_frame_per_block` 与 Self-Forcing 那个同名参数是一回事：干净上下文按这么多帧一块、依时间序提交。区别只有两点——尾块吸收余数所以不要求整除（继承的 rollout 会断言整除），且 transformer 上的块宽取 `max(它, 编辑区间宽度)`。
|
||||
|
||||
任意序编辑需要完整 KV cache：两个入口都设 `local_attn_size = -1`。金字塔还需要 `stochastic_sampling = True`——两级之间的缓冲步本身就是调度自己的重加噪。任何一项不满足时 pipeline 会直接报错，而不是默默降级。
|
||||
任意序编辑需要完整 KV cache：两个入口都设 `local_attn_size = -1`。任何逐步换划分的 ladder（`denoise_mode="pyramid"` 或 `"full_then_blocks"`）受同一条约束，并额外要求 `stochastic_sampling = True`——两级之间的缓冲步本身就是调度自己的重加噪。任何一项不满足时 pipeline 会直接报错，而不是默默降级。
|
||||
|
||||
---
|
||||
|
||||
## 五、更多资源
|
||||
|
||||
- **Flex-Forcing 论文**：https://arxiv.org/abs/2607.03509
|
||||
- **阶段 1**：[README_TRAIN_ODE_zh-CN.md](./README_TRAIN_ODE_zh-CN.md)
|
||||
- **共享分块 helper**：`videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`、`sample_flexible_chunks`、`refine_partition`、`build_pyramid_partitions`、`validate_nested_partitions`、`broadcast_chunk_sizes`
|
||||
- **模型 / pipeline**：`videox_fun/models/wan_transformer3d_flex_forcing.py`、`videox_fun/pipeline/pipeline_wan_flex_forcing.py`
|
||||
- **上游阶段**：[README_TRAIN_ODE.md](../wan2.1_self_forcing/README_TRAIN_ODE.md)、[README_TRAIN_DISTILL.md](../wan2.1_self_forcing/README_TRAIN_DISTILL.md)
|
||||
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
|
||||
|
||||
@@ -7,9 +7,9 @@ This document provides the complete workflow for **Flex-Forcing Stage 1 — ODE
|
||||
> ODE regression is the **pre-training stage** of the two-stage Flex-Forcing pipeline. The pipeline is:
|
||||
>
|
||||
> 1. **Step 1 — Generate ODE pairs** (`generate_ode_pairs.py`): Use the **bidirectional teacher** (Wan2.1-T2V-1.3B) to perform full multi-step CFG denoising on a list of text prompts. Save the intermediate latents along the ODE trajectory together with the encoded prompt embeddings as `.safetensors` files.
|
||||
> 2. **Step 2 — Train ODE regression** (`train_ode.py`): Load the generated ODE pairs and train a **flexible-chunking generator** to predict the clean endpoint `x0` at multiple sampled trajectory points, drawing a fresh random frame partition every iteration. The output checkpoint serves as the few-step initialization for the subsequent **Stage 2 DMD distillation** (`train_distill.py`, see [README_TRAIN_FLEX_DMD_DISTILL.md](./README_TRAIN_FLEX_DMD_DISTILL.md)).
|
||||
> 2. **Step 2 — Train ODE regression** (`train_ode.py`): Load the generated ODE pairs and train a **flexible-chunking generator** to predict the clean endpoint `x0` at multiple sampled trajectory points, drawing a fresh frame partition every iteration from the **same four-arm mixture Stage 2 takes its level 0 from**, so that no layout Stage 2 distills on is one this stage never regressed. The output checkpoint serves as the few-step initialization for the subsequent **Stage 2 DMD distillation** (`train_distill.py`, see [README_TRAIN_FLEX_DMD_DISTILL.md](./README_TRAIN_FLEX_DMD_DISTILL.md)).
|
||||
>
|
||||
> Flex-Forcing itself replaces Self-Forcing's scalar `num_frame_per_block` with a **partition of the frame axis** `a = (a_0, ..., a_K)`: attention is bidirectional *inside* a chunk and autoregressive *across* chunks. `[1] * F` recovers pure autoregression, `[F]` recovers pure bidirectional attention, and everything in between is covered by one model. The random partitions are first seen in this stage.
|
||||
> Flex-Forcing itself replaces Self-Forcing's scalar `num_frame_per_block` with a **partition of the frame axis** `a = (a_0, ..., a_K)`: attention is bidirectional *inside* a chunk and autoregressive *across* chunks. `[1] * F` recovers pure autoregression, `[F]` recovers pure bidirectional attention, and everything in between is covered by one model. Both extremes are drawn here by explicit arms rather than left to chance: a random draw reaches `[F]` at best `1/19` of the time (and never at all while `--flex_chunk_max` sits below the frame count), and `[1] * F` with probability `1/F!`.
|
||||
|
||||
---
|
||||
|
||||
@@ -208,9 +208,22 @@ After Step 1 completes and `datasets/ode_pairs_output/outputs.json` is generated
|
||||
|
||||
For each training sample the script:
|
||||
1. Loads the saved sparse trajectory (5 points) and prompt embedding from one `.safetensors` file.
|
||||
2. Draws the frame-axis partition and gives **one timestep per chunk** (`get_timestep_for_ode_flexible`) — this is what makes a chunk internally bidirectional. 90% of iterations draw a **fresh random partition** (chunk sizes `--flex_chunk_min`..`--flex_chunk_max`); the other `UNIFORM_BLOCK_PROB = 0.1` pin it to the uniform `--num_frame_per_block` layout, which a random draw would reach with probability `1/183708` at 21 latent frames and which is exactly what `log_validation` renders.
|
||||
2. Draws the frame-axis partition and gives **one timestep per chunk** (`get_timestep_for_ode_flexible`) — this is what makes a chunk internally bidirectional.
|
||||
3. Feeds the noisy latent and the per-frame timestep through the generator, converts the predicted flow into an `x0` prediction, and computes MSE loss against the **clean endpoint** of the trajectory.
|
||||
|
||||
Step 2's partition comes from `sample_flex_level0` in `videox_fun/utils/flex_chunking.py` — the **same four-arm mixture Stage 2 takes its level 0 from** (shown for 21 frames with `--num_frame_per_block=3`):
|
||||
|
||||
| Arm | Partition drawn | Stage 2 share |
|
||||
| --- | --- | --- |
|
||||
| coarse | `[21]` — one chunk over the whole clip, fully bidirectional | 30% |
|
||||
| uniform | The uniform grid, itself three equiprobable sub-layouts: `[3]*7` / `[1]*21` / `[3]*7` planning followed by `[1]*21` | 30% |
|
||||
| full_then_blocks | `[21]` (the ladder's level 0) | 30% |
|
||||
| random | A fresh random partition, chunk sizes `--flex_chunk_min`..`--flex_chunk_max` | 10% |
|
||||
|
||||
**Why the draw is shared**: this stage's checkpoint is Stage 2's `--ode_transformer_path` initialisation, so every layout Stage 2 spends iterations on has to be one Stage 1 actually regressed.
|
||||
|
||||
Stage 1 keeps only the level-0 partition, so the two ladder-building arms collapse onto `[21]`: at `--flex_pyramid_levels>1` that is `[21]` 60% / `[3]*7` 20% / `[1]*21` 10% / random 10%; at `=1` there is no coarse arm and the other two take 45% each, i.e. `[21]` 45% / `[3]*7` 30% / `[1]*21` 15% / random 10%. `log_validation` pins `[3]*7`, which the uniform arm covers at 20% of iterations instead of the `1/183708` a purely random draw would give it.
|
||||
|
||||
### 4.1 Quick Start
|
||||
|
||||
The ready-to-use launcher is [train_ode.sh](./train_ode.sh):
|
||||
@@ -257,6 +270,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
|
||||
--flex_forcing \
|
||||
--flex_chunk_min=2 \
|
||||
--flex_chunk_max=10 \
|
||||
--flex_pyramid_levels=4 \
|
||||
--num_frame_per_block=3
|
||||
```
|
||||
|
||||
@@ -268,7 +282,9 @@ bash scripts/wan2.1_flex_forcing/train_ode.sh
|
||||
|
||||
> 💡 Because the ODE trajectory and prompt embeddings are pre-computed in Step 1, **no VAE / text encoder is invoked during ODE training** — training is fast, memory-efficient, and `train_data_dir` can be left empty when `outputs.json` already contains absolute paths.
|
||||
|
||||
> Stage 1 carries `--flex_forcing` and `--flex_chunk_min/max`, but **no** `--flex_pyramid_levels`: it regresses `x0` in a single forward per block, so there is no multi-step schedule to ladder over. The pyramid timestep schedule arrives in Stage 2.
|
||||
> Stage 1 carries `--flex_forcing`, `--flex_chunk_min/max` **and** `--flex_pyramid_levels`, but builds no ladder: it regresses `x0` in a single forward per chunk, so only level 0 of the shared draw is used. `--flex_pyramid_levels` must still **match Stage 2's**, because it decides whether the coarse (whole-clip) arm exists and therefore what the other three arms weigh — leaving it at the default `1` while Stage 2 runs `4` would cut the `[F]` layout's share of Stage 1 from 60% to 45%, under-training exactly the layout Stage 2 starts most of its iterations from. The pyramid *timestep schedule* itself arrives in Stage 2.
|
||||
>
|
||||
> Stage 1 also does **not** pass `--use_kv_cache_training`, so it runs the one-shot block-mask path. That forward supplies neither `kv_cache` nor `clean_x`, which is what arms the K-Projection (§3.3), so `flex_kproj.*` receives **zero gradient** for the whole of Stage 1 and reaches Stage 2 at its identity initialisation. That is the documented cold start rather than a defect — but it does mean Stage 1's steps buy Π nothing. Add the flag if you want Π trained here too; Stage 1's KV branch is already partition-aware and passes one timestep per block, which is what `_first_frame_timestep` assumes.
|
||||
|
||||
Output: `output_dir_wan2.1_flex_forcing_ode_regression/`.
|
||||
|
||||
@@ -304,9 +320,9 @@ Output: `output_dir_wan2.1_flex_forcing_ode_regression/`.
|
||||
| Parameter | Description | Example Value |
|
||||
|-----------|-------------|---------------|
|
||||
| `--train_sampling_steps` | Size of the scheduler grid that `denoising_step_indices_list` indexes into | 1000 |
|
||||
| `--denoising_step_indices_list` | Discrete timestep indices used during ODE regression (corresponds to the 5 sparse points sampled in Step 1). Read off the scheduler's own grid rather than a `set_timesteps()` rebuild, so at `shift=8.0` they map to t = `[1000, 960, 888.89, 727.27]` - bit-identical to what `stochastic_sampling_timesteps` feeds validation and inference. The rebuild was off by up to `5.44`, because it ends its grid at a different `sigma_min`. | `1000 750 500 250` |
|
||||
| `--denoising_step_indices_list` | Discrete timestep indices used during ODE regression (corresponds to the 5 sparse points sampled in Step 1). Read off the scheduler's own grid rather than a `set_timesteps()` rebuild, so at `shift=8.0` they map to t = `[1000, 960, 888.89, 727.27]` - bit-identical to what `stochastic_sampling_timesteps` feeds validation and inference. The rebuild ends its grid at a different `sigma_min`, so indexing it would land away from the t the pipeline actually denoises at. | `1000 750 500 250` |
|
||||
| `--shift` | Shift for `FlowMatchEulerDiscreteScheduler` — **must match `--shift` used in Step 1** | 8.0 |
|
||||
| `--num_frame_per_block` | The uniform fallback for the inherited path; the partition `UNIFORM_BLOCK_PROB = 0.1` of iterations pin, so that layout is genuinely trained rather than merely reachable; and the *fixed* partition `log_validation` samples with so checkpoints stay comparable | 3 |
|
||||
| `--num_frame_per_block` | The uniform fallback for the inherited path; the width the uniform arm of the shared level-0 mixture pins (`[3]*7`, which the uniform arm gives 20% of iterations at `--flex_pyramid_levels>1`, so that layout is genuinely trained rather than merely reachable); the `{1, nfpb}` set the `full_then_blocks` arm draws its block width from; and the *fixed* partition `log_validation` samples with so checkpoints stay comparable | 3 |
|
||||
| `--independent_first_frame` | First frame is independent (`[1, N, N, ...]` block pattern) | - |
|
||||
| `--context_noise` | Context noise level (matches downstream distillation config) | 0 |
|
||||
|
||||
@@ -315,8 +331,9 @@ Output: `output_dir_wan2.1_flex_forcing_ode_regression/`.
|
||||
| Parameter | Paper | Description | Default |
|
||||
|-----------|-------|-------------|---------|
|
||||
| `--flex_forcing` | §3.1 | Instantiate `WanTransformer3DModel_FlexForcing` and validate with `WanFlexForcingPipeline`. Off = the inherited path, unchanged. | off |
|
||||
| `--flex_chunk_min` | §3.1 | Smallest chunk drawn per iteration. `1` is legal and lets single-frame (strictly causal) chunks appear mid-clip; `2` is the paper's and keeps the sampler's tail from degenerating into a sliver. | 2 |
|
||||
| `--flex_chunk_max` | §3.1 | Largest chunk drawn per iteration. Set equal to `--flex_chunk_min` to pin one fixed layout. Unlike stage 2 there is no coarse mixture here, so raising this to the latent frame count is the **only** way stage 1 ever regresses the whole clip as one fully bidirectional chunk (about 5% of draws at 21 latent frames); below the frame count that layout is never drawn. | 10 |
|
||||
| `--flex_chunk_min` | §5 | Smallest chunk drawn by the *random* arm (10% of iterations). `1` is legal and lets single-frame (strictly causal) chunks appear mid-clip; `2` is the paper's and keeps the sampler's tail from degenerating into a sliver. Unlike Stage 2, nothing here refines a coarse chunk down to 1-frame leaves, so `1` really does add coverage — though the fully causal `[1]*F` layout is a separate arm and is drawn either way. | 2 |
|
||||
| `--flex_chunk_max` | §5 | Largest chunk drawn by the *random* arm. Set equal to `--flex_chunk_min` to pin that arm to one fixed layout. The whole-clip layout is its own arm, so this knob does not decide whether Stage 1 ever regresses the clip as one fully bidirectional chunk; raising it to the latent frame count only spends random-arm draws on a layout the coarse arm already covers — at 21 latent frames it makes `[21]` 1/19 of that arm and cuts `P([3]*7)` from `1/183708` to `1/1106560`. | 10 |
|
||||
| `--flex_pyramid_levels` | §3.2 | **Must match Stage 2's.** Stage 1 builds no ladder, but the level-0 mixture is shared and this flag decides whether its coarse arm exists: `>1` gives 30% whole-clip / 30% uniform / 30% `full_then_blocks` / 10% random (the random arm holds a fixed `FLEX_RANDOM_PROB` and the other three split the rest equally), while `=1` divides the coarse arm's share between the two survivors - random stays at 10% - for 45 / 45 / 10. | 1 |
|
||||
|
||||
**Validation Parameters (Optional)**:
|
||||
|
||||
@@ -368,6 +385,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
|
||||
--flex_forcing \
|
||||
--flex_chunk_min=2 \
|
||||
--flex_chunk_max=10 \
|
||||
--flex_pyramid_levels=4 \
|
||||
--num_frame_per_block=3
|
||||
```
|
||||
|
||||
@@ -419,6 +437,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
|
||||
--flex_forcing \
|
||||
--flex_chunk_min=2 \
|
||||
--flex_chunk_max=10 \
|
||||
--flex_pyramid_levels=4 \
|
||||
--num_frame_per_block=3
|
||||
```
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
> ODE 回归是 Flex-Forcing 两阶段流程的 **预训练阶段**，整体流程分两步：
|
||||
>
|
||||
> 1. **第一步 — 生成 ODE 对**（`generate_ode_pairs.py`）：使用 **双向教师模型** Wan2.1-T2V-1.3B，对一组文本提示词执行完整的多步 CFG 去噪，将 ODE 轨迹上的中间 latent 与编码后的 prompt embedding 一起保存为 `.safetensors` 文件。
|
||||
> 2. **第二步 — ODE 回归训练**（`train_ode.py`）：加载第一步生成的 ODE 对，训练一个 **灵活分块生成器**，在轨迹上随机抽样多个噪声等级，预测干净的终点 `x0`；每次迭代都会抽取一个帧划分（九成是全新一轮的随机划分，一成钉成 `--num_frame_per_block` 的均匀块）。训练得到的权重作为 **阶段 2 DMD 蒸馏**（`train_distill.py`，参见 [README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md](./README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md)）的少步初始化。
|
||||
> 2. **第二步 — ODE 回归训练**（`train_ode.py`）：加载第一步生成的 ODE 对，训练一个 **灵活分块生成器**，在轨迹上随机抽样多个噪声等级，预测干净的终点 `x0`；每次迭代都会抽取一个帧划分，用的是**阶段 2 取 level 0 时的同一套四档混合**，因此阶段 2 要蒸馏的布局没有一种是本阶段没回归过的。训练得到的权重作为 **阶段 2 DMD 蒸馏**（`train_distill.py`，参见 [README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md](./README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md)）的少步初始化。
|
||||
>
|
||||
> Flex-Forcing 本身把 Self-Forcing 的标量 `num_frame_per_block` 换成**帧轴上的一个划分** `a = (a_0, ..., a_K)`：chunk（分块）内部是双向注意力，chunk 之间是自回归。`[1] * F` 退化为纯自回归，`[F]` 退化为纯双向注意力，中间的一切形态都由同一个模型覆盖。随机分块正是在本阶段首次引入的。
|
||||
|
||||
@@ -208,9 +208,22 @@ accelerate launch --multi_gpu --num_processes=8 --mixed_precision="bf16" \
|
||||
|
||||
每个训练样本上，训练脚本会：
|
||||
1. 从一个 `.safetensors` 文件中加载稀疏的 5 点轨迹与 prompt embedding；
|
||||
2. 抽取帧划分，并给**每个 chunk 配一个统一的时间步**（`get_timestep_for_ode_flexible`）——这正是"chunk 内部双向"的来源；其中 90% 的迭代抽取一个**全新的随机划分**（chunk 大小在 `--flex_chunk_min`..`--flex_chunk_max` 之间），另外 `UNIFORM_BLOCK_PROB = 0.1` 钉成 `--num_frame_per_block` 的均匀划分——21 个 latent 帧下随机抽取碰到 `[3]*7` 的概率只有 `1/183708`，而 `log_validation` 渲染的正是它；
|
||||
2. 抽取帧划分，并给**每个 chunk 配一个统一的时间步**（`get_timestep_for_ode_flexible`）——这正是"chunk 内部双向"的来源；
|
||||
3. 将带噪 latent 与逐帧时间步送入生成器，把输出的 flow 转换为 `x0` 预测，与轨迹的**干净终点**计算 MSE 损失。
|
||||
|
||||
第 2 步的划分来自 `videox_fun/utils/flex_chunking.py` 的 `sample_flex_level0`，与阶段 2 取 level 0 时用的是**同一套四档混合**（以 21 帧、`--num_frame_per_block=3` 为例）：
|
||||
|
||||
| 档 | 抽出的划分 | 阶段 2 占比 |
|
||||
| --- | --- | --- |
|
||||
| coarse | `[21]` —— 整段一个 chunk，完全双向 | 30% |
|
||||
| uniform | 均匀网格，内部再三等分：`[3]*7` / `[1]*21` / `[3]*7` 规划后接 `[1]*21` | 30% |
|
||||
| full_then_blocks | `[21]`（阶梯的 level 0） | 30% |
|
||||
| random | 全新随机划分，chunk 大小在 `--flex_chunk_min`..`--flex_chunk_max` 之间 | 10% |
|
||||
|
||||
**为什么共用这次抽取**：本阶段的 checkpoint 就是阶段 2 的 `--ode_transformer_path` 初始化，所以阶段 2 会花迭代去跑的每一种布局，都必须是阶段 1 真正回归过的。
|
||||
|
||||
阶段 1 只取 level-0 那一层，两个搭阶梯的档因此都塌到 `[21]` 上：`--flex_pyramid_levels>1` 时 `[21]` 60% / `[3]*7` 20% / `[1]*21` 10% / 随机 10%；`=1` 时没有 coarse 档，其余两档各 45%，即 `[21]` 45% / `[3]*7` 30% / `[1]*21` 15% / 随机 10%。`log_validation` 固定用 `[3]*7`，均匀档给了它 20% 的迭代，而纯随机抽取只有 `1/183708`。
|
||||
|
||||
### 4.1 快速开始
|
||||
|
||||
直接复用启动脚本 [train_ode.sh](./train_ode.sh)：
|
||||
@@ -257,6 +270,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
|
||||
--flex_forcing \
|
||||
--flex_chunk_min=2 \
|
||||
--flex_chunk_max=10 \
|
||||
--flex_pyramid_levels=4 \
|
||||
--num_frame_per_block=3
|
||||
```
|
||||
|
||||
@@ -268,7 +282,9 @@ bash scripts/wan2.1_flex_forcing/train_ode.sh
|
||||
|
||||
> 💡 因为 ODE 轨迹与 prompt embedding 已在第一步预先计算完毕，**ODE 训练阶段不会再调用 VAE / 文本编码器**，训练速度快、显存占用低。当 `outputs.json` 中已经使用绝对路径时，`train_data_dir` 可以留空。
|
||||
|
||||
> 阶段 1 带 `--flex_forcing` 与 `--flex_chunk_min/max`，但**不带** `--flex_pyramid_levels`：金字塔是**多步**调度，而 ODE 阶段对每个 block 只做一次前向来回归 `x0`，没有可以逐级细化的调度。金字塔时间步在阶段 2 才出现。
|
||||
> 阶段 1 带 `--flex_forcing`、`--flex_chunk_min/max` **和** `--flex_pyramid_levels`，但**不搭阶梯**：每个 chunk 只做一次前向来回归 `x0`，所以共用抽取里只有 level 0 被用到。`--flex_pyramid_levels` 仍必须**与阶段 2 一致**，因为它决定 coarse（整段）档是否存在、也就决定了其余三档各占多少 —— 阶段 2 跑 `4` 而这里留默认 `1`，会把 `[F]` 布局在阶段 1 的占比从 60% 砍到 45%，而 `[F]` 恰恰是阶段 2 大多数迭代起步的布局，于是被训得不足。阶梯的*时间步排布*本身要到阶段 2 才出现。
|
||||
>
|
||||
> 阶段 1 也**不传** `--use_kv_cache_training`，因此走的是 one-shot block-mask 路径。该次前向既不给 `kv_cache` 也不给 `clean_x`，而这两者正是启用 K-Projection（§3.3）的条件，所以 `flex_kproj.*` 在整个阶段 1 **梯度恒为 0**，交到阶段 2 时仍停在恒等初始化。这是文档写明的冷启动而非缺陷 —— 但确实意味着阶段 1 这些步数对 Π 毫无贡献。若希望 Π 也在本阶段被训起来，加上该 flag 即可：阶段 1 的 KV 分支已经是划分感知的，且按 block 传单一时间步，正好符合 `_first_frame_timestep` 的假设。
|
||||
|
||||
输出目录：`output_dir_wan2.1_flex_forcing_ode_regression/`。
|
||||
|
||||
@@ -304,9 +320,9 @@ bash scripts/wan2.1_flex_forcing/train_ode.sh
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|------|------|-------|
|
||||
| `--train_sampling_steps` | `denoising_step_indices_list` 所索引的调度器网格大小 | 1000 |
|
||||
| `--denoising_step_indices_list` | ODE 回归使用的离散时间步索引（与第一步抽样的 5 个稀疏轨迹点对应）。直接读调度器自带的网格，而不是 `set_timesteps()` 重建的那张，所以 `shift=8.0` 下映射为 t = `[1000, 960, 888.89, 727.27]` —— 与 `stochastic_sampling_timesteps` 喂给校验和推理的值逐位一致。重建的那张最多偏 `5.44`，因为它的网格末端 `sigma_min` 不同。 | `1000 750 500 250` |
|
||||
| `--denoising_step_indices_list` | ODE 回归使用的离散时间步索引（与第一步抽样的 5 个稀疏轨迹点对应）。直接读调度器自带的网格，而不是 `set_timesteps()` 重建的那张，所以 `shift=8.0` 下映射为 t = `[1000, 960, 888.89, 727.27]` —— 与 `stochastic_sampling_timesteps` 喂给校验和推理的值逐位一致。重建的那张网格末端 `sigma_min` 不同，索引它会落在 pipeline 实际去噪的 t 之外。 | `1000 750 500 250` |
|
||||
| `--shift` | `FlowMatchEulerDiscreteScheduler` 的 shift —— **必须与第一步生成时使用的 `--shift` 一致** | 8.0 |
|
||||
| `--num_frame_per_block` | 继承路径的均匀 fallback；`UNIFORM_BLOCK_PROB = 0.1` 的迭代会把划分钉成它，使该布局是真正被训练而不只是理论上可达；也是 `log_validation` 采样验证时用的固定划分，保证 checkpoint 之间可比较 | 3 |
|
||||
| `--num_frame_per_block` | 继承路径的均匀 fallback；共用 level-0 混合里均匀档所钉的块宽（`[3]*7`，`--flex_pyramid_levels>1` 时均匀档给它 20% 的迭代，使该布局是真正被训练而不只是理论上可达）；`full_then_blocks` 档抽块宽时的 `{1, nfpb}` 取值集合；也是 `log_validation` 采样验证时用的固定划分，保证 checkpoint 之间可比较 | 3 |
|
||||
| `--independent_first_frame` | 第一帧是否独立（`[1, N, N, ...]` 块模式） | - |
|
||||
| `--context_noise` | 上下文噪声等级（与下游蒸馏配置匹配） | 0 |
|
||||
|
||||
@@ -315,8 +331,9 @@ bash scripts/wan2.1_flex_forcing/train_ode.sh
|
||||
| 参数 | 论文 | 说明 | 默认值 |
|
||||
|-----------|-------|-------------|---------|
|
||||
| `--flex_forcing` | §3.1 | 实例化 `WanTransformer3DModel_FlexForcing` 并用 `WanFlexForcingPipeline` 校验。关闭 = 走继承的原路径，行为不变。 | 关 |
|
||||
| `--flex_chunk_min` | §3.1 | 每次迭代抽取的最小 chunk 大小。`1` 合法，会让单帧（严格因果）chunk 出现在片段中间；`2` 是论文取值，也能避免采样器把尾部退化成 1 帧的碎块。 | 2 |
|
||||
| `--flex_chunk_max` | §3.1 | 每次迭代抽取的最大 chunk 大小。设为与 `--flex_chunk_min` 相等即固定一种布局。与阶段 2 不同，这里没有粗粒度混合机制，所以把这个值抬到 latent 帧数是阶段 1 **唯一**能把“整段作为一个完全双向 chunk”回归出来的途径（21 个 latent 帧时约占 5% 的抽取）；低于帧数则该布局永远不会被抽到。 | 10 |
|
||||
| `--flex_chunk_min` | §5 | **随机档**（占 10% 迭代）抽取的最小 chunk 大小。`1` 合法，会让单帧（严格因果）chunk 出现在片段中间；`2` 是论文取值，也能避免采样器把尾部退化成 1 帧的碎块。与阶段 2 不同，这里没有任何东西会把粗 chunk 细化成 1 帧叶子，所以 `1` 确实能增加覆盖——不过完全因果的 `[1]*F` 布局本身是独立一档，无论怎么设都会被抽到。 | 2 |
|
||||
| `--flex_chunk_max` | §5 | **随机档**抽取的最大 chunk 大小。设为与 `--flex_chunk_min` 相等即把该档固定成一种布局。整段布局本身就是独立的一档，所以这个旋钮并不决定阶段 1 能否把“整段作为一个完全双向 chunk”回归出来；把它抬到 latent 帧数只是让随机档把抽取花在 coarse 档已经覆盖的布局上——21 个 latent 帧时它会让 `[21]` 占该档的 1/19，并把 `P([3]*7)` 从 `1/183708` 摊薄到 `1/1106560`。 | 10 |
|
||||
| `--flex_pyramid_levels` | §3.2 | **必须与阶段 2 一致。**阶段 1 不搭阶梯，但 level-0 混合是共用的，而这个 flag 决定 coarse 档是否存在：`>1` 时为 30% 整段 / 30% 均匀 / 30% `full_then_blocks` / 10% 随机（随机档占固定的 `FLEX_RANDOM_PROB`，其余三档平分剩下的 90%）；`=1` 时 coarse 档的那份由两个幸存档平分——随机档恒为 10%——变成 45 / 45 / 10。 | 1 |
|
||||
|
||||
**验证参数（可选）**：
|
||||
|
||||
@@ -368,6 +385,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
|
||||
--flex_forcing \
|
||||
--flex_chunk_min=2 \
|
||||
--flex_chunk_max=10 \
|
||||
--flex_pyramid_levels=4 \
|
||||
--num_frame_per_block=3
|
||||
```
|
||||
|
||||
@@ -419,6 +437,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
|
||||
--flex_forcing \
|
||||
--flex_chunk_min=2 \
|
||||
--flex_chunk_max=10 \
|
||||
--flex_pyramid_levels=4 \
|
||||
--num_frame_per_block=3
|
||||
```
|
||||
|
||||
|
||||
@@ -77,16 +77,12 @@ from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
|
||||
WanTransformer3DModel,
|
||||
WanTransformer3DModel_FlexForcing,
|
||||
WanTransformer3DModel_SelfForcing)
|
||||
from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
|
||||
WanFlexForcingPipeline,
|
||||
WanSelfForcingPipeline)
|
||||
from videox_fun.pipeline import (WanFlexForcingPipeline, WanI2VPipeline,
|
||||
WanPipeline, WanSelfForcingPipeline)
|
||||
from videox_fun.utils.discrete_sampler import DiscreteSampling
|
||||
from videox_fun.utils.flex_chunking import (broadcast_chunk_sizes,
|
||||
build_full_then_blocks_partitions,
|
||||
build_pyramid_partitions,
|
||||
from videox_fun.utils.flex_chunking import (build_pyramid_partitions,
|
||||
chunk_boundaries,
|
||||
sample_flexible_chunks,
|
||||
uniform_chunks,
|
||||
sample_flex_level0,
|
||||
validate_nested_partitions)
|
||||
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
|
||||
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
|
||||
@@ -299,6 +295,7 @@ def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer
|
||||
# the layout inference builds - otherwise the samples say
|
||||
# nothing about the model being trained.
|
||||
flex_kwargs = {}
|
||||
val_stochastic_sampling = not args.flow_euler_rollout
|
||||
if args.flex_forcing:
|
||||
if args.flex_pyramid_levels > 1:
|
||||
# 3.2: exactly what predict_t2v.py passes. `chunk_spec=None`
|
||||
@@ -311,14 +308,6 @@ def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer
|
||||
min_num_frame_per_block=args.flex_min_num_frame_per_block,
|
||||
)
|
||||
else:
|
||||
# No ladder in training, and with no pyramid the
|
||||
# whole-clip band is empty - `[21]` is never drawn at
|
||||
# `flex_chunk_max` below the frame count, so "bidir"
|
||||
# would validate out of distribution too. The uniform
|
||||
# block layout is the one band that is both fixed
|
||||
# across checkpoints and actually trained:
|
||||
# FLEX_ARM_PROB of iterations use it, against
|
||||
# ~1.3% for the most common random partition.
|
||||
flex_kwargs = dict(
|
||||
chunk_spec=args.num_frame_per_block,
|
||||
denoise_mode="fixed",
|
||||
@@ -336,6 +325,7 @@ def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer
|
||||
num_frame_per_block = args.num_frame_per_block,
|
||||
independent_first_frame = args.independent_first_frame,
|
||||
context_noise = args.context_noise,
|
||||
stochastic_sampling = val_stochastic_sampling,
|
||||
**flex_kwargs,
|
||||
).videos
|
||||
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
|
||||
@@ -379,21 +369,20 @@ def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=Non
|
||||
return torch.clip(t.to(torch.int32), low, high - 1)
|
||||
|
||||
|
||||
# Each training iteration draws its level-0 layout from four equally likely arms
|
||||
# (FLEX_ARM_PROB = 25% of iterations each): the whole-clip planning chunk (3.2's
|
||||
# coarse step, `[[F], ...]` at inference), the launcher's own uniform
|
||||
# `num_frame_per_block`, the fixed "first step full, then block-major" ladder
|
||||
# (`denoise_mode="full_then_blocks"`), and a random 3.1 partition. Equal shares
|
||||
# keep the coarse, causal and every in-between layout covered while giving the
|
||||
# full_then_blocks ladder enough iterations to train the block-major-at-high-noise
|
||||
# trajectory the binary pyramid never reaches. Deliberately not a CLI flag - it is
|
||||
# a property of the schedule mixture, not a knob the launcher keeps in sync.
|
||||
FLEX_ARM_PROB = 0.25
|
||||
# The level-0 four-arm mixture (30% of iterations each for the whole-clip
|
||||
# planning chunk, the launcher's own uniform `num_frame_per_block` and the fixed
|
||||
# "first step full, then block-major" ladder, plus 10% for a random §3.3/§5
|
||||
# partition - see `FLEX_RANDOM_PROB`) now lives in
|
||||
# `videox_fun.utils.flex_chunking.sample_flex_level0`, because `train_ode.py`
|
||||
# draws from the SAME law: its checkpoint is this stage's `--ode_transformer_path`
|
||||
# initialisation, so a layout this stage spends most of its iterations on must
|
||||
# not be one stage 1 never regressed. What stays here is only the ladder built on
|
||||
# top of the draw.
|
||||
|
||||
# How many steps of a launch report the partition they drew. Counted from this
|
||||
# process rather than from `global_step`, so a resumed run still gets its own
|
||||
# look at the layout. Two lines come out per step (the generator and the critic
|
||||
# each draw one), which is enough to watch the three arms show up in their
|
||||
# each draw one), which is enough to watch the arms show up in their
|
||||
# stated shares; after that they are only noise in a multi-day run.
|
||||
FLEX_LAYOUT_LOG_STEPS = 20
|
||||
|
||||
@@ -402,134 +391,74 @@ def sample_flex_partitions(args, num_frames, num_denoising_steps, torch_rng,
|
||||
device, verbose=False):
|
||||
"""Draw the Flex-Forcing partition ladder for one training iteration.
|
||||
|
||||
arXiv 2607.03509 3.1 asks for a *random* partition per rollout (chunk sizes
|
||||
2..10) so a single model covers the whole causal-to-bidirectional spectrum,
|
||||
and 3.2 then nests that partition into a coarse-to-fine ladder with one level
|
||||
per denoising step. Returns ``None`` when Flex-Forcing is off, which leaves
|
||||
the inherited uniform ``num_frame_per_block`` masks untouched.
|
||||
arXiv 2607.03509 §3.3 asks for a *random* partition per rollout so a single
|
||||
model covers the whole causal-to-bidirectional spectrum, §5 fixes that draw's
|
||||
chunk sizes at 2..10, and §3.2 then nests the partition into a coarse-to-fine
|
||||
ladder with one level per denoising step. Returns ``None`` when Flex-Forcing
|
||||
is off, which leaves the inherited uniform ``num_frame_per_block`` masks
|
||||
untouched.
|
||||
|
||||
The level-0 layout is a four-way mixture decided by a single uniform draw,
|
||||
each arm taking ``FLEX_ARM_PROB`` (25%) of iterations: one chunk over the
|
||||
whole clip (inference's coarse planning step), a single-level uniform block
|
||||
grid, the fixed "first step full, then block-major" ladder
|
||||
(`denoise_mode="full_then_blocks"`), and a random 2..10 partition. The coarse
|
||||
and random arms are refined into the same binary pyramid, so a single set of
|
||||
weights covers `denoise_mode="pyramid"` whether or not the caller pins
|
||||
`chunk_spec`. The uniform and full_then_blocks arms bypass that refinement:
|
||||
uniform emits a pure-Self-Forcing ladder in one of three layouts drawn from
|
||||
The level-0 layout is a four-way mixture decided by a single uniform draw:
|
||||
one chunk over the whole clip (inference's coarse planning step), a
|
||||
single-level uniform block grid, the fixed "first step full, then
|
||||
block-major" ladder (`denoise_mode="full_then_blocks"`), and a random 2..10
|
||||
partition. The random arm takes ``FLEX_RANDOM_PROB`` (10%) and the other
|
||||
three split what is left equally, i.e. 30% each. The coarse and random arms
|
||||
are refined into the same binary pyramid, so a single set of weights covers
|
||||
`denoise_mode="pyramid"` whether or not the caller pins `chunk_spec`. The
|
||||
uniform and full_then_blocks arms bypass that refinement: uniform emits a
|
||||
pure-Self-Forcing ladder in one of three layouts drawn from
|
||||
`num_frame_per_block` (`[3x7]`) and fully causal (`[1x21]`) -- `[3x7]` alone,
|
||||
`[1x21]` alone, or `[3x7]` planning followed by `[1x21]` for every remaining
|
||||
step; full_then_blocks emits its two-level ladder whole, drawing the same
|
||||
block width. With no pyramid the coarse arm is empty (a `[F]` level 0 would be
|
||||
a degenerate single-chunk rollout), so the split is 25% uniform / 25%
|
||||
full_then_blocks / 50% random.
|
||||
step, 10% of iterations each; full_then_blocks emits its two-level ladder
|
||||
whole, drawing the same block width. With no pyramid the coarse arm is empty
|
||||
(a `[F]` level 0 would be a degenerate single-chunk rollout), so the two
|
||||
survivors widen and the split becomes 45% uniform / 45% full_then_blocks /
|
||||
10% random.
|
||||
|
||||
The draw itself is :func:`videox_fun.utils.flex_chunking.sample_flex_level0`,
|
||||
shared with ``train_ode.py`` so that stage 1 - whose checkpoint this stage
|
||||
loads via ``--ode_transformer_path`` - regresses the same level-0 law it is
|
||||
then distilled on. What this function adds is only the *ladder*: how the drawn
|
||||
layout nests into one level per denoising step.
|
||||
"""
|
||||
if not args.flex_forcing:
|
||||
return None
|
||||
# The arm selector must be IDENTICAL on every rank, not just its broadcast
|
||||
# count. Only the coarse/random arms funnel through the shared refine tail
|
||||
# below; the uniform and full_then_blocks arms each set the ladder directly
|
||||
# from a different tensor, so a per-rank `u` would leave ranks in different
|
||||
# arms building different ladders
|
||||
# -> different walk forward counts -> FSDP all-gather desync (NCCL hang).
|
||||
# Broadcast rank 0's draw so all ranks take the same arm; the arm's existing
|
||||
# single broadcast then reconciles its within-arm randomness (the random base,
|
||||
# the block width) to rank 0, making every ladder byte-identical.
|
||||
u = torch.rand((), generator=torch_rng, device=device)
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
dist.broadcast(u.reshape(1), src=0)
|
||||
u = u.item()
|
||||
pyramid = args.flex_pyramid_levels > 1
|
||||
# Four equally likely arms; the coarse (whole-clip planning) one only when a
|
||||
# pyramid is on, else its 25% folds into the random arm.
|
||||
coarse_hi = FLEX_ARM_PROB if pyramid else 0.0
|
||||
uniform_hi = coarse_hi + FLEX_ARM_PROB
|
||||
ftb_hi = uniform_hi + FLEX_ARM_PROB
|
||||
ladder = None
|
||||
if u < coarse_hi:
|
||||
# Coarse end of 3.1: reuse the ladder builder's own level-0 rule so the
|
||||
# `independent_first_frame` handling cannot drift from inference's.
|
||||
base = build_pyramid_partitions(
|
||||
num_frames, num_levels=1, base_chunks=None,
|
||||
independent_first_frame=args.independent_first_frame)[0]
|
||||
arm = "whole clip, the 3.2 coarse planning layout"
|
||||
elif u < uniform_hi:
|
||||
# Pure Self-Forcing (no binary-pyramid refinement), three layouts drawn
|
||||
# from `num_frame_per_block` (-> [3x7]) and fully causal (1 frame/block,
|
||||
# -> [1x21]): [3x7] alone, [1x21] alone, and [3x7] planning followed by
|
||||
# [1x21] for every remaining step (the coarse-block -> per-frame descent
|
||||
# the pyramid never samples because it halves through [2,1x14] first).
|
||||
# The three change the ladder DEPTH (hence the walk's forward count), so
|
||||
# the choice must be rank-consistent; it is free here because `u` was
|
||||
# already broadcast to rank 0, so slicing its position within the uniform
|
||||
# band gives every rank the same sub-arm with no extra collective. Like
|
||||
# every arm this issues EXACTLY ONE broadcast (the level-0 layout; the
|
||||
# trailing [1x21] levels are deterministic from `num_frames`). Setting
|
||||
# `ladder` here skips the refine tail. `uniform_chunks`, not
|
||||
# `normalize_chunk_spec`: the latter encodes `independent_first_frame` as
|
||||
# a leading 1 and would silently disagree with the bands around it.
|
||||
band = (u - coarse_hi) / FLEX_ARM_PROB # [0, 1) in the uniform arm
|
||||
sub = min(int(band * 3), 2) # 0 / 1 / 2, equiprobable
|
||||
nfpb = max(1, int(args.num_frame_per_block))
|
||||
iff = args.independent_first_frame
|
||||
causal = uniform_chunks(num_frames, 1, independent_first_frame=iff)
|
||||
if sub == 0:
|
||||
base = broadcast_chunk_sizes(causal, device=device)
|
||||
ladder = [base]
|
||||
arm = "uniform 1-frame blocks, single level (fully causal)"
|
||||
elif sub == 1:
|
||||
base = broadcast_chunk_sizes(
|
||||
uniform_chunks(num_frames, nfpb, independent_first_frame=iff),
|
||||
device=device)
|
||||
ladder = [base]
|
||||
arm = f"uniform {nfpb}-frame blocks, single level"
|
||||
else:
|
||||
base = broadcast_chunk_sizes(
|
||||
uniform_chunks(num_frames, nfpb, independent_first_frame=iff),
|
||||
device=device)
|
||||
ladder = ([base] + [list(causal)
|
||||
for _ in range(max(0, num_denoising_steps - 1))])
|
||||
arm = (f"uniform {nfpb}-frame plan -> fully causal [1x" f"{num_frames}] "
|
||||
"for the remaining steps")
|
||||
elif u < ftb_hi:
|
||||
# "First step full, every later step block-major": a fixed two-level
|
||||
# ladder that does NOT go through the binary pyramid refinement, so the
|
||||
# block-major level is reached at high noise (the 2nd step) instead of as
|
||||
# a deep refinement - the one trajectory the pyramid arms never sample.
|
||||
# The block-major width is drawn between fully causal (1 frame / block)
|
||||
# and the launcher's `num_frame_per_block`, so both the tight-AR and the
|
||||
# coarse-block refinement of the whole-clip plan get trained. The block
|
||||
# width is a per-rank draw, so it has to be reconciled with EXACTLY ONE
|
||||
# broadcast - the same collective count every other arm issues (they each
|
||||
# broadcast their single base) - or the ranks desync and NCCL deadlocks.
|
||||
# Level 0 is just [F], identical on every rank and needing no sync, so we
|
||||
# broadcast only the block-major level; rank 0's draw wins and every rank
|
||||
# rebuilds the same two-level ladder locally.
|
||||
block_choices = sorted({1, max(1, int(args.num_frame_per_block))})
|
||||
pick = int(torch.rand((), generator=torch_rng, device=device).item()
|
||||
* len(block_choices))
|
||||
block = block_choices[min(pick, len(block_choices) - 1)]
|
||||
block = min(block, int(num_frames))
|
||||
ftb = build_full_then_blocks_partitions(
|
||||
num_frames, block,
|
||||
independent_first_frame=args.independent_first_frame)
|
||||
ftb[1] = broadcast_chunk_sizes(ftb[1], device=device)
|
||||
ladder = ftb[:num_denoising_steps]
|
||||
arm = (f"full clip -> uniform {max(ftb[1])}-frame blocks, "
|
||||
"the full_then_blocks ladder")
|
||||
draw = sample_flex_level0(
|
||||
num_frames,
|
||||
num_frame_per_block=args.num_frame_per_block,
|
||||
chunk_min=args.flex_chunk_min,
|
||||
chunk_max=args.flex_chunk_max,
|
||||
pyramid=pyramid,
|
||||
generator=torch_rng,
|
||||
device=device,
|
||||
independent_first_frame=args.independent_first_frame)
|
||||
base = draw["sizes"]
|
||||
arm = draw["arm_text"]
|
||||
if draw["arm"] in ("uniform_causal", "uniform_blocks"):
|
||||
# Pure Self-Forcing grids: a single level, and they deliberately bypass
|
||||
# the binary refinement below - refining `[3x7]` would halve it through
|
||||
# `[2, 1x14]`, which is the pyramid's own trajectory, not this arm's.
|
||||
ladder = [base]
|
||||
elif draw["arm"] == "uniform_plan_then_causal":
|
||||
# `[3x7]` planning, then fully causal for every remaining step: the
|
||||
# coarse-block -> per-frame descent the pyramid never samples. The
|
||||
# trailing `[1x21]` levels are deterministic from `num_frames`, so the
|
||||
# arm's single broadcast (already issued) covered everything random.
|
||||
ladder = ([base] + [list(draw["causal"])
|
||||
for _ in range(max(0, num_denoising_steps - 1))])
|
||||
elif draw["arm"] == "full_then_blocks":
|
||||
# Two levels, taken whole. The walk clamps to the finest level for every
|
||||
# later step, so a ladder shorter than the denoising schedule is the
|
||||
# intended shape here, not wasted steps.
|
||||
ladder = [list(base), list(draw["block_level"])][:num_denoising_steps]
|
||||
else:
|
||||
base = sample_flexible_chunks(
|
||||
num_frames, min_chunk=args.flex_chunk_min, max_chunk=args.flex_chunk_max,
|
||||
generator=torch_rng, device=device,
|
||||
independent_first_frame=args.independent_first_frame)
|
||||
arm = f"random {args.flex_chunk_min}..{args.flex_chunk_max}-frame blocks, the 3.1 spectrum"
|
||||
if ladder is None:
|
||||
# Every rank has to train the same layout: the FlexAttention mask, and the
|
||||
# `coarse` / `random`: the two arms the binary pyramid refines. Every rank
|
||||
# has to train the same layout - the FlexAttention mask, and the
|
||||
# `num_frame_per_block` derived from it, must agree across the SP/FSDP
|
||||
# group. Only the coarse and random arms reach here; the uniform and
|
||||
# full_then_blocks arms above already set `ladder` (each with its single
|
||||
# broadcast), so they skip this refine path.
|
||||
base = broadcast_chunk_sizes(base, device=device)
|
||||
# group - which `sample_flex_level0` already guaranteed by broadcasting
|
||||
# rank 0's draw.
|
||||
ladder = [base]
|
||||
if pyramid:
|
||||
ladder = build_pyramid_partitions(
|
||||
@@ -601,16 +530,24 @@ def flex_self_context_enabled(args, clean_x, final_step_index):
|
||||
"""Whether 3.3's K-Projection gets a self-generated clean context this round.
|
||||
|
||||
This is a block-mask-path concern only, and this flag is inert under
|
||||
``--use_kv_cache_training``. The KV-cache rollout passes no ``flex_state``,
|
||||
so ``_flex_project`` short-circuits and :math:`\\Pi` stays at its identity
|
||||
initialisation there - which matches KV inference, which never supplies a
|
||||
timestep either, so the two sides agree rather than one of them drifting.
|
||||
The block-mask path is the one that exercises :math:`\\Pi`, and it builds no
|
||||
KV cache, so with prompt-only data - ``clean_latents is None``, which the
|
||||
prompt-only branch sets unconditionally - no clean context exists anywhere
|
||||
in the sequence, :math:`\\Pi_{t\\leftarrow 0}` has nothing to project and
|
||||
never enters the autograd graph: its parameters sit in an optimizer group
|
||||
and still stay at the identity initialisation for the entire run.
|
||||
``--use_kv_cache_training`` - but *not* because :math:`\\Pi` is idle there.
|
||||
:class:`WanTransformer3DModel_FlexForcing` arms :math:`\\Pi` by itself
|
||||
whenever a projection was built and the forward has a clean context to read
|
||||
(``kv_cache is not None or clean_x is not None``), taking the noise level
|
||||
from ``t``; a caller that passes no ``flex_state`` still projects. So
|
||||
:math:`\\Pi` **is** invoked on the KV-cache path, in training and in
|
||||
inference alike, and the exit step that carries gradient gives it real
|
||||
gradient. All this flag does *not* do there is supply a ``clean_x`` half:
|
||||
the KV rollout already commits its own clean context into the cache ("feed
|
||||
denoised_pred directly"), so it has no use for one and the switch has no
|
||||
lever to pull.
|
||||
On the block-mask path the switch is what decides whether :math:`\\Pi` has
|
||||
anything to project at all: that path builds no KV cache, so with
|
||||
prompt-only data - ``clean_latents is None``, which the prompt-only branch
|
||||
sets unconditionally - no clean context exists anywhere in the sequence,
|
||||
:math:`\\Pi_{t\\leftarrow 0}` has nothing to project and never enters the
|
||||
autograd graph: its parameters sit in an optimizer group and still stay at
|
||||
the identity initialisation for the entire run.
|
||||
|
||||
The denoising loop already re-rolls the model under ``no_grad`` and converts
|
||||
every non-final step to an x0 prediction, so the previous step's prediction
|
||||
@@ -1137,21 +1074,23 @@ def parse_args():
|
||||
"--flex_chunk_min",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Smallest chunk drawn when sampling a partition (paper 3.1 uses 2). "
|
||||
"1 is allowed but adds no coverage: refining a 2-frame chunk already "
|
||||
"yields 1-frame leaves."
|
||||
help="Smallest chunk drawn by the level-0 mixture's *random* arm (paper "
|
||||
"§5 uses 2). 1 is allowed but adds no coverage: refining a 2-frame "
|
||||
"chunk already yields 1-frame leaves."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flex_chunk_max",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Largest chunk drawn when sampling a partition (paper 3.1 uses 10; "
|
||||
"setting it equal to --flex_chunk_min pins one fixed layout). Keep it "
|
||||
help="Largest chunk drawn by the level-0 mixture's *random* arm (paper §5 "
|
||||
"uses 10; setting it equal to --flex_chunk_min pins that arm to one "
|
||||
"fixed layout). Keep it "
|
||||
"below the latent frame count: the whole-clip layout is already "
|
||||
"covered by the pyramid mixture, so raising max to the frame count "
|
||||
"only spends draws on a layout you already have and thins out 3.1 - "
|
||||
"at 21 latent frames max=21 leaves 3649 distinct partitions versus "
|
||||
"4882 at max=10."
|
||||
"covered by the coarse arm, so raising max to the frame count spends "
|
||||
"draws on a layout you already have and thins every other one - at 21 "
|
||||
"latent frames exhaustive enumeration gives 6441 distinct partitions "
|
||||
"at max=10 and 6765 at max=21, but max=21 makes [21] alone 1/19 of "
|
||||
"draws and cuts P([3]*7) from 1/183708 to 1/1106560."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flex_pyramid_levels",
|
||||
@@ -1178,10 +1117,14 @@ def parse_args():
|
||||
"block-mask path trains on prompts alone: it builds no KV cache, so "
|
||||
"without this there is no clean context at all and the K-Projection "
|
||||
"never enters the autograd graph. Doubles the attention sequence. "
|
||||
"Inert under --use_kv_cache_training: that path passes no flex_state, "
|
||||
"so the K-Projection is never invoked on either side and stays at its "
|
||||
"identity init. Also ignored when real --use_teacher_forcing data is "
|
||||
"present, and with --flex_forcing off."
|
||||
"Inert under --use_kv_cache_training for a different reason than it "
|
||||
"looks like: that path already commits its own clean context into the "
|
||||
"KV cache, so it needs no clean_x half. The K-Projection itself IS "
|
||||
"active there - the Flex backbone arms it automatically whenever a "
|
||||
"projection was built and kv_cache or clean_x is present - so this "
|
||||
"flag is not evidence that the projection sits untrained under "
|
||||
"--use_kv_cache_training. Also ignored when real --use_teacher_forcing "
|
||||
"data is present, and with --flex_forcing off."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--independent_first_frame",
|
||||
|
||||
@@ -63,11 +63,7 @@ from videox_fun.models import (AutoencoderKLWan, WanT5EncoderModel,
|
||||
WanTransformer3DModel_FlexForcing,
|
||||
WanTransformer3DModel_SelfForcing)
|
||||
from videox_fun.pipeline import WanFlexForcingPipeline, WanSelfForcingPipeline
|
||||
from videox_fun.utils.flex_chunking import (UNIFORM_BLOCK_PROB,
|
||||
broadcast_chunk_sizes,
|
||||
chunk_boundaries,
|
||||
sample_flexible_chunks,
|
||||
uniform_chunks)
|
||||
from videox_fun.utils.flex_chunking import chunk_boundaries, sample_flex_level0
|
||||
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
|
||||
from videox_fun.utils.utils import save_videos_grid
|
||||
|
||||
@@ -234,12 +230,27 @@ def get_timestep_for_ode(
|
||||
|
||||
|
||||
def get_timestep_for_ode_flexible(min_timestep, max_timestep, batch_size,
|
||||
chunk_sizes, device, generator=None):
|
||||
chunk_sizes, device, generator=None,
|
||||
clean_index=None):
|
||||
"""`get_timestep_for_ode` for a Flex-Forcing partition (arXiv 2607.03509 3.1).
|
||||
|
||||
Same contract - every frame of a chunk shares one timestep index - but the
|
||||
chunk widths come from the partition instead of a single scalar, so the
|
||||
reshape trick of the uniform version does not apply.
|
||||
|
||||
``clean_index`` is the index whose entry in ``train_denoising_step_list`` is
|
||||
t=0, or ``None`` when the schedule has no clean entry. When given, each
|
||||
sample is guaranteed at least one chunk that did *not* draw it. An all-clean
|
||||
draw makes the loss mask ``(timestep != 0)`` empty, and the caller's
|
||||
fallback ``mse(pred_x0, target)`` is then *exactly* zero rather than merely
|
||||
small: those frames gather the clean endpoint itself and sigma is pinned to
|
||||
0, so ``pred_x0 == target`` bitwise. The iteration contributes no gradient
|
||||
at all. The probability is ``len(train_denoising_step_list) ** -K`` in the
|
||||
number of chunks K, so it is dominated by the coarsest layouts: the
|
||||
whole-clip ``[F]`` partition is a single chunk and on its own carries
|
||||
``1 / len(train_denoising_step_list)`` of iterations. A single chunk also
|
||||
has no later block to serve as clean context *for*, which is the only reason
|
||||
the clean index exists, so redrawing it loses nothing.
|
||||
"""
|
||||
num_frames = sum(int(c) for c in chunk_sizes)
|
||||
per_chunk = torch.randint(
|
||||
@@ -248,6 +259,21 @@ def get_timestep_for_ode_flexible(min_timestep, max_timestep, batch_size,
|
||||
device=device, dtype=torch.long,
|
||||
generator=generator,
|
||||
)
|
||||
if clean_index is not None and per_chunk.shape[1] > 0 \
|
||||
and max_timestep - min_timestep > 1:
|
||||
all_clean = (per_chunk == clean_index).all(dim=1)
|
||||
if bool(all_clean.any()):
|
||||
rows = all_clean.nonzero(as_tuple=True)[0]
|
||||
# One chunk per affected sample, redrawn uniformly over every index
|
||||
# *except* the clean one: draw from the shrunken range, then shift
|
||||
# past the hole. Works for any clean_index, not just the last one.
|
||||
pick = torch.randint(
|
||||
0, per_chunk.shape[1], [rows.numel()],
|
||||
device=device, dtype=torch.long, generator=generator)
|
||||
redraw = torch.randint(
|
||||
min_timestep, max_timestep - 1, [rows.numel()],
|
||||
device=device, dtype=torch.long, generator=generator)
|
||||
per_chunk[rows, pick] = redraw + (redraw >= clean_index).long()
|
||||
timestep = torch.empty(
|
||||
batch_size, num_frames, device=device, dtype=per_chunk.dtype)
|
||||
for chunk_idx, (start_frame, end_frame) in enumerate(chunk_boundaries(chunk_sizes)):
|
||||
@@ -256,38 +282,48 @@ def get_timestep_for_ode_flexible(min_timestep, max_timestep, batch_size,
|
||||
|
||||
|
||||
def sample_flex_partition(args, num_frames, torch_rng, device):
|
||||
"""Draw this iteration's Flex-Forcing partition (3.1/3.3).
|
||||
"""Draw this iteration's Flex-Forcing partition (§3.1 / §3.3).
|
||||
|
||||
Returns ``None`` when Flex-Forcing is off, so every caller keeps using the
|
||||
uniform ``num_frame_per_block`` layout. 3.2's pyramid is a *multi-step*
|
||||
schedule and the ODE stage regresses x0 in a single forward per block, so
|
||||
there is nothing to ladder over here.
|
||||
uniform ``num_frame_per_block`` layout.
|
||||
|
||||
``UNIFORM_BLOCK_PROB`` of iterations pin the launcher's own uniform
|
||||
``num_frame_per_block`` layout instead of drawing a random one. Random draws
|
||||
would almost never produce it - the uniform ``[3] * 7`` at 21 latent frames
|
||||
has probability 1/183708, i.e. 0.11 expected hits in 20000 iterations - yet
|
||||
`log_validation` renders exactly that partition, so without this band stage
|
||||
1's samples would come from a layout it never regressed.
|
||||
The draw is :func:`videox_fun.utils.flex_chunking.sample_flex_level0` - the
|
||||
*same* four-arm mixture stage 2 uses for its level 0, random arm at
|
||||
``FLEX_RANDOM_PROB`` (10%) and the three deterministic arms at 30% each.
|
||||
Sharing it is the whole point: this stage's checkpoint is stage 2's
|
||||
``--ode_transformer_path`` initialisation, so every layout stage 2 spends
|
||||
iterations on has to be one stage 1 actually regressed. The two-band sampler
|
||||
this replaced drew 10% uniform / 90% random - coincidentally the mirror image
|
||||
of today's split, and beside the point, since what made it unusable was its
|
||||
chunk support being exactly ``{--flex_chunk_min .. --flex_chunk_max}``: the
|
||||
whole-clip ``[F]`` layout that 60% of stage 2's level-0 iterations start from
|
||||
(at ``--flex_pyramid_levels=4``) had probability *zero* there (``max_chunk``
|
||||
below the frame count makes it unreachable, and at most 1/19 of random draws
|
||||
even with ``max_chunk`` raised to it), and so did every 1-frame chunk -
|
||||
``[1] * F`` is ``1/F!`` under a random draw, i.e. unreachable at any setting
|
||||
of these two knobs.
|
||||
|
||||
§3.2's pyramid is a *multi-step* schedule and this stage regresses x0 in a
|
||||
single forward per chunk, so there is no ladder to build here and only
|
||||
level 0 is used. ``--flex_pyramid_levels`` is still read, and must match
|
||||
stage 2's, because it decides whether the coarse arm exists and therefore
|
||||
how the other arms split the 90% the random one leaves behind.
|
||||
|
||||
``sample_flex_level0`` broadcasts rank 0's draw, so every rank regresses the
|
||||
same layout; otherwise the FlexAttention mask and the per-chunk timesteps
|
||||
would disagree across the group.
|
||||
"""
|
||||
if not args.flex_forcing:
|
||||
return None
|
||||
if torch.rand(
|
||||
(), generator=torch_rng, device=device).item() < UNIFORM_BLOCK_PROB:
|
||||
# `uniform_chunks` rather than `normalize_chunk_spec`: the latter takes no
|
||||
# `independent_first_frame` argument (it encodes that as a leading 1), so
|
||||
# it would silently drop the leading 1 the random branch below adds.
|
||||
sizes = uniform_chunks(
|
||||
num_frames, args.num_frame_per_block,
|
||||
independent_first_frame=args.independent_first_frame)
|
||||
else:
|
||||
sizes = sample_flexible_chunks(
|
||||
num_frames, min_chunk=args.flex_chunk_min, max_chunk=args.flex_chunk_max,
|
||||
generator=torch_rng, device=device,
|
||||
independent_first_frame=args.independent_first_frame)
|
||||
# Every rank must regress the same layout, otherwise the FlexAttention mask
|
||||
# and the per-chunk timesteps disagree across the group.
|
||||
return broadcast_chunk_sizes(sizes, device=device)
|
||||
return sample_flex_level0(
|
||||
num_frames,
|
||||
num_frame_per_block=args.num_frame_per_block,
|
||||
chunk_min=args.flex_chunk_min,
|
||||
chunk_max=args.flex_chunk_max,
|
||||
pyramid=args.flex_pyramid_levels > 1,
|
||||
generator=torch_rng,
|
||||
device=device,
|
||||
independent_first_frame=args.independent_first_frame)["sizes"]
|
||||
|
||||
|
||||
def initialize_kv_cache_for_training(batch_size, num_frames, frame_seq_length,
|
||||
@@ -566,20 +602,39 @@ def parse_args():
|
||||
"--flex_chunk_min",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Smallest chunk drawn when sampling a partition (paper 3.1 uses 2). "
|
||||
"1 is allowed but adds no coverage: refining a 2-frame chunk already "
|
||||
"yields 1-frame leaves."
|
||||
help="Smallest chunk drawn by the *random* arm of the level-0 mixture "
|
||||
"(paper §5 uses 2). Unlike stage 2, nothing here refines a coarse "
|
||||
"chunk down to 1-frame leaves, so 1 is a real setting: it lets the "
|
||||
"random arm emit single-frame chunks itself. The fully causal "
|
||||
"[1] * F layout is a separate arm and is drawn either way."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flex_chunk_max",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Largest chunk drawn when sampling a partition (paper 3.1 uses 10; "
|
||||
"setting it equal to --flex_chunk_min pins one fixed layout). Unlike "
|
||||
"stage 2 there is no coarse mixture here, so raising this to the "
|
||||
"latent frame count is the *only* way stage 1 ever sees the whole "
|
||||
"clip as one fully bidirectional chunk - about 5% of draws at 21 "
|
||||
"latent frames. Below the frame count that layout is never drawn."
|
||||
help="Largest chunk drawn by the *random* arm of the level-0 mixture "
|
||||
"(paper §5 uses 10; setting it equal to --flex_chunk_min pins that "
|
||||
"arm to one fixed layout). The whole-clip layout is a separate arm, "
|
||||
"so this knob does not decide whether stage 1 ever regresses the "
|
||||
"clip as one fully bidirectional chunk. Raising it to the latent "
|
||||
"frame count only spends random-arm draws on a layout the coarse arm "
|
||||
"already covers - at 21 latent frames it makes [21] 1/19 of that arm "
|
||||
"and thins every other one, cutting P([3]*7) from 1/183708 to "
|
||||
"1/1106560."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flex_pyramid_levels",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Must match stage 2's --flex_pyramid_levels. Stage 1 builds no ladder "
|
||||
"(it regresses x0 in one forward per chunk, so only level 0 is used), "
|
||||
"but the level-0 mixture is shared with stage 2 and this flag decides "
|
||||
"whether its coarse arm exists: >1 gives 30%% whole-clip / 30%% "
|
||||
"uniform / 30%% full_then_blocks / 10%% random, while =1 divides the "
|
||||
"coarse arm's share between the two survivors - random stays at 10%% - "
|
||||
"for 45/45/10. Leaving it at 1 while stage 2 "
|
||||
"runs 4 would under-train exactly the [F] layout stage 2 starts most "
|
||||
"of its iterations from."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--independent_first_frame",
|
||||
@@ -1044,14 +1099,13 @@ def main():
|
||||
|
||||
# Compute denoising_step_list from the scheduler's own grid.
|
||||
# Deliberately *not* calling `set_timesteps(...)`: it rebuilds the grid with
|
||||
# a different `sigma_min`, so indexing it lands up to 5.4 away from the t the
|
||||
# pipeline actually denoises at (measured 5.44 at shift=8.0, 3.71 at 5.0).
|
||||
# The `__init__` grid maps index k to exactly t = 1 - k / num_train_timesteps,
|
||||
# which is what `stochastic_sampling_timesteps` hardcodes, so training,
|
||||
# validation and inference read the same numbers - |diff| = 0.0 at both
|
||||
# shifts. `train_distill.py` relies on the same grid. The explicit `.to()`
|
||||
# keeps these on the accelerator, which `set_timesteps(device=...)` used to
|
||||
# do for us and which the per-frame `timestep` below requires.
|
||||
# a different `sigma_min`, so indexing it lands away from the t the pipeline
|
||||
# actually denoises at. The `__init__` grid maps index k to exactly
|
||||
# t = 1 - k / num_train_timesteps, which is what
|
||||
# `stochastic_sampling_timesteps` hardcodes, so training, validation and
|
||||
# inference read the same numbers. `train_distill.py` relies on the same
|
||||
# grid. The explicit `.to()` keeps these on the accelerator, which the
|
||||
# per-frame `timestep` below requires.
|
||||
denoising_step_list = noise_scheduler.timesteps.to(accelerator.device)[
|
||||
args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)
|
||||
]
|
||||
@@ -1064,6 +1118,13 @@ def main():
|
||||
denoising_step_list, torch.tensor([0], device=denoising_step_list.device)
|
||||
])
|
||||
num_denoising_steps = len(train_denoising_step_list)
|
||||
# Index of the clean (t=0) entry, derived from the list itself rather than
|
||||
# assumed to be the last one. Passed to `get_timestep_for_ode_flexible`,
|
||||
# which keeps every iteration from drawing it for *all* chunks - an
|
||||
# all-clean draw zeroes the loss mask and makes the fallback MSE exactly 0,
|
||||
# i.e. a zero-gradient iteration. See that function for the probability.
|
||||
_clean_hits = (train_denoising_step_list == 0).nonzero(as_tuple=True)[0]
|
||||
clean_step_index = int(_clean_hits[-1]) if _clean_hits.numel() > 0 else None
|
||||
if accelerator.is_main_process:
|
||||
print(f"Denoising step list (inference): {denoising_step_list.tolist()}")
|
||||
print(f"Denoising step list (training): {train_denoising_step_list.tolist()}")
|
||||
@@ -1221,6 +1282,7 @@ def main():
|
||||
0, num_denoising_steps, bsz, all_num_frames,
|
||||
accelerator.device,
|
||||
generator=torch_rng,
|
||||
clean_index=clean_step_index,
|
||||
) # [B, F]
|
||||
else:
|
||||
index = get_timestep_for_ode(
|
||||
@@ -1350,6 +1412,7 @@ def main():
|
||||
0, num_denoising_steps, bsz, flex_partition,
|
||||
accelerator.device,
|
||||
generator=torch_rng,
|
||||
clean_index=clean_step_index,
|
||||
) # [B, F]
|
||||
else:
|
||||
index = get_timestep_for_ode(
|
||||
|
||||
@@ -39,4 +39,5 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
|
||||
--flex_forcing \
|
||||
--flex_chunk_min=2 \
|
||||
--flex_chunk_max=10 \
|
||||
--flex_pyramid_levels=4 \
|
||||
--num_frame_per_block=3
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 时，表示视频和图像分辨率的最大值。
|
||||
|
||||
@@ -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):
|
||||
|
||||
|
||||
@@ -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 特有参数**（除非清楚后果，否则需与第一步保持一致）：
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import gc
|
||||
import json
|
||||
import logging
|
||||
@@ -74,10 +75,12 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512,
|
||||
get_random_mask)
|
||||
from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
|
||||
WanTransformer3DModel,
|
||||
WanTransformer3DModel_SelfForcing)
|
||||
WanTransformer3DModel_SelfForcing,
|
||||
add_wan_gan_discriminator)
|
||||
from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
|
||||
WanSelfForcingPipeline)
|
||||
from videox_fun.utils.discrete_sampler import DiscreteSampling
|
||||
from videox_fun.utils.fsdp_ema import FSDPEMA
|
||||
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
|
||||
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
|
||||
get_image_to_video_latent,
|
||||
@@ -165,6 +168,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi
|
||||
return pred_for_score, score_num_frames, True
|
||||
|
||||
|
||||
def repeat_batch(value, repeats):
|
||||
"""Repeat a batched tensor or per-sample list along the batch dimension."""
|
||||
if value is None or repeats == 1:
|
||||
return value
|
||||
if torch.is_tensor(value):
|
||||
return torch.cat([value] * repeats)
|
||||
return list(value) * repeats
|
||||
|
||||
|
||||
def filter_kwargs(cls, kwargs):
|
||||
import inspect
|
||||
sig = inspect.signature(cls.__init__)
|
||||
@@ -492,6 +504,8 @@ def parse_args():
|
||||
parser.add_argument(
|
||||
"--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
|
||||
)
|
||||
parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.")
|
||||
parser.add_argument("--ema_decay", type=float, default=0.9999, help="EMA decay applied per generator update (Causal-Forcing uses 0.99).")
|
||||
parser.add_argument(
|
||||
"--allow_tf32",
|
||||
action="store_true",
|
||||
@@ -523,6 +537,9 @@ def parse_args():
|
||||
parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
|
||||
parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
|
||||
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
|
||||
parser.add_argument(
|
||||
"--fsdp_reduce_fp32", action="store_true", help="Reduce gradients in fp32 under FSDP mixed precision (params are still gathered in low precision)."
|
||||
)
|
||||
parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
|
||||
parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
|
||||
parser.add_argument(
|
||||
@@ -722,6 +739,12 @@ def parse_args():
|
||||
default=None,
|
||||
help=("If you want to load the ode-trained weight into generator transformer3d, input its path."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--real_score_model_name_or_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help=("Path to a separate model (e.g. Wan2.1-T2V-14B) used as real score. Defaults to pretrained_model_name_or_path."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae_path",
|
||||
type=str,
|
||||
@@ -782,6 +805,11 @@ def parse_args():
|
||||
default=4.5,
|
||||
help="The cfg scale for real score.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--score_timestep_clamp",
|
||||
action="store_true",
|
||||
help="Clamp DMD and critic timesteps to [0.02, 0.98] * num_train_timesteps.",
|
||||
)
|
||||
parser.add_argument(
|
||||
'--denoising_step_indices_list',
|
||||
nargs='+',
|
||||
@@ -842,6 +870,67 @@ def parse_args():
|
||||
default=1.0,
|
||||
help="Probability of applying teacher forcing per step (1.0 = always)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gan_g_weight",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="One-Forcing: weight of the non-saturating GAN loss added to the DMD generator loss (0 disables). "
|
||||
"Any positive GAN/R1/R2 weight attaches a discriminator head to fake_score and requires real videos."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gan_d_weight",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="One-Forcing: weight of the GAN discriminator loss added to the fake_score denoising loss (0 disables)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gan_feature_layers",
|
||||
nargs="+",
|
||||
type=int,
|
||||
default=[21, 29],
|
||||
help="fake_score block indices whose outputs feed the discriminator head (one register token per layer)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gan_block_ffn_dim",
|
||||
type=int,
|
||||
default=2048,
|
||||
help="FFN dim of the discriminator cross-attention blocks."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gan_block_num_heads",
|
||||
type=int,
|
||||
default=12,
|
||||
help="Number of heads of the discriminator cross-attention blocks."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--relativistic_discriminator",
|
||||
action="store_true",
|
||||
help="Use the relativistic GAN loss (fake vs. real logit difference)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--r1_weight",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Finite-difference R1 penalty on noisy real latents (APT form, not divided by sigma)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--r1_sigma",
|
||||
type=float,
|
||||
default=0.01,
|
||||
help="Perturbation std of the R1 penalty."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--r2_weight",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Finite-difference R2 penalty on noisy generated latents."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--r2_sigma",
|
||||
type=float,
|
||||
default=0.01,
|
||||
help="Perturbation std of the R2 penalty."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
|
||||
@@ -914,6 +1003,9 @@ def main():
|
||||
else:
|
||||
fsdp_stage = 0
|
||||
print(f"Using FSDP stage: {fsdp_stage}")
|
||||
if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
|
||||
fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
|
||||
print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
|
||||
|
||||
args.use_fsdp = True
|
||||
if fsdp_stage == 3:
|
||||
@@ -970,6 +1062,9 @@ def main():
|
||||
args.mixed_precision = accelerator.mixed_precision
|
||||
|
||||
args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list]
|
||||
use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight))
|
||||
if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1):
|
||||
raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.")
|
||||
# Load scheduler, tokenizer and models.
|
||||
noise_scheduler = FlowMatchEulerDiscreteScheduler(
|
||||
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
|
||||
@@ -1030,7 +1125,7 @@ def main():
|
||||
low_cpu_mem_usage=True,
|
||||
).to(weight_dtype)
|
||||
real_score_transformer3d = WanTransformer3DModel.from_pretrained(
|
||||
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
|
||||
os.path.join(args.real_score_model_name_or_path or args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
|
||||
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
||||
low_cpu_mem_usage=True,
|
||||
).to(weight_dtype)
|
||||
@@ -1062,7 +1157,8 @@ def main():
|
||||
state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()}
|
||||
|
||||
m, u = generator_transformer3d.load_state_dict(state_dict, strict=False)
|
||||
m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False)
|
||||
if args.real_score_model_name_or_path is None:
|
||||
m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False)
|
||||
m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False)
|
||||
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
|
||||
assert len(u) == 0
|
||||
@@ -1116,6 +1212,33 @@ def main():
|
||||
if trainable_module_name in name:
|
||||
param.requires_grad = True
|
||||
break
|
||||
if use_gan:
|
||||
add_wan_gan_discriminator(
|
||||
fake_score_transformer3d, args.gan_feature_layers, args.gan_block_ffn_dim, args.gan_block_num_heads,
|
||||
weight_dtype, seed=args.seed or 0,
|
||||
)
|
||||
|
||||
# Create EMA for the generator_transformer3d.
|
||||
if args.use_ema:
|
||||
if zero_stage == 3:
|
||||
raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.")
|
||||
|
||||
if fsdp_plugin is not None:
|
||||
if getattr(fsdp_plugin, "fsdp_version", 1) != 1:
|
||||
raise NotImplementedError("FSDPEMA only supports FSDP1.")
|
||||
ema_module = WanTransformer3DModel_SelfForcing.from_config(generator_transformer3d.config)
|
||||
ema_generator_transformer3d = FSDPEMA(
|
||||
ema_module, source=generator_transformer3d,
|
||||
accelerator=accelerator, fsdp_plugin=fsdp_plugin,
|
||||
decay=args.ema_decay,
|
||||
)
|
||||
else:
|
||||
ema_generator_transformer3d = EMAModel(
|
||||
generator_transformer3d.parameters(),
|
||||
decay=args.ema_decay,
|
||||
model_cls=WanTransformer3DModel_SelfForcing,
|
||||
model_config=generator_transformer3d.config,
|
||||
)
|
||||
|
||||
# `accelerate` 0.16.0 will have better support for customized saving
|
||||
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
|
||||
@@ -1178,6 +1301,32 @@ def main():
|
||||
accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
|
||||
accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)
|
||||
|
||||
if args.use_ema:
|
||||
def save_ema_hook(models, weights, output_dir):
|
||||
if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process:
|
||||
ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
|
||||
|
||||
def load_ema_hook(models, input_dir):
|
||||
ema_path = os.path.join(input_dir, "transformer_ema")
|
||||
if isinstance(ema_generator_transformer3d, FSDPEMA):
|
||||
ema_generator_transformer3d.load_pretrained(ema_path)
|
||||
else:
|
||||
_, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True)
|
||||
load_model = WanTransformer3DModel_SelfForcing.from_pretrained(
|
||||
input_dir, subfolder="transformer_ema",
|
||||
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
||||
low_cpu_mem_usage=True,
|
||||
torch_dtype=torch.float32,
|
||||
)
|
||||
load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config)
|
||||
load_model.load_state_dict(ema_kwargs)
|
||||
ema_generator_transformer3d.load_state_dict(load_model.state_dict())
|
||||
ema_generator_transformer3d.to(accelerator.device)
|
||||
del load_model
|
||||
|
||||
accelerator.register_save_state_pre_hook(save_ema_hook)
|
||||
accelerator.register_load_state_pre_hook(load_ema_hook)
|
||||
|
||||
if args.gradient_checkpointing:
|
||||
generator_transformer3d.enable_gradient_checkpointing()
|
||||
fake_score_transformer3d.enable_gradient_checkpointing()
|
||||
@@ -1269,6 +1418,11 @@ def main():
|
||||
if accelerator.is_main_process:
|
||||
print(f"Set {name} to lr : {args.learning_rate_critic / 2}")
|
||||
break
|
||||
if use_gan:
|
||||
fake_trainable_params_optim[0]['params'].extend(
|
||||
param for name, param in fake_score_transformer3d.named_parameters()
|
||||
if name.startswith("gan_head.") and name not in in_already
|
||||
)
|
||||
|
||||
if args.use_came:
|
||||
optimizer = optimizer_cls(
|
||||
@@ -1311,7 +1465,7 @@ def main():
|
||||
args.random_hw_adapt = False
|
||||
|
||||
# Get the dataset
|
||||
if args.train_mode != "normal" or args.use_teacher_forcing:
|
||||
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
|
||||
train_dataset = ImageVideoDataset(
|
||||
args.train_data_meta, args.train_data_dir,
|
||||
video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames,
|
||||
@@ -1339,7 +1493,7 @@ def main():
|
||||
|
||||
return length_to_frame_num
|
||||
|
||||
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing:
|
||||
if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan:
|
||||
aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
|
||||
batch_sampler_generator = torch.Generator().manual_seed(args.seed)
|
||||
batch_sampler = AspectRatioBatchImageVideoSampler(
|
||||
@@ -1682,6 +1836,9 @@ def main():
|
||||
if args.train_mode != "normal":
|
||||
clip_image_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
|
||||
|
||||
if args.use_ema:
|
||||
ema_generator_transformer3d.to(accelerator.device)
|
||||
|
||||
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
|
||||
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
|
||||
if overrode_max_train_steps:
|
||||
@@ -1772,6 +1929,8 @@ def main():
|
||||
vae_stream_2 = None
|
||||
|
||||
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
|
||||
min_score_timestep = int(0.02 * noise_scheduler.config.num_train_timesteps)
|
||||
max_score_timestep = int(0.98 * noise_scheduler.config.num_train_timesteps)
|
||||
|
||||
def randomize_denoising_step_indices(
|
||||
denoising_step_indices_list,
|
||||
@@ -1832,6 +1991,8 @@ def main():
|
||||
# must be averaged by contribution count, not by gradient_accumulation_steps.
|
||||
train_gen_log_count = 0
|
||||
train_denoising_loss = 0.0
|
||||
train_gan_g_loss = 0.0
|
||||
train_gan_d_loss = 0.0
|
||||
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
|
||||
for step, batch in enumerate(train_dataloader):
|
||||
# Data batch sanity check
|
||||
@@ -1852,7 +2013,7 @@ def main():
|
||||
save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.mp4", rescale=True)
|
||||
|
||||
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
|
||||
if args.train_mode != "normal" or args.use_teacher_forcing:
|
||||
if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
|
||||
# Convert images to latent space
|
||||
pixel_values = batch["pixel_values"].to(weight_dtype)
|
||||
|
||||
@@ -1964,7 +2125,7 @@ def main():
|
||||
pixel_values_bs = pixel_values_bs.sample()
|
||||
new_pixel_values.append(pixel_values_bs)
|
||||
return torch.cat(new_pixel_values, dim = 0)
|
||||
if args.use_teacher_forcing:
|
||||
if args.use_teacher_forcing or use_gan:
|
||||
clean_latents = _batch_encode_vae(pixel_values)
|
||||
else:
|
||||
clean_latents = None
|
||||
@@ -1996,7 +2157,7 @@ def main():
|
||||
clip_context.append(_clip_context)
|
||||
clip_context = torch.cat(clip_context)
|
||||
|
||||
if args.use_teacher_forcing:
|
||||
if args.use_teacher_forcing or use_gan:
|
||||
target_shape = clean_latents.size()
|
||||
else:
|
||||
target_shape = mask_latents.size()
|
||||
@@ -2504,7 +2665,7 @@ def main():
|
||||
generator_noise_list = [generator_noise[i] for i in range(bsz)]
|
||||
clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None
|
||||
|
||||
# Use block_mask for causal training (一次性处理整个视频)
|
||||
# Use block_mask for causal training (process the whole video at once)
|
||||
generator_pred = generator_transformer3d(
|
||||
x=generator_noise_list,
|
||||
context=prompt_embeds,
|
||||
@@ -2537,6 +2698,8 @@ def main():
|
||||
# Common code for both KV cache and block mask training
|
||||
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
|
||||
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
|
||||
if args.score_timestep_clamp:
|
||||
generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep)
|
||||
generator_denoised_input = add_noise(
|
||||
generator_pred,
|
||||
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
|
||||
@@ -2642,17 +2805,53 @@ def main():
|
||||
train_dmd_loss += avg_dmd_loss.item()
|
||||
train_gen_log_count += 1
|
||||
|
||||
# One-Forcing GAN generator loss: gradients flow through fake_score into generator_pred
|
||||
generator_loss = dmd_loss
|
||||
if args.gan_g_weight > 0:
|
||||
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
|
||||
gan_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
|
||||
if args.score_timestep_clamp:
|
||||
gan_timestep = gan_timestep.clamp(min_score_timestep, max_score_timestep)
|
||||
gan_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng)
|
||||
gan_inputs = [add_noise(generator_pred, gan_noise, gan_timestep)]
|
||||
if args.relativistic_discriminator:
|
||||
real_latents = clean_latents[:, :, :generator_pred.size(2)].to(generator_pred.dtype)
|
||||
gan_inputs.append(add_noise(real_latents, gan_noise, gan_timestep))
|
||||
num_gan_inputs = len(gan_inputs)
|
||||
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
|
||||
_, gan_logits = fake_score_transformer3d(
|
||||
x=torch.cat(gan_inputs).to(dtype=weight_dtype),
|
||||
context=repeat_batch(prompt_embeds, num_gan_inputs),
|
||||
t=gan_timestep.repeat(num_gan_inputs),
|
||||
seq_len=seq_len,
|
||||
y=repeat_batch(inpaint_latents, num_gan_inputs) if args.train_mode != "normal" else None,
|
||||
clip_fea=repeat_batch(clip_context, num_gan_inputs) if args.train_mode != "normal" else None,
|
||||
gan_classify=True,
|
||||
)
|
||||
gan_logits = gan_logits.float().chunk(num_gan_inputs)
|
||||
if args.relativistic_discriminator:
|
||||
gan_g_loss = F.softplus(-(gan_logits[0] - gan_logits[1])).mean()
|
||||
else:
|
||||
gan_g_loss = F.softplus(-gan_logits[0]).mean()
|
||||
gan_g_loss = gan_g_loss * args.gan_g_weight
|
||||
train_gan_g_loss += accelerator.gather(gan_g_loss.detach().repeat(args.train_batch_size)).mean().item()
|
||||
generator_loss = generator_loss + gan_g_loss
|
||||
|
||||
if args.low_vram:
|
||||
real_score_transformer3d = real_score_transformer3d.to("cpu")
|
||||
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
|
||||
if args.gan_g_weight == 0:
|
||||
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
accelerator.backward(dmd_loss)
|
||||
accelerator.backward(generator_loss)
|
||||
if accelerator.sync_gradients:
|
||||
accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
|
||||
optimizer.step()
|
||||
lr_scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
if args.gan_g_weight > 0:
|
||||
# Drop the fake_score grads produced by the generator GAN loss
|
||||
critic_optimizer.zero_grad()
|
||||
|
||||
if args.low_vram:
|
||||
fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device)
|
||||
@@ -2941,6 +3140,8 @@ def main():
|
||||
|
||||
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
|
||||
critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
|
||||
if args.score_timestep_clamp:
|
||||
critic_timestep = critic_timestep.clamp(min_score_timestep, max_score_timestep)
|
||||
critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng)
|
||||
|
||||
fake_score_denoised_input = add_noise(
|
||||
@@ -2949,15 +3150,25 @@ def main():
|
||||
critic_timestep
|
||||
)
|
||||
|
||||
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
|
||||
fake_score_denoised_output = fake_score_transformer3d(
|
||||
x=fake_score_denoised_input,
|
||||
context=prompt_embeds,
|
||||
t=critic_timestep,
|
||||
seq_len=seq_len,
|
||||
y=inpaint_latents if args.train_mode != "normal" else None,
|
||||
clip_fea=clip_context if args.train_mode != "normal" else None,
|
||||
)
|
||||
# Generated and real latents share critic_noise / critic_timestep. Prepare the
|
||||
# real input and the R1/R2 perturbations once; they are consumed by the critic
|
||||
# forward below and by the R1/R2 penalties after the main backward(s).
|
||||
real_noisy = None
|
||||
r1_perturbation = None
|
||||
r2_perturbation = None
|
||||
if use_gan:
|
||||
real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype)
|
||||
real_noisy = add_noise(real_latents, critic_noise, critic_timestep)
|
||||
del real_latents
|
||||
if args.r1_weight > 0:
|
||||
r1_perturbation = args.r1_sigma * torch.randn(
|
||||
real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng
|
||||
)
|
||||
if args.r2_weight > 0:
|
||||
r2_perturbation = args.r2_sigma * torch.randn(
|
||||
fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype,
|
||||
device=fake_score_denoised_input.device, generator=torch_rng
|
||||
)
|
||||
|
||||
def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
|
||||
noise_pred = noise_pred.float()
|
||||
@@ -2971,11 +3182,132 @@ def main():
|
||||
final_loss = masked_loss.mean()
|
||||
return final_loss
|
||||
|
||||
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
|
||||
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
|
||||
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
|
||||
|
||||
accelerator_fake_score_transformer3d.backward(denoising_loss)
|
||||
if use_gan and not args.relativistic_discriminator:
|
||||
# Memory-optimized critic: the non-relativistic gan_d_loss =
|
||||
# softplus(-real) + softplus(fake) is separable, so run fake and real as two
|
||||
# sequential 1x-batch forwards and backward the fake pass first. Backwarding
|
||||
# fake frees its gradient-checkpointing block inputs before the real forward,
|
||||
# dropping the peak activation from 2x batch to 1x batch. Both backward(s)
|
||||
# accumulate into fake_score.grad, giving the identical update as the combined
|
||||
# 2x-batch pass. R1/R2 below stay as their own 2x-batch re-forwards.
|
||||
# -- Fake pass (1x batch): denoising output + fake discriminator logit --
|
||||
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
|
||||
fake_score_denoised_output, fake_gan_logits = fake_score_transformer3d(
|
||||
x=fake_score_denoised_input,
|
||||
context=prompt_embeds,
|
||||
t=critic_timestep,
|
||||
seq_len=seq_len,
|
||||
y=inpaint_latents if args.train_mode != "normal" else None,
|
||||
clip_fea=clip_context if args.train_mode != "normal" else None,
|
||||
gan_classify=True,
|
||||
)
|
||||
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
|
||||
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
|
||||
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
|
||||
|
||||
fake_logit = fake_gan_logits.float()
|
||||
fake_d_loss = F.softplus(fake_logit).mean() * args.gan_d_weight
|
||||
critic_loss = denoising_loss + fake_d_loss
|
||||
accelerator_fake_score_transformer3d.backward(critic_loss)
|
||||
# Fake graph freed: the transformer blocks' stored inputs are released here.
|
||||
fake_d_detached = fake_d_loss.detach()
|
||||
del critic_loss, fake_gan_logits, fake_logit, fake_d_loss, fake_score_denoised_output
|
||||
|
||||
# -- Real pass (1x batch): discriminator logit only --
|
||||
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
|
||||
_, real_gan_logits = fake_score_transformer3d(
|
||||
x=real_noisy,
|
||||
context=prompt_embeds,
|
||||
t=critic_timestep,
|
||||
seq_len=seq_len,
|
||||
y=inpaint_latents if args.train_mode != "normal" else None,
|
||||
clip_fea=clip_context if args.train_mode != "normal" else None,
|
||||
gan_classify=True,
|
||||
)
|
||||
real_logit = real_gan_logits.float()
|
||||
real_d_loss = F.softplus(-real_logit).mean() * args.gan_d_weight
|
||||
accelerator_fake_score_transformer3d.backward(real_d_loss)
|
||||
# Real graph freed. Combine the two detached halves for logging; gan_d_loss is
|
||||
# kept alive because the progress-bar logging below reads it.
|
||||
gan_d_loss = fake_d_detached + real_d_loss.detach()
|
||||
train_gan_d_loss += accelerator.gather(gan_d_loss.repeat(args.train_batch_size)).mean().item()
|
||||
del fake_d_detached, real_gan_logits, real_logit, real_d_loss
|
||||
else:
|
||||
# Combined path: 2x-batch forward for relativistic GAN (both logits coupled by
|
||||
# softplus(-(real - fake))), or 1x-batch forward when GAN is disabled.
|
||||
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
|
||||
if use_gan:
|
||||
fake_score_denoised_output, gan_logits = fake_score_transformer3d(
|
||||
x=torch.cat([fake_score_denoised_input, real_noisy]),
|
||||
context=repeat_batch(prompt_embeds, 2),
|
||||
t=critic_timestep.repeat(2),
|
||||
seq_len=seq_len,
|
||||
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
|
||||
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
|
||||
gan_classify=True,
|
||||
)
|
||||
fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0]
|
||||
gan_logits = gan_logits.float().chunk(2)
|
||||
else:
|
||||
fake_score_denoised_output = fake_score_transformer3d(
|
||||
x=fake_score_denoised_input,
|
||||
context=prompt_embeds,
|
||||
t=critic_timestep,
|
||||
seq_len=seq_len,
|
||||
y=inpaint_latents if args.train_mode != "normal" else None,
|
||||
clip_fea=clip_context if args.train_mode != "normal" else None,
|
||||
)
|
||||
|
||||
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
|
||||
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
|
||||
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
|
||||
|
||||
critic_loss = denoising_loss
|
||||
if use_gan:
|
||||
fake_logit, real_logit = gan_logits[0], gan_logits[1]
|
||||
if args.relativistic_discriminator:
|
||||
gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
|
||||
else:
|
||||
gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
|
||||
gan_d_loss = gan_d_loss * args.gan_d_weight
|
||||
train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
|
||||
critic_loss = critic_loss + gan_d_loss
|
||||
|
||||
accelerator_fake_score_transformer3d.backward(critic_loss)
|
||||
del critic_loss
|
||||
|
||||
if use_gan and args.r1_weight > 0:
|
||||
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
|
||||
_, r1_logits = fake_score_transformer3d(
|
||||
x=torch.cat([real_noisy + r1_perturbation, real_noisy]),
|
||||
context=repeat_batch(prompt_embeds, 2),
|
||||
t=critic_timestep.repeat(2),
|
||||
seq_len=seq_len,
|
||||
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
|
||||
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
|
||||
gan_classify=True,
|
||||
)
|
||||
r1_logits = r1_logits.float().chunk(2)
|
||||
r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean()
|
||||
accelerator_fake_score_transformer3d.backward(r1_loss)
|
||||
del r1_loss
|
||||
|
||||
if use_gan and args.r2_weight > 0:
|
||||
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
|
||||
_, r2_logits = fake_score_transformer3d(
|
||||
x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]),
|
||||
context=repeat_batch(prompt_embeds, 2),
|
||||
t=critic_timestep.repeat(2),
|
||||
seq_len=seq_len,
|
||||
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
|
||||
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
|
||||
gan_classify=True,
|
||||
)
|
||||
r2_logits = r2_logits.float().chunk(2)
|
||||
r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean()
|
||||
accelerator_fake_score_transformer3d.backward(r2_loss)
|
||||
del r2_loss
|
||||
|
||||
if accelerator_fake_score_transformer3d.sync_gradients:
|
||||
accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm)
|
||||
critic_optimizer.step()
|
||||
@@ -2988,13 +3320,22 @@ def main():
|
||||
|
||||
# Checks if the accelerator has performed an optimization step behind the scenes
|
||||
if accelerator.sync_gradients:
|
||||
# EMA tracks the generator only on rounds where it actually ran optimizer.step();
|
||||
# with gen_update_interval > 1 the critic-only rounds leave the generator unchanged.
|
||||
if args.use_ema and generator_update:
|
||||
ema_generator_transformer3d.step(generator_transformer3d.parameters())
|
||||
|
||||
progress_bar.update(1)
|
||||
global_step += 1
|
||||
accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step)
|
||||
train_log = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}
|
||||
if use_gan:
|
||||
train_log.update({"train_gan_g_loss": train_gan_g_loss / max(train_gen_log_count, 1), "train_gan_d_loss": train_gan_d_loss})
|
||||
accelerator.log(train_log, step=global_step)
|
||||
train_dmd_loss = 0.0
|
||||
train_gen_log_count = 0
|
||||
train_denoising_loss = 0.0
|
||||
train_gan_g_loss = 0.0
|
||||
train_gan_d_loss = 0.0
|
||||
|
||||
if global_step % args.checkpointing_steps == 0:
|
||||
if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
|
||||
@@ -3038,6 +3379,10 @@ def main():
|
||||
|
||||
if args.validation_prompts is not None and global_step % args.validation_steps == 0:
|
||||
with progress_bar.paused():
|
||||
if args.use_ema:
|
||||
# Store the generator parameters temporarily and load the EMA parameters to perform inference.
|
||||
ema_generator_transformer3d.store(generator_transformer3d.parameters())
|
||||
ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
|
||||
log_validation(
|
||||
vae,
|
||||
text_encoder,
|
||||
@@ -3050,8 +3395,15 @@ def main():
|
||||
weight_dtype,
|
||||
global_step,
|
||||
)
|
||||
if args.use_ema:
|
||||
# Switch back to the original generator parameters.
|
||||
ema_generator_transformer3d.restore(generator_transformer3d.parameters())
|
||||
|
||||
logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
|
||||
if use_gan:
|
||||
logs["gan_d_loss"] = gan_d_loss.detach().item()
|
||||
if args.gan_g_weight > 0:
|
||||
logs["gan_g_loss"] = gan_g_loss.detach().item()
|
||||
progress_bar.set_postfix(**logs)
|
||||
|
||||
if global_step >= args.max_train_steps:
|
||||
@@ -3059,6 +3411,10 @@ def main():
|
||||
|
||||
if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
|
||||
with progress_bar.paused():
|
||||
if args.use_ema:
|
||||
# Store the generator parameters temporarily and load the EMA parameters to perform inference.
|
||||
ema_generator_transformer3d.store(generator_transformer3d.parameters())
|
||||
ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
|
||||
log_validation(
|
||||
vae,
|
||||
text_encoder,
|
||||
@@ -3071,6 +3427,9 @@ def main():
|
||||
weight_dtype,
|
||||
global_step,
|
||||
)
|
||||
if args.use_ema:
|
||||
# Switch back to the original generator parameters.
|
||||
ema_generator_transformer3d.restore(generator_transformer3d.parameters())
|
||||
|
||||
# Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever
|
||||
# something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -0,0 +1,374 @@
|
||||
# Wan2.2 Reward LoRA Training Guide
|
||||
|
||||
This document provides a complete workflow for **Reward Backpropagation** fine-tuning of Wan2.2, i.e. optimizing a LoRA so that the generated videos better align with human preferences. It covers environment configuration, data preparation, the reward-specific training arguments (single / multi reward, differentiable VAE decoding, backprop strategy), training validation, and inference testing.
|
||||
|
||||
> **Note**: Wan2.2 is a video generation model supporting text-to-video (T2V), image-to-video (I2V), and text-image-to-video (TI2V). The A14B variants adopt a dual-Transformer architecture (high-noise / low-noise models), while the TI2V-5B variant uses a single-Transformer architecture. Reward LoRA training runs a differentiable rollout (a short denoising trajectory) and back-propagates the reward model's loss into the trainable expert, so it needs more VRAM than standard LoRA fine-tuning. This guide covers training your own reward LoRA for any of the Wan2.2 variants.
|
||||
|
||||
We explore the Reward Backpropagation technique <sup>[1](#ref1) [2](#ref2)</sup> to optimize the generated videos by [Wan2.2](https://github.com/aigc-apps/VideoX-Fun) for better alignment with human preferences. The training script is [train_reward_lora.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/wan2.2/train_reward_lora.py); you can train your own reward LoRA and load it onto the corresponding base model as a plug-in.
|
||||
|
||||
> [!NOTE]
|
||||
> Reward backpropagation tends to reward-hack: as the reward keeps rising the perceptual quality can plateau or degrade (adversarial artifacts in the background). Prefer a **video reward** (e.g. `VideoAlignReward`) or a **multi-reward** mix (e.g. `HPSv3Reward` + `VideoAlignReward`) over a pure image reward to mitigate the loss of temporal dynamism — see [3.4](#34-single-vs-multi-reward-training) and [5. Limitations](#5-limitations).
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
- [1. Environment Configuration](#1-environment-configuration)
|
||||
- [2. Data Preparation](#2-data-preparation)
|
||||
- [2.1 Quick Test Dataset](#21-quick-test-dataset)
|
||||
- [2.2 Reward Model Weights](#22-reward-model-weights)
|
||||
- [2.3 Dataset Structure and metadata.json](#23-dataset-structure-and-metadatajson)
|
||||
- [3. Reward LoRA Training](#3-reward-lora-training)
|
||||
- [3.1 Download Pre-trained Base Model](#31-download-pre-trained-base-model)
|
||||
- [3.2 Quick Start (FSDP)](#32-quick-start-fsdp)
|
||||
- [3.3 Reward-specific Parameter Explanation](#33-reward-specific-parameter-explanation)
|
||||
- [3.4 Single vs Multi-Reward Training](#34-single-vs-multi-reward-training)
|
||||
- [3.5 Differentiable Decoding and `--use_1steps_vae`](#35-differentiable-decoding-and---use_1steps_vae)
|
||||
- [3.6 Training Validation](#36-training-validation)
|
||||
- [4. Inference Testing](#4-inference-testing)
|
||||
- [5. Limitations](#5-limitations)
|
||||
- [6. More Resources](#6-more-resources)
|
||||
- [References](#references)
|
||||
|
||||
---
|
||||
|
||||
## 1. Environment Configuration
|
||||
|
||||
**Method 1: Using requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**Method 2: Manual Dependency Installation**
|
||||
|
||||
```bash
|
||||
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
|
||||
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
|
||||
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
|
||||
pip install yunchang xfuser modelscope openpyxl
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
pip install deepspeed==0.17.0 numpy==1.26.4
|
||||
```
|
||||
|
||||
> Some video reward models (e.g. `VideoAlignReward`, `HPSv3Reward`) are built on a Qwen2-VL backbone and require `transformers==4.51.3`. Install the matching version when using those rewards.
|
||||
|
||||
**Method 3: Using Docker**
|
||||
|
||||
When using Docker, please ensure that the graphics card driver and CUDA environment are correctly installed on your machine, then execute the following commands:
|
||||
|
||||
```
|
||||
# pull image
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
|
||||
# enter image
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Data Preparation
|
||||
|
||||
### 2.1 Quick Test Dataset
|
||||
|
||||
Reward LoRA training runs a differentiable rollout from text prompts, so a **prompt-only** dataset is sufficient for the standard T2V reward path (`train_mode="normal"`), which is the recommended setup. For the I2V / inpaint path, a video dataset provides the reference conditioning.
|
||||
|
||||
```bash
|
||||
# Download official example dataset
|
||||
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
|
||||
```
|
||||
|
||||
### 2.2 Reward Model Weights
|
||||
|
||||
The backbone / checkpoint source differs per reward:
|
||||
|
||||
- **`HPSReward` (HPS v2.1) and `MPSReward`**: the score heads are tiny and are **downloaded automatically on first run** into the torch hub cache (`~/.cache/torch/hub`); note that **HPS additionally needs two special config files patched into `hpsv2`'s bundled open_clip** (see below), while MPS needs no extra setup;
|
||||
- **`VideoAlignReward` (VideoReward) and `HPSv3Reward`**: built on Qwen2-VL backbones and fairly large, so **stage them locally in advance** and point `--reward_fn_kwargs` at the local paths.
|
||||
|
||||
The dependencies and sources are summarized below (see [3.3](#33-reward-specific-parameter-explanation) for the `--reward_fn` values):
|
||||
|
||||
| Reward | Backbone / dependency | Score head / checkpoint | Manual download |
|
||||
|--|--|--|--|
|
||||
| `HPSReward` (`version="v2.1"`) | open_clip `ViT-H-14`, requires `pip install hpsv2` + the `factory.py` patch and BPE vocab | `HPS_v2.1_compressed.pt` (auto-downloaded) | No, override with `model_path` |
|
||||
| `MPSReward` | `laion/CLIP-ViT-H-14-laion2B-s32B-b79K` (config/tokenizer auto-downloaded); MPS code is vendored under `videox_fun/reward/MPS/` | `MPS_overall.pth` (auto-downloaded) | No, override with `model_path` |
|
||||
| `VideoAlignReward` | `Qwen/Qwen2-VL-2B-Instruct` | `KlingTeam/VideoReward` (contains `model_config.json` + weights) | **Yes**, `model_path` is required |
|
||||
| `HPSv3Reward` | `Qwen/Qwen2-VL-7B-Instruct` | `HPSv3.safetensors` from `MizzenAI/HPSv3` | Recommended; auto-downloaded from HF when omitted |
|
||||
|
||||
**HPS v2.1 (`HPSReward`)** — besides `pip install hpsv2`, you must drop two special config files into `hpsv2`'s bundled open_clip: `create_model_and_transforms("ViT-H-14", ...)` uses arguments added by the HPS fork, so it needs the patched `factory.py`; `get_tokenizer("ViT-H-14")` needs the BPE vocab `bpe_simple_vocab_16e6.txt.gz`. Loading fails if either is missing. The `ViT-H-14` backbone and the `HPS_v2.1_compressed.pt` head download automatically to the torch hub cache on first run:
|
||||
|
||||
```bash
|
||||
pip install hpsv2
|
||||
# Patch hpsv2's bundled open_clip factory.py and add the BPE vocab file
|
||||
site_packages=$(python -c "import site; print(site.getsitepackages()[0])")
|
||||
wget -O $site_packages/hpsv2/src/open_clip/factory.py https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/package/patches/hpsv2_src_open_clip_factory_patches.py
|
||||
wget -O $site_packages/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz https://github.com/tgxs002/HPSv2/raw/refs/heads/master/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="HPSReward" \
|
||||
--reward_fn_kwargs='{"version": "v2.1"}'
|
||||
```
|
||||
|
||||
**MPS (`MPSReward`)** — the MPS code is vendored under `videox_fun/reward/MPS/` and needs no patching; the `MPS_overall.pth` head downloads automatically from OSS, and the config + tokenizer of `laion/CLIP-ViT-H-14-laion2B-s32B-b79K` download automatically from HuggingFace (if HF is unreachable, point `processor_name_or_path` to a local mirror):
|
||||
|
||||
```bash
|
||||
--reward_fn="MPSReward" \
|
||||
--reward_fn_kwargs='{"max_reward": 1, "loss_scale": 1}'
|
||||
```
|
||||
|
||||
**VideoReward (`VideoAlignReward`)** — needs the Qwen2-VL-2B backbone + a VideoReward head directory (containing `model_config.json` and weights). Download both in advance; `model_path` is required:
|
||||
|
||||
```bash
|
||||
modelscope download --model Qwen/Qwen2-VL-2B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-2B-Instruct
|
||||
huggingface-cli download KlingTeam/VideoReward --local-dir models/Personalized_Model/VideoReward
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="VideoAlignReward" \
|
||||
--reward_fn_kwargs='{"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "reward_dim": "TA+MQ", "fps": 16, "use_norm": true, "differentiable": true, "max_reward": 2.0}'
|
||||
```
|
||||
|
||||
**HPSv3 (`HPSv3Reward`)** — needs the Qwen2-VL-7B backbone + `HPSv3.safetensors`. When `checkpoint_path` is omitted it downloads automatically from HF `MizzenAI/HPSv3`; staging it locally is recommended:
|
||||
|
||||
```bash
|
||||
modelscope download --model Qwen/Qwen2-VL-7B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-7B-Instruct
|
||||
modelscope download --model MizzenAI/HPSv3 --local_dir models/Personalized_Model/HPSv3
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="HPSv3Reward" \
|
||||
--reward_fn_kwargs='{"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}'
|
||||
```
|
||||
|
||||
> With `--backprop`, `VideoAlignReward` / `HPSv3Reward` must set `"differentiable": true`, otherwise their `@torch.no_grad()` path cuts the gradient chain (see [3.3](#33-reward-specific-parameter-explanation)).
|
||||
|
||||
### 2.3 Dataset Structure and metadata.json
|
||||
|
||||
```
|
||||
📦 datasets/
|
||||
├── 📂 my_dataset/
|
||||
│ ├── 📂 train/
|
||||
│ │ ├── 📄 video001.mp4
|
||||
│ │ ├── 📄 video002.mp4
|
||||
│ │ └── 📄 ...
|
||||
│ └── 📄 metadata.json
|
||||
```
|
||||
|
||||
**Relative Path Format** (example):
|
||||
```json
|
||||
[
|
||||
{
|
||||
"file_path": "train/video001.mp4",
|
||||
"text": "A beautiful sunset over the ocean, golden hour lighting",
|
||||
"type": "video",
|
||||
"width": 1024,
|
||||
"height": 1024
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
**Key Field Descriptions**:
|
||||
- `file_path`: Video path (relative or absolute path).
|
||||
- `text`: Video description (English prompt).
|
||||
- `type`: Data type, fixed as `"video"`.
|
||||
- `width` / `height`: Video width and height (**highly recommended**, used for bucket training). You can use `scripts/process_json_add_width_and_height.py` to extract these fields.
|
||||
|
||||
**Relative vs Absolute Path**:
|
||||
- Relative: `export DATASET_NAME="datasets/X-Fun-Videos-Demo/"` and `export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"`.
|
||||
- Absolute: `export DATASET_NAME=""` and `export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"`.
|
||||
|
||||
---
|
||||
|
||||
## 3. Reward LoRA Training
|
||||
|
||||
### 3.1 Download Pre-trained Base Model
|
||||
|
||||
```bash
|
||||
# Create model directory
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# Download Wan2.2 official weights
|
||||
# T2V model (text-to-video, dual-Transformer)
|
||||
modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-T2V-A14B
|
||||
# or I2V model (image-to-video, dual-Transformer)
|
||||
# modelscope download --model Wan-AI/Wan2.2-I2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-I2V-A14B
|
||||
# or TI2V model (text-image-to-video, single-Transformer)
|
||||
# modelscope download --model Wan-AI/Wan2.2-TI2V-5B --local_dir models/Diffusion_Transformer/Wan2.2-TI2V-5B
|
||||
```
|
||||
|
||||
### 3.2 Quick Start (FSDP)
|
||||
|
||||
Reward LoRA training keeps the full denoising rollout and the reward model's forward graph in memory, so FSDP (`FULL_SHARD`) is recommended. With `--boundary_type="high"`, the *other* (frozen) expert is already sharded and resident on GPU and is reused during validation, so no second copy is re-loaded from disk.
|
||||
|
||||
**Wan2.2 T2V Reward LoRA Training Example (single HPS v2.1 reward, FSDP)**:
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
export NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False \
|
||||
scripts/wan2.2/train_reward_lora.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=1 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=50 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_wan2.2_reward_lora" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--enable_bucket \
|
||||
--boundary_type="high" \
|
||||
--rank=128 \
|
||||
--network_alpha=64 \
|
||||
--target_name="q,k,v,ffn.0,ffn.2" \
|
||||
--train_mode="normal" \
|
||||
--guidance_scale=5.0 \
|
||||
--num_inference_steps=40 \
|
||||
--num_decoded_latents=3 \
|
||||
--reward_fn="HPSReward" \
|
||||
--reward_fn_kwargs='{"version": "v2.1"}' \
|
||||
--backprop \
|
||||
--backprop_strategy="tail" \
|
||||
--backprop_num_steps=1 \
|
||||
--low_vram
|
||||
```
|
||||
|
||||
> The `train_reward_lora.sh` script in this directory provides a ready-to-run template. Adjust `MODEL_NAME`, dataset paths, `--reward_fn`, and `--boundary_type` to your setup.
|
||||
|
||||
### 3.3 Reward-specific Parameter Explanation
|
||||
|
||||
Besides the standard LoRA arguments (see [README_TRAIN_LORA.md](./README_TRAIN_LORA.md)), the reward trainer adds the following:
|
||||
|
||||
| Parameter | Description | Example Value |
|
||||
|-----------|-------------|---------------|
|
||||
| `--reward_fn` | Reward model class name(s). A **comma-separated list** enables multi-reward training; the differentiable losses are combined by a weighted sum. Available: `AestheticReward`, `HPSReward`, `PickScoreReward`, `MPSReward`, `HPSv3Reward`, `VideoAlignReward`. | `"HPSv3Reward,VideoAlignReward"` |
|
||||
| `--reward_fn_kwargs` | Keyword args for the reward model(s). A **flat JSON dict** for a single reward; a **nested JSON dict keyed by reward name** for multiple rewards. | `'{"VideoAlignReward": {"max_reward": 2.0}, "HPSv3Reward": {"max_reward": 12}}'` |
|
||||
| `--multi_reward_weights` | JSON dict of per-reward loss weights. Normalized to sum to 1; if omitted with multiple rewards, equal weights are used. Ignored for a single reward. | `'{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}'` |
|
||||
| `--num_inference_steps` | Number of denoising steps in the differentiable rollout and validation. | 40 |
|
||||
| `--num_decoded_latents` | Number of latent frames decoded back to pixels for the reward. The chunked causal VAE limits this to `<= 3` (gradient-checkpointing constraint) unless `--use_1steps_vae` is set. | 3 |
|
||||
| `--num_reward_frames` | Randomly select this many decoded frames for reward evaluation. `None` = use all. Only applied to **image** rewards (see below). | 2 |
|
||||
| `--use_1steps_vae` | Use the single-pass Wan2.2 VAE (`AutoencoderKLWan3_8_1Steps`) for differentiable decoding; numerically equivalent, decodes the whole latent in one forward and lifts the `<= 3` latent limit. | - |
|
||||
| `--backprop` | Enable the reward-backpropagation training mode. | - |
|
||||
| `--backprop_strategy` | Which rollout steps keep a gradient graph: `last`, `tail`, `uniform`, `random`. | `tail` |
|
||||
| `--backprop_num_steps` | Number of backprop steps, used when `--backprop_strategy` is `tail`/`uniform`/`random`. | 1 |
|
||||
| `--first_frame_mse_weight` | Weight for the first-frame latent MSE auxiliary loss in inpaint mode. | 1.0 |
|
||||
|
||||
**Image vs Video rewards (`is_image_reward`)**:
|
||||
- **Image rewards** (`AestheticReward`, `HPSReward`, `PickScoreReward`, `MPSReward`, `HPSv3Reward`) score individual frames, so frame subsampling (`--num_reward_frames`) is safe.
|
||||
- **Video rewards** (`VideoAlignReward`, `is_image_reward=False`) need the full temporal sequence to judge motion quality; collapsing frames would degrade the video to a near-static input and break the motion dimension. For these, all decoded frames are kept (count bounded by `--num_decoded_latents`).
|
||||
- In inpaint mode (`train_mode="inpaint"`), the first frame is dropped before reward evaluation for every reward.
|
||||
|
||||
> **IMPORTANT**: when using `--backprop`, a video reward such as `VideoAlignReward` / `HPSv3Reward` must be constructed with `"differentiable": true` inside `--reward_fn_kwargs`; otherwise its `@torch.no_grad()` path cuts the gradient chain entirely.
|
||||
|
||||
### 3.4 Single vs Multi-Reward Training
|
||||
|
||||
A single reward uses `--reward_fn="HPSv3Reward"` with a flat `--reward_fn_kwargs`. To combine several objectives, pass a comma-separated list, a nested kwargs dict, and (optionally) weights. Per-reward loss and reward values are logged separately (e.g. `train/loss_HPSv3Reward`, `train/reward_VideoAlignReward`), while the optimized `loss = Σ wᵢ · lossᵢ`.
|
||||
|
||||
**Multi-reward example (HPSv3 + VideoAlign, T2V)** — the strongest alignment setup, since the video reward supplies the temporal (motion / consistency) signal that image rewards cannot:
|
||||
|
||||
```bash
|
||||
--num_decoded_latents=13 \
|
||||
--num_reward_frames=2 \
|
||||
--use_1steps_vae \
|
||||
--reward_fn="VideoAlignReward,HPSv3Reward" \
|
||||
--reward_fn_kwargs='{"VideoAlignReward": {"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "fps": 16, "reward_dim": "TA+MQ", "differentiable": true, "use_norm": true, "max_reward": 2.0, "loss_scale": 1.0}, "HPSv3Reward": {"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}}' \
|
||||
--multi_reward_weights='{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}' \
|
||||
--backprop \
|
||||
--backprop_strategy="tail" \
|
||||
--backprop_num_steps=1
|
||||
```
|
||||
|
||||
### 3.5 Differentiable Decoding and `--use_1steps_vae`
|
||||
|
||||
The default chunked causal VAE (`AutoencoderKLWan3_8`) applies gradient checkpointing per chunk, which limits differentiable decoding to the **first 3 latent frames** (`--num_decoded_latents <= 3`). Setting `--num_decoded_latents > 3` without `--use_1steps_vae` raises a `ValueError`.
|
||||
|
||||
Enable `--use_1steps_vae` to switch to the single-pass variant, which decodes the whole latent in one forward, applies checkpointing cleanly per decoder block, and lifts the limit — useful when a video reward needs more temporal context. It is numerically equivalent and only affects the `AutoencoderKLWan3_8` (2.2 VAE) path.
|
||||
|
||||
### 3.6 Training Validation
|
||||
|
||||
You can configure validation parameters to periodically generate test videos during training so you can monitor alignment progress. Under FSDP, `log_validation` runs a pipeline forward that is a cross-rank collective, so **all processes must enter the validation call** (it is intentionally *not* wrapped in `is_main_process`); the trainable expert already carries the LoRA, and the frozen expert is reused rather than re-loaded.
|
||||
|
||||
| Parameter | Description | Default Value |
|
||||
|-----------|-------------|---------------|
|
||||
| `--validation_steps` | Run validation every N steps | 2000 |
|
||||
| `--validation_epochs` | Run validation every N epochs | 5 |
|
||||
| `--validation_prompts` | Prompts for validating video generation | None |
|
||||
| `--validation_paths` | Reference image paths for I2V validation (inpaint mode only) | None |
|
||||
|
||||
**Validation Example**:
|
||||
|
||||
```bash
|
||||
--validation_prompts="A stylish woman strolls down a bustling Tokyo street at night, warm neon glow and animated city signage reflecting on the wet pavement; she wears a black leather jacket over a red dress, confident casual demeanor, medium shot, cinematic." \
|
||||
--validation_steps=50 \
|
||||
--validation_epochs=500
|
||||
```
|
||||
|
||||
> Validation videos are saved to the `output_dir` directory. For T2V validation (without a reference image), keep `train_mode="normal"` and omit `--validation_paths`.
|
||||
|
||||
---
|
||||
|
||||
## 4. Inference Testing
|
||||
|
||||
The trained reward LoRA is a normal LoRA and is loaded the same way as [README_TRAIN_LORA.md](./README_TRAIN_LORA.md). For an A14B dual-Transformer model, set `lora_path` / `lora_weight` for the low-noise reward LoRA and `lora_high_path` / `lora_high_weight` for the high-noise reward LoRA in [examples/wan2.2/predict_t2v.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.2/predict_t2v.py) (or `predict_i2v.py`):
|
||||
|
||||
```python
|
||||
# Path to the trained reward LoRA weights, e.g.
|
||||
# "output_dir_wan2.2_reward_lora/checkpoint-xxx/diffusion_pytorch_model.safetensors"
|
||||
lora_path = None
|
||||
lora_high_path = None
|
||||
lora_weight = 0.55
|
||||
lora_high_weight = 0.55
|
||||
```
|
||||
|
||||
Run single-GPU inference with:
|
||||
|
||||
```bash
|
||||
python examples/wan2.2/predict_t2v.py
|
||||
```
|
||||
|
||||
For memory-management modes and multi-GPU parallel inference, refer to the Inference Testing section of [README_TRAIN_LORA.md](./README_TRAIN_LORA.md).
|
||||
|
||||
---
|
||||
|
||||
## 5. Limitations
|
||||
|
||||
1. After training to a certain extent, the reward keeps increasing but the quality of the generated videos no longer improves. The model trickily learns shortcuts (adding artifacts in the background, i.e. adversarial patches) to raise the reward.
|
||||
2. There is still a lack of suitable preference models for video generation. Directly using image preference models cannot evaluate preferences along the temporal dimension (such as dynamism and consistency), and tends to reduce the dynamism of generated videos. This can be partially mitigated by computing the reward using only the first decoded frame, but the impact persists — which is why a video reward (e.g. `VideoAlignReward`) with `--use_1steps_vae` and a larger `--num_decoded_latents`, or a multi-reward mix, is recommended for motion-sensitive alignment.
|
||||
|
||||
---
|
||||
|
||||
## 6. More Resources
|
||||
|
||||
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
|
||||
- **Standard LoRA training**: [README_TRAIN_LORA.md](./README_TRAIN_LORA.md)
|
||||
- **Chinese version**: [README_TRAIN_REWARD_zh-CN.md](./README_TRAIN_REWARD_zh-CN.md)
|
||||
|
||||
---
|
||||
|
||||
## References
|
||||
<ol>
|
||||
<li id="ref1">Clark, Kevin, et al. "Directly fine-tuning diffusion models on differentiable rewards.". In ICLR 2024.</li>
|
||||
<li id="ref2">Prabhudesai, Mihir, et al. "Aligning text-to-image diffusion models with reward backpropagation." arXiv preprint arXiv:2310.03739 (2023).</li>
|
||||
</ol>
|
||||
@@ -0,0 +1,374 @@
|
||||
# Wan2.2 Reward LoRA 训练指南
|
||||
|
||||
本文档提供 Wan2.2 **Reward 反向传播（Reward Backpropagation）** 微调的完整流程，即用 LoRA 优化生成视频，使其更符合人类偏好。内容涵盖环境配置、数据准备、Reward 专用训练参数（单/多 reward、可微 VAE 解码、backprop 策略）、训练验证与推理测试。
|
||||
|
||||
> **说明**：Wan2.2 是一个支持文本到视频（T2V）、图像到视频（I2V）、文本-图像到视频（TI2V）的视频生成模型。A14B 系列采用双 Transformer 架构（高噪声/低噪声模型），TI2V-5B 使用单 Transformer 架构。Reward LoRA 训练会跑一段可微的 rollout（一条简短的去噪轨迹），并把 reward 模型的 loss 反传到可训练的专家上，因此比标准 LoRA 微调更耗显存。本指南覆盖为任意 Wan2.2 变体训练自有 reward LoRA 的流程。
|
||||
|
||||
我们采用 Reward Backpropagation 技术 <sup>[1](#ref1) [2](#ref2)</sup> 来优化 [Wan2.2](https://github.com/aigc-apps/VideoX-Fun) 生成的视频，使其更好地对齐人类偏好。训练脚本为 [train_reward_lora.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/wan2.2/train_reward_lora.py)；你可以训练自己的 reward LoRA，并作为插件加载到对应的基座模型上。
|
||||
|
||||
> [!NOTE]
|
||||
> Reward 反向传播容易出现 reward-hacking：reward 持续上升，但感知质量可能停滞甚至退化（背景出现对抗伪影）。相比纯图像 reward，更推荐使用**视频 reward**（如 `VideoAlignReward`）或**多 reward 组合**（如 `HPSv3Reward` + `VideoAlignReward`），以缓解时序运动性的下降——见 [3.4](#34-单-reward-与多-reward-训练) 与 [五、局限性](#五局限性)。
|
||||
|
||||
---
|
||||
|
||||
## 目录
|
||||
- [一、环境配置](#一环境配置)
|
||||
- [二、数据准备](#二数据准备)
|
||||
- [2.1 快速测试数据集](#21-快速测试数据集)
|
||||
- [2.2 Reward 模型权重](#22-reward-模型权重)
|
||||
- [2.3 数据集结构与 metadata.json](#23-数据集结构与-metadatajson)
|
||||
- [三、Reward LoRA 训练](#三reward-lora-训练)
|
||||
- [3.1 下载预训练基座模型](#31-下载预训练基座模型)
|
||||
- [3.2 快速开始（FSDP）](#32-快速开始fsdp)
|
||||
- [3.3 Reward 专用参数解析](#33-reward-专用参数解析)
|
||||
- [3.4 单 Reward 与多 Reward 训练](#34-单-reward-与多-reward-训练)
|
||||
- [3.5 可微解码与 `--use_1steps_vae`](#35-可微解码与---use_1steps_vae)
|
||||
- [3.6 训练验证](#36-训练验证)
|
||||
- [四、推理测试](#四推理测试)
|
||||
- [五、局限性](#五局限性)
|
||||
- [六、更多资源](#六更多资源)
|
||||
- [参考文献](#参考文献)
|
||||
|
||||
---
|
||||
|
||||
## 一、环境配置
|
||||
|
||||
**方式 1：使用 requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**方式 2：手动安装依赖**
|
||||
|
||||
```bash
|
||||
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
|
||||
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
|
||||
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
|
||||
pip install yunchang xfuser modelscope openpyxl
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
pip install deepspeed==0.17.0 numpy==1.26.4
|
||||
```
|
||||
|
||||
> 部分视频 reward 模型（如 `VideoAlignReward`、`HPSv3Reward`）基于 Qwen2-VL 骨干，需要 `transformers==4.51.3`。使用这些 reward 时请安装对应版本。
|
||||
|
||||
**方式 3：使用 docker**
|
||||
|
||||
使用 docker 的情况下，请保证机器中已经正确安装显卡驱动与 CUDA 环境，然后以此执行以下命令：
|
||||
|
||||
```
|
||||
# pull image
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
|
||||
# enter image
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 二、数据准备
|
||||
|
||||
### 2.1 快速测试数据集
|
||||
|
||||
Reward LoRA 训练从文本 prompt 出发跑可微 rollout，因此标准的 T2V reward 路径（`train_mode="normal"`）只需一个**纯 prompt 数据集**即可，这也是推荐配置。对于 I2V/inpaint 路径，则使用视频数据集来提供参考条件。
|
||||
|
||||
```bash
|
||||
# 下载官方示例数据集
|
||||
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
|
||||
```
|
||||
|
||||
### 2.2 Reward 模型权重
|
||||
|
||||
不同 reward 的骨干/checkpoint 来源不同：
|
||||
|
||||
- **`HPSReward`（HPS v2.1）与 `MPSReward`**：打分头很小，会在**首次运行时自动下载**到 torch hub 缓存（`~/.cache/torch/hub`）；其中 **HPS 还需给 `hpsv2` 内置的 open_clip 补两个特殊配置文件**（见下），MPS 无需额外准备；
|
||||
- **`VideoAlignReward`（VideoReward）与 `HPSv3Reward`**：基于 Qwen2-VL 骨干，体积较大，**建议提前下载到本地**，再用 `--reward_fn_kwargs` 指向本地路径。
|
||||
|
||||
各 reward 的依赖与来源汇总如下（`--reward_fn` 取值见 [3.3](#33-reward-专用参数解析)）：
|
||||
|
||||
| Reward | 骨干 / 依赖 | 打分头 / checkpoint | 手动下载 |
|
||||
|--|--|--|--|
|
||||
| `HPSReward`（`version="v2.1"`） | open_clip `ViT-H-14`，需 `pip install hpsv2` + `factory.py` 补丁与 BPE 词表 | `HPS_v2.1_compressed.pt`（自动下载） | 否，可用 `model_path` 覆盖 |
|
||||
| `MPSReward` | `laion/CLIP-ViT-H-14-laion2B-s32B-b79K`（config/tokenizer 自动下载），MPS 代码已内置于 `videox_fun/reward/MPS/` | `MPS_overall.pth`（自动下载） | 否，可用 `model_path` 覆盖 |
|
||||
| `VideoAlignReward` | `Qwen/Qwen2-VL-2B-Instruct` | `KlingTeam/VideoReward`（含 `model_config.json` 与权重） | **是**，`model_path` 必填 |
|
||||
| `HPSv3Reward` | `Qwen/Qwen2-VL-7B-Instruct` | `MizzenAI/HPSv3` 的 `HPSv3.safetensors` | 建议，缺省时自动从 HF 下载 |
|
||||
|
||||
**HPS v2.1（`HPSReward`）**——除了 `pip install hpsv2`，还要给 `hpsv2` 内置的 open_clip 补两个特殊配置文件：`create_model_and_transforms("ViT-H-14", ...)` 用到了 HPS 分支新增的参数，需要打过补丁的 `factory.py`；`get_tokenizer("ViT-H-14")` 需要 BPE 词表 `bpe_simple_vocab_16e6.txt.gz`。缺任一都会在加载时报错。`ViT-H-14` 骨干与 `HPS_v2.1_compressed.pt` 打分头会在首次运行时自动下载到 torch hub 缓存：
|
||||
|
||||
```bash
|
||||
pip install hpsv2
|
||||
# 覆盖 hpsv2 内置 open_clip 的 factory.py 补丁，并补上 BPE 词表
|
||||
site_packages=$(python -c "import site; print(site.getsitepackages()[0])")
|
||||
wget -O $site_packages/hpsv2/src/open_clip/factory.py https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/package/patches/hpsv2_src_open_clip_factory_patches.py
|
||||
wget -O $site_packages/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz https://github.com/tgxs002/HPSv2/raw/refs/heads/master/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="HPSReward" \
|
||||
--reward_fn_kwargs='{"version": "v2.1"}'
|
||||
```
|
||||
|
||||
**MPS（`MPSReward`）**——MPS 代码已内置于 `videox_fun/reward/MPS/`，无需打补丁；`MPS_overall.pth` 打分头会自动从 OSS 下载，`laion/CLIP-ViT-H-14-laion2B-s32B-b79K` 的 config 与 tokenizer 会自动从 HuggingFace 下载（若无法访问 HF，可用 `processor_name_or_path` 指向本地镜像）：
|
||||
|
||||
```bash
|
||||
--reward_fn="MPSReward" \
|
||||
--reward_fn_kwargs='{"max_reward": 1, "loss_scale": 1}'
|
||||
```
|
||||
|
||||
**VideoReward（`VideoAlignReward`）**——需 Qwen2-VL-2B 骨干 + VideoReward 打分头目录（内含 `model_config.json` 与权重），二者都要提前下载，且 `model_path` 必填：
|
||||
|
||||
```bash
|
||||
modelscope download --model Qwen/Qwen2-VL-2B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-2B-Instruct
|
||||
huggingface-cli download KlingTeam/VideoReward --local-dir models/Personalized_Model/VideoReward
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="VideoAlignReward" \
|
||||
--reward_fn_kwargs='{"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "reward_dim": "TA+MQ", "fps": 16, "use_norm": true, "differentiable": true, "max_reward": 2.0}'
|
||||
```
|
||||
|
||||
**HPSv3（`HPSv3Reward`）**——需 Qwen2-VL-7B 骨干 + `HPSv3.safetensors`；不指定 `checkpoint_path` 时会自动从 HF `MizzenAI/HPSv3` 下载，建议提前落盘：
|
||||
|
||||
```bash
|
||||
modelscope download --model Qwen/Qwen2-VL-7B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-7B-Instruct
|
||||
modelscope download --model MizzenAI/HPSv3 --local_dir models/Personalized_Model/HPSv3
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="HPSv3Reward" \
|
||||
--reward_fn_kwargs='{"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}'
|
||||
```
|
||||
|
||||
> 使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 必须设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会切断梯度链（见 [3.3](#33-reward-专用参数解析)）。
|
||||
|
||||
### 2.3 数据集结构与 metadata.json
|
||||
|
||||
```
|
||||
📦 datasets/
|
||||
├── 📂 my_dataset/
|
||||
│ ├── 📂 train/
|
||||
│ │ ├── 📄 video001.mp4
|
||||
│ │ ├── 📄 video002.mp4
|
||||
│ │ └── 📄 ...
|
||||
│ └── 📄 metadata.json
|
||||
```
|
||||
|
||||
**相对路径格式**（示例）：
|
||||
```json
|
||||
[
|
||||
{
|
||||
"file_path": "train/video001.mp4",
|
||||
"text": "A beautiful sunset over the ocean, golden hour lighting",
|
||||
"type": "video",
|
||||
"width": 1024,
|
||||
"height": 1024
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
**关键字段说明**：
|
||||
- `file_path`：视频路径（相对路径或绝对路径）。
|
||||
- `text`：视频描述（英文 prompt）。
|
||||
- `type`：数据类型，固定为 `"video"`。
|
||||
- `width` / `height`：视频宽高（**强烈建议提供**，用于分桶训练）。可用 `scripts/process_json_add_width_and_height.py` 提取这两个字段。
|
||||
|
||||
**相对路径与绝对路径**：
|
||||
- 相对路径：`export DATASET_NAME="datasets/X-Fun-Videos-Demo/"` 与 `export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"`。
|
||||
- 绝对路径：`export DATASET_NAME=""` 与 `export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"`。
|
||||
|
||||
---
|
||||
|
||||
## 三、Reward LoRA 训练
|
||||
|
||||
### 3.1 下载预训练基座模型
|
||||
|
||||
```bash
|
||||
# 创建模型目录
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# 下载 Wan2.2 官方权重
|
||||
# T2V 模型（文本到视频，双 Transformer）
|
||||
modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-T2V-A14B
|
||||
# 或 I2V 模型（图像到视频，双 Transformer）
|
||||
# modelscope download --model Wan-AI/Wan2.2-I2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-I2V-A14B
|
||||
# 或 TI2V 模型（文本-图像到视频，单 Transformer）
|
||||
# modelscope download --model Wan-AI/Wan2.2-TI2V-5B --local_dir models/Diffusion_Transformer/Wan2.2-TI2V-5B
|
||||
```
|
||||
|
||||
### 3.2 快速开始（FSDP）
|
||||
|
||||
Reward LoRA 训练会把完整的去噪 rollout 与 reward 模型的前向计算图保持在显存中，因此推荐使用 FSDP（`FULL_SHARD`）。当 `--boundary_type="high"` 时，*另一个*（冻结的）专家在训练期已被分片并常驻 GPU，验证阶段会直接复用它，而不会从磁盘再加载一份副本。
|
||||
|
||||
**Wan2.2 T2V Reward LoRA 训练示例（单一 HPS v2.1 reward，FSDP）**：
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# NCCL_IB_DISABLE=1 与 NCCL_P2P_DISABLE=1 用于没有 RDMA 的多机环境。
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
export NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False \
|
||||
scripts/wan2.2/train_reward_lora.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=1 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=50 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_wan2.2_reward_lora" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--enable_bucket \
|
||||
--boundary_type="high" \
|
||||
--rank=128 \
|
||||
--network_alpha=64 \
|
||||
--target_name="q,k,v,ffn.0,ffn.2" \
|
||||
--train_mode="normal" \
|
||||
--guidance_scale=5.0 \
|
||||
--num_inference_steps=40 \
|
||||
--num_decoded_latents=3 \
|
||||
--reward_fn="HPSReward" \
|
||||
--reward_fn_kwargs='{"version": "v2.1"}' \
|
||||
--backprop \
|
||||
--backprop_strategy="tail" \
|
||||
--backprop_num_steps=1 \
|
||||
--low_vram
|
||||
```
|
||||
|
||||
> 本目录下的 `train_reward_lora.sh` 提供了一份可直接运行的模板。请根据你的环境调整 `MODEL_NAME`、数据集路径、`--reward_fn` 与 `--boundary_type`。
|
||||
|
||||
### 3.3 Reward 专用参数解析
|
||||
|
||||
除标准 LoRA 参数外（见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)），reward 训练脚本还新增以下参数：
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|-----------|-------------|---------------|
|
||||
| `--reward_fn` | Reward 模型类名，可用**逗号分隔的列表**启用多 reward 训练；各 reward 的可微 loss 以加权和合并。可选：`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`、`VideoAlignReward`。 | `"HPSv3Reward,VideoAlignReward"` |
|
||||
| `--reward_fn_kwargs` | Reward 模型的关键词参数。单 reward 用**扁平 JSON dict**；多 reward 用**以 reward 名为 key 的嵌套 JSON dict**。 | `'{"VideoAlignReward": {"max_reward": 2.0}, "HPSv3Reward": {"max_reward": 12}}'` |
|
||||
| `--multi_reward_weights` | 多 reward 各 loss 权重的 JSON dict。会归一化到和为 1；多 reward 时若省略则使用等权重；单 reward 时忽略。 | `'{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}'` |
|
||||
| `--num_inference_steps` | 可微 rollout 与验证中的去噪步数。 | 40 |
|
||||
| `--num_decoded_latents` | 解码回像素以供 reward 评估的潜帧数。分块因果 VAE 受梯度检查点限制，除非设置 `--use_1steps_vae`，否则要求 `<= 3`。 | 3 |
|
||||
| `--num_reward_frames` | 随机选取该数量的解码帧用于 reward 评估。`None` = 使用全部。仅对**图像** reward 生效（见下）。 | 2 |
|
||||
| `--use_1steps_vae` | 使用单段式 Wan2.2 VAE（`AutoencoderKLWan3_8_1Steps`）做可微解码；数值等价，一次前向解码整个 latent，干净地按解码块应用梯度检查点，并解除 `<= 3` 的限制。仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。 | - |
|
||||
| `--backprop` | 启用 reward 反向传播训练模式。 | - |
|
||||
| `--backprop_strategy` | rollout 中哪些步保留梯度图：`last`、`tail`、`uniform`、`random`。 | `tail` |
|
||||
| `--backprop_num_steps` | backprop 步数，仅当 `--backprop_strategy` 为 `tail`/`uniform`/`random` 时生效。 | 1 |
|
||||
| `--first_frame_mse_weight` | inpaint 模式下首帧 latent MSE 辅助 loss 的权重。 | 1.0 |
|
||||
|
||||
**图像 reward 与视频 reward（`is_image_reward`）**：
|
||||
- **图像 reward**（`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`）逐帧打分，因此帧抽稀（`--num_reward_frames`）是安全的。
|
||||
- **视频 reward**（`VideoAlignReward`，`is_image_reward=False`）需要完整时序序列来评判运动质量；若压缩帧会把视频退化为近乎静态的输入并破坏运动维度。因此这类 reward 保留全部解码帧（数量受 `--num_decoded_latents` 约束）。
|
||||
- 在 inpaint 模式（`train_mode="inpaint"`）下，所有 reward 在评估前都会丢弃首帧。
|
||||
|
||||
> **重要**：使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 这类视频 reward 必须在 `--reward_fn_kwargs` 中设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会彻底切断梯度链。
|
||||
|
||||
### 3.4 单 Reward 与多 Reward 训练
|
||||
|
||||
单 reward 使用 `--reward_fn="HPSv3Reward"` 配一个扁平的 `--reward_fn_kwargs`。若要组合多个目标，传入逗号分隔的列表、嵌套的 kwargs dict，以及（可选的）权重。每个 reward 的 loss 与 reward 值会分开记录（如 `train/loss_HPSv3Reward`、`train/reward_VideoAlignReward`），而实际优化的 `loss = Σ wᵢ · lossᵢ`。
|
||||
|
||||
**多 reward 示例（HPSv3 + VideoAlign，T2V）**——最强的对齐配置，因为视频 reward 能补上图像 reward 无法覆盖的时序（运动/一致性）信号：
|
||||
|
||||
```bash
|
||||
--num_decoded_latents=13 \
|
||||
--num_reward_frames=2 \
|
||||
--use_1steps_vae \
|
||||
--reward_fn="VideoAlignReward,HPSv3Reward" \
|
||||
--reward_fn_kwargs='{"VideoAlignReward": {"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "fps": 16, "reward_dim": "TA+MQ", "differentiable": true, "use_norm": true, "max_reward": 2.0, "loss_scale": 1.0}, "HPSv3Reward": {"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}}' \
|
||||
--multi_reward_weights='{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}' \
|
||||
--backprop \
|
||||
--backprop_strategy="tail" \
|
||||
--backprop_num_steps=1
|
||||
```
|
||||
|
||||
### 3.5 可微解码与 `--use_1steps_vae`
|
||||
|
||||
默认的分块因果 VAE（`AutoencoderKLWan3_8`）按块应用梯度检查点，从而把可微解码限制在**前 3 个潜帧**（`--num_decoded_latents <= 3`）。不设置 `--use_1steps_vae` 而使用 `--num_decoded_latents > 3` 会抛出 `ValueError`。
|
||||
|
||||
启用 `--use_1steps_vae` 可切换到单段式变体，它在一次前向中解码整个 latent、按解码块干净地应用梯度检查点，并解除该限制——当视频 reward 需要更多时序上下文时很有用。它数值等价，且仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。
|
||||
|
||||
### 3.6 训练验证
|
||||
|
||||
你可以配置验证参数，在训练过程中周期性生成测试视频，以监控对齐进度。在 FSDP 下，`log_validation` 里的 pipeline 前向是一次跨 rank 的集合通信，因此**所有进程都必须进入验证调用**（它刻意**没有**包在 `is_main_process` 里）；可训练专家已带上 LoRA，冻结专家则被复用而非重新加载。
|
||||
|
||||
| 参数 | 说明 | 默认值 |
|
||||
|-----------|-------------|---------------|
|
||||
| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
|
||||
| `--validation_epochs` | 每 N 个 epoch 运行一次验证 | 5 |
|
||||
| `--validation_prompts` | 用于验证视频生成的 prompt | None |
|
||||
| `--validation_paths` | I2V 验证的参考图路径（仅 inpaint 模式） | None |
|
||||
|
||||
**验证示例**：
|
||||
|
||||
```bash
|
||||
--validation_prompts="A stylish woman strolls down a bustling Tokyo street at night, warm neon glow and animated city signage reflecting on the wet pavement; she wears a black leather jacket over a red dress, confident casual demeanor, medium shot, cinematic." \
|
||||
--validation_steps=50 \
|
||||
--validation_epochs=500
|
||||
```
|
||||
|
||||
> 验证视频会保存到 `output_dir` 目录。若做 T2V 验证（无参考图），保持 `train_mode="normal"` 并省略 `--validation_paths`。
|
||||
|
||||
---
|
||||
|
||||
## 四、推理测试
|
||||
|
||||
训练好的 reward LoRA 就是一个普通 LoRA，加载方式与 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 完全一致。对于 A14B 双 Transformer 模型，在 [examples/wan2.2/predict_t2v.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.2/predict_t2v.py)（或 `predict_i2v.py`）中，为 low-noise reward LoRA 设置 `lora_path` / `lora_weight`，为 high-noise reward LoRA 设置 `lora_high_path` / `lora_high_weight`：
|
||||
|
||||
```python
|
||||
# 训练好的 reward LoRA 权重路径，例如
|
||||
# "output_dir_wan2.2_reward_lora/checkpoint-xxx/diffusion_pytorch_model.safetensors"
|
||||
lora_path = None
|
||||
lora_high_path = None
|
||||
lora_weight = 0.55
|
||||
lora_high_weight = 0.55
|
||||
```
|
||||
|
||||
单卡推理运行：
|
||||
|
||||
```bash
|
||||
python examples/wan2.2/predict_t2v.py
|
||||
```
|
||||
|
||||
显存管理模式与多卡并行推理的详细说明，请参见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 的推理测试章节。
|
||||
|
||||
---
|
||||
|
||||
## 五、局限性
|
||||
|
||||
1. 训练到一定程度后，reward 持续上升，但生成视频的质量不再进一步提升。模型会取巧地学到一些捷径（在背景里添加伪影，即对抗补丁）来抬高 reward。
|
||||
2. 目前仍缺乏适合视频生成的偏好模型。直接使用图像偏好模型无法沿时序维度评估偏好（如运动性与一致性），并且会降低生成视频的运动性。可以仅用解码视频的第一帧来计算 reward 来部分缓解，但影响依然存在——这也正是对于运动敏感的对齐，推荐使用视频 reward（如 `VideoAlignReward`）并配合 `--use_1steps_vae` 与更大的 `--num_decoded_latents`、或采用多 reward 组合的原因。
|
||||
|
||||
---
|
||||
|
||||
## 六、更多资源
|
||||
|
||||
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
|
||||
- **标准 LoRA 训练**：[README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)
|
||||
- **英文版**：[README_TRAIN_REWARD.md](./README_TRAIN_REWARD.md)
|
||||
|
||||
---
|
||||
|
||||
## 参考文献
|
||||
<ol>
|
||||
<li id="ref1">Clark, Kevin, et al. "Directly fine-tuning diffusion models on differentiable rewards.". In ICLR 2024.</li>
|
||||
<li id="ref2">Prabhudesai, Mihir, et al. "Aligning text-to-image diffusion models with reward backpropagation." arXiv preprint arXiv:2310.03739 (2023).</li>
|
||||
</ol>
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,50 @@
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# NCCL_IB_DISABLE=1 与 NCCL_P2P_DISABLE=1 用于没有 RDMA 的多机环境。
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
export NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False \
|
||||
scripts/wan2.2/train_reward_lora.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=1 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=50 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_wan2.2_reward_lora" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--enable_bucket \
|
||||
--boundary_type="high" \
|
||||
--rank=128 \
|
||||
--network_alpha=64 \
|
||||
--target_name="q,k,v,ffn.0,ffn.2" \
|
||||
--train_mode="normal" \
|
||||
--guidance_scale=5.0 \
|
||||
--num_inference_steps=40 \
|
||||
--num_decoded_latents=3 \
|
||||
--reward_fn="HPSReward" \
|
||||
--reward_fn_kwargs='{"version": "v2.1"}' \
|
||||
--backprop \
|
||||
--backprop_strategy="tail" \
|
||||
--backprop_num_steps=1 \
|
||||
--low_vram
|
||||
@@ -1,40 +1,380 @@
|
||||
# Wan2.2-Fun-Reward-LoRAs
|
||||
## Introduction
|
||||
We explore the Reward Backpropagation technique <sup>[1](#ref1) [2](#ref2)</sup> to optimized the generated videos by [Wan2.2-Fun](https://github.com/aigc-apps/VideoX-Fun) for better alignment with human preferences.
|
||||
We provide the following pre-trained models (i.e. LoRAs) along with [the training script](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/wan2.2_fun/train_reward_lora.py). You can use these LoRAs to enhance the corresponding base model as a plug-in or train your own reward LoRA.
|
||||
# Wan2.2 Fun Reward LoRA 训练指南
|
||||
|
||||
For more details, please refer to our [GitHub repo](https://github.com/aigc-apps/VideoX-Fun).
|
||||
本文档提供 Wan2.2 Fun **Reward 反向传播（Reward Backpropagation）** 微调的完整流程，即用 LoRA 优化生成视频，使其更符合人类偏好。内容涵盖环境配置、数据准备、Reward 专用训练参数（单/多 reward、可微 VAE 解码、backprop 策略）、训练验证与推理测试。
|
||||
|
||||
| Name | Base Model | Reward Model | Hugging Face | Description |
|
||||
> **说明**：Wan2.2 Fun 是一个基于 Wan2.2 架构的视频修复模型。Wan2.2 采用双 Transformer 架构（高噪声/低噪声模型），5B 版本使用单 Transformer 架构。Reward LoRA 训练会跑一段可微的 rollout（一条简短的去噪轨迹），并把 reward 模型的 loss 反传到可训练的专家上，因此比标准 LoRA 微调更耗显存。本指南支持 A14B 与 5B 两种变体。
|
||||
|
||||
我们采用 Reward Backpropagation 技术 <sup>[1](#ref1) [2](#ref2)</sup> 来优化 [Wan2.2-Fun](https://github.com/aigc-apps/VideoX-Fun) 生成的视频，使其更好地对齐人类偏好。我们随[训练脚本](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/wan2.2_fun/train_reward_lora.py)一并提供以下预训练模型（即 LoRA）。你可以用这些 LoRA 作为插件增强对应的基座模型，也可以训练自己的 reward LoRA。
|
||||
|
||||
| 名称 | 基座模型 | Reward 模型 | Hugging Face | 说明 |
|
||||
|--|--|--|--|--|
|
||||
| Wan2.2-Fun-A14B-InP-high-noise-HPS2.1.safetensors | [Wan2.2-Fun-A14B-InP (high noise)](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/high_noise_model) | [HPS v2.1](https://github.com/tgxs002/HPSv2) | [🤗Link](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-high-noise-HPS2.1.safetensors) | Official HPS v2.1 reward LoRA (`rank=128` and `network_alpha=64`) for Wan2.2-Fun-A14B-InP (high noise). It is trained with a batch size of 8 for 5,000 steps.|
|
||||
| Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors | [Wan2.2-Fun-A14B-InP (low noise)](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/low_noise_model) | [MPS](https://github.com/Kwai-Kolors/MPS) | [🤗Link](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors) | Official HPS v2.1 reward LoRA (`rank=128` and `network_alpha=64`) for Wan2.2-Fun-A14B-InP (low noise). It is trained with a batch size of 8 for 2,700 steps.|
|
||||
| Wan2.2-Fun-A14B-InP-high-noise-MPS.safetensors | [Wan2.2-Fun-A14B-InP (high noise)](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/high_noise_model) | [HPS v2.1](https://github.com/tgxs002/HPSv2) | [🤗Link](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-high-noise-MPS.safetensors) | Official MPS reward LoRA (`rank=128` and `network_alpha=64`) for Wan2.2-Fun-A14B-InP (high noise). It is trained with a batch size of 8 for 5,000 steps.|
|
||||
| Wan2.2-Fun-A14B-InP-low-noise-MPS.safetensors | [Wan2.2-Fun-A14B-InP (low noise)](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/low_noise_model) | [MPS](https://github.com/Kwai-Kolors/MPS) | [🤗Link](https://huggingface.co/alibaba-pai/Wan2.1-Fun-Reward-LoRAs/resolve/main/Wan2.1-Fun-14B-InP-MPS.safetensors) | Official MPS reward LoRA (`rank=128` and `network_alpha=64`) for Wan2.2-Fun-A14B-InP (low noise). It is trained with a batch size of 8 for xxx steps.|
|
||||
| Wan2.2-Fun-A14B-InP-high-noise-HPS2.1.safetensors | [Wan2.2-Fun-A14B-InP（high noise）](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/high_noise_model) | [HPS v2.1](https://github.com/tgxs002/HPSv2) | [🤗链接](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-high-noise-HPS2.1.safetensors) | 面向 Wan2.2-Fun-A14B-InP（high noise）的官方 HPS v2.1 reward LoRA（`rank=128`、`network_alpha=64`）。 |
|
||||
| Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors | [Wan2.2-Fun-A14B-InP（low noise）](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/low_noise_model) | [HPS v2.1](https://github.com/tgxs002/HPSv2) | [🤗链接](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors) | 面向 Wan2.2-Fun-A14B-InP（low noise）的官方 HPS v2.1 reward LoRA（`rank=128`、`network_alpha=64`）。 |
|
||||
| Wan2.2-Fun-A14B-InP-high-noise-MPS.safetensors | [Wan2.2-Fun-A14B-InP（high noise）](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/high_noise_model) | [MPS](https://github.com/Kwai-Kolors/MPS) | [🤗链接](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-high-noise-MPS.safetensors) | 面向 Wan2.2-Fun-A14B-InP（high noise）的官方 MPS reward LoRA（`rank=128`、`network_alpha=64`）。 |
|
||||
| Wan2.2-Fun-A14B-InP-low-noise-MPS.safetensors | [Wan2.2-Fun-A14B-InP（low noise）](https://huggingface.co/alibaba-pai/Wan2.2-Fun-A14B-InP/tree/main/low_noise_model) | [MPS](https://github.com/Kwai-Kolors/MPS) | [🤗链接](https://huggingface.co/alibaba-pai/Wan2.1-Fun-Reward-LoRAs/resolve/main/Wan2.1-Fun-14B-InP-MPS.safetensors) | 面向 Wan2.2-Fun-A14B-InP（low noise）的官方 MPS reward LoRA（`rank=128`、`network_alpha=64`）。 |
|
||||
|
||||
> [!NOTE]
|
||||
> We found that, MPS reward LoRA for the low-noise model converges significantly more slowly than on the other models, and may not deliver satisfactory results. Therefore, for the low-noise model, we recommend using HPSv2.1 reward LoRA.
|
||||
> 我们发现，low-noise 模型的 MPS reward LoRA 收敛速度明显慢于其他模型，且效果可能不理想。因此对于 low-noise 模型，我们推荐使用 HPS v2.1 reward LoRA。
|
||||
|
||||
## Demo
|
||||
Please refer to [here](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs#demo).
|
||||
## 演示
|
||||
请参考[此处](https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs#demo)。
|
||||
|
||||
## Quick Start
|
||||
Set `lora_path` along with `lora_weight` for the low noise reward LoRA, while specifying `lora_high_path` and `lora_high_weight` for high noise reward LoRA in [examples/wan2.2_fun/predict_t2v.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.1_fun/predict_t2v.py).
|
||||
---
|
||||
|
||||
## Training
|
||||
The training code is based on [train_lora.py](./train_lora.py). We provide a shell script to train the HPS v2.1 reward LoRA for the low noise model of Wan2.2-Fun-A14B-InP, which can be trained on a single 8*A100 node with 80GB VRAM. To train reward LoRA for the high noise model, Deepspeed Zero3 with CPU offload is required.
|
||||
## 目录
|
||||
- [一、环境配置](#一环境配置)
|
||||
- [二、数据准备](#二数据准备)
|
||||
- [2.1 快速测试数据集](#21-快速测试数据集)
|
||||
- [2.2 Reward 模型权重](#22-reward-模型权重)
|
||||
- [2.3 数据集结构与 metadata.json](#23-数据集结构与-metadatajson)
|
||||
- [三、Reward LoRA 训练](#三reward-lora-训练)
|
||||
- [3.1 下载预训练基座模型](#31-下载预训练基座模型)
|
||||
- [3.2 快速开始（FSDP）](#32-快速开始fsdp)
|
||||
- [3.3 Reward 专用参数解析](#33-reward-专用参数解析)
|
||||
- [3.4 单 Reward 与多 Reward 训练](#34-单-reward-与多-reward-训练)
|
||||
- [3.5 可微解码与 `--use_1steps_vae`](#35-可微解码与---use_1steps_vae)
|
||||
- [3.6 训练验证](#36-训练验证)
|
||||
- [四、推理测试](#四推理测试)
|
||||
- [五、局限性](#五局限性)
|
||||
- [六、更多资源](#六更多资源)
|
||||
- [参考文献](#参考文献)
|
||||
|
||||
Please refer to [Setup](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/cogvideox_fun/README_TRAIN_REWARD.md#setup) and [Important Args](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/cogvideox_fun/README_TRAIN_REWARD.md#important-args) before training.
|
||||
---
|
||||
|
||||
## 一、环境配置
|
||||
|
||||
## Limitations
|
||||
1. We observe after training to a certain extent, the reward continues to increase, but the quality of the generated videos does not further improve.
|
||||
The model trickly learns some shortcuts (by adding artifacts in the background, i.e., adversarial patches) to increase the reward.
|
||||
2. Currently, there is still a lack of suitable preference models for video generation. Directly using image preference models cannot
|
||||
evaluate preferences along the temporal dimension (such as dynamism and consistency). Further more, We find using image preference models leads to a decrease
|
||||
in the dynamism of generated videos. Although this can be mitigated by computing the reward using only the first frame of the decoded video, the impact still persists.
|
||||
**方式 1：使用 requirements.txt**
|
||||
|
||||
## Reference
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**方式 2：手动安装依赖**
|
||||
|
||||
```bash
|
||||
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
|
||||
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
|
||||
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
|
||||
pip install yunchang xfuser modelscope openpyxl
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
pip install deepspeed==0.17.0 numpy==1.26.4
|
||||
```
|
||||
|
||||
> 部分视频 reward 模型（如 `VideoAlignReward`、`HPSv3Reward`）基于 Qwen2-VL 骨干，需要 `transformers==4.51.3`。使用这些 reward 时请安装对应版本。
|
||||
|
||||
**方式 3：使用 docker**
|
||||
|
||||
使用 docker 的情况下，请保证机器中已经正确安装显卡驱动与 CUDA 环境，然后以此执行以下命令：
|
||||
|
||||
```
|
||||
# pull image
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
|
||||
# enter image
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 二、数据准备
|
||||
|
||||
### 2.1 快速测试数据集
|
||||
|
||||
Reward LoRA 训练从文本 prompt 出发跑可微 rollout，因此标准的 T2V reward 路径只需一个**纯 prompt 数据集**即可。对于 I2V/inpaint 路径，则使用视频数据集来提供参考条件。
|
||||
|
||||
```bash
|
||||
# 下载官方示例数据集
|
||||
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
|
||||
```
|
||||
|
||||
### 2.2 Reward 模型权重
|
||||
|
||||
不同 reward 的骨干/checkpoint 来源不同：
|
||||
|
||||
- **`HPSReward`（HPS v2.1）与 `MPSReward`**：打分头很小，会在**首次运行时自动下载**到 torch hub 缓存（`~/.cache/torch/hub`）；其中 **HPS 还需给 `hpsv2` 内置的 open_clip 补两个特殊配置文件**（见下），MPS 无需额外准备；
|
||||
- **`VideoAlignReward`（VideoReward）与 `HPSv3Reward`**：基于 Qwen2-VL 骨干，体积较大，**建议提前下载到本地**，再用 `--reward_fn_kwargs` 指向本地路径。
|
||||
|
||||
各 reward 的依赖与来源汇总如下（`--reward_fn` 取值见 [3.3](#33-reward-专用参数解析)）：
|
||||
|
||||
| Reward | 骨干 / 依赖 | 打分头 / checkpoint | 手动下载 |
|
||||
|--|--|--|--|
|
||||
| `HPSReward`（`version="v2.1"`） | open_clip `ViT-H-14`，需 `pip install hpsv2` + `factory.py` 补丁与 BPE 词表 | `HPS_v2.1_compressed.pt`（自动下载） | 否，可用 `model_path` 覆盖 |
|
||||
| `MPSReward` | `laion/CLIP-ViT-H-14-laion2B-s32B-b79K`（config/tokenizer 自动下载），MPS 代码已内置于 `videox_fun/reward/MPS/` | `MPS_overall.pth`（自动下载） | 否，可用 `model_path` 覆盖 |
|
||||
| `VideoAlignReward` | `Qwen/Qwen2-VL-2B-Instruct` | `KlingTeam/VideoReward`（含 `model_config.json` 与权重） | **是**，`model_path` 必填 |
|
||||
| `HPSv3Reward` | `Qwen/Qwen2-VL-7B-Instruct` | `MizzenAI/HPSv3` 的 `HPSv3.safetensors` | 建议，缺省时自动从 HF 下载 |
|
||||
|
||||
**HPS v2.1（`HPSReward`）**——除了 `pip install hpsv2`，还要给 `hpsv2` 内置的 open_clip 补两个特殊配置文件：`create_model_and_transforms("ViT-H-14", ...)` 用到了 HPS 分支新增的参数，需要打过补丁的 `factory.py`；`get_tokenizer("ViT-H-14")` 需要 BPE 词表 `bpe_simple_vocab_16e6.txt.gz`。缺任一都会在加载时报错。`ViT-H-14` 骨干与 `HPS_v2.1_compressed.pt` 打分头会在首次运行时自动下载到 torch hub 缓存：
|
||||
|
||||
```bash
|
||||
pip install hpsv2
|
||||
# 覆盖 hpsv2 内置 open_clip 的 factory.py 补丁，并补上 BPE 词表
|
||||
site_packages=$(python -c "import site; print(site.getsitepackages()[0])")
|
||||
wget -O $site_packages/hpsv2/src/open_clip/factory.py https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/package/patches/hpsv2_src_open_clip_factory_patches.py
|
||||
wget -O $site_packages/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz https://github.com/tgxs002/HPSv2/raw/refs/heads/master/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="HPSReward" \
|
||||
--reward_fn_kwargs='{"version": "v2.1"}'
|
||||
```
|
||||
|
||||
**MPS（`MPSReward`）**——MPS 代码已内置于 `videox_fun/reward/MPS/`，无需打补丁；`MPS_overall.pth` 打分头会自动从 OSS 下载，`laion/CLIP-ViT-H-14-laion2B-s32B-b79K` 的 config 与 tokenizer 会自动从 HuggingFace 下载（若无法访问 HF，可用 `processor_name_or_path` 指向本地镜像）：
|
||||
|
||||
```bash
|
||||
--reward_fn="MPSReward" \
|
||||
--reward_fn_kwargs='{"max_reward": 1, "loss_scale": 1}'
|
||||
```
|
||||
|
||||
**VideoReward（`VideoAlignReward`）**——需 Qwen2-VL-2B 骨干 + VideoReward 打分头目录（内含 `model_config.json` 与权重），二者都要提前下载，且 `model_path` 必填：
|
||||
|
||||
```bash
|
||||
modelscope download --model Qwen/Qwen2-VL-2B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-2B-Instruct
|
||||
huggingface-cli download KlingTeam/VideoReward --local-dir models/Personalized_Model/VideoReward
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="VideoAlignReward" \
|
||||
--reward_fn_kwargs='{"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "reward_dim": "TA+MQ", "fps": 16, "use_norm": true, "differentiable": true, "max_reward": 2.0}'
|
||||
```
|
||||
|
||||
**HPSv3（`HPSv3Reward`）**——需 Qwen2-VL-7B 骨干 + `HPSv3.safetensors`；不指定 `checkpoint_path` 时会自动从 HF `MizzenAI/HPSv3` 下载，建议提前落盘：
|
||||
|
||||
```bash
|
||||
modelscope download --model Qwen/Qwen2-VL-7B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-7B-Instruct
|
||||
modelscope download --model MizzenAI/HPSv3 --local_dir models/Personalized_Model/HPSv3
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="HPSv3Reward" \
|
||||
--reward_fn_kwargs='{"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}'
|
||||
```
|
||||
|
||||
> 使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 必须设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会切断梯度链（见 [3.3](#33-reward-专用参数解析)）。
|
||||
|
||||
### 2.3 数据集结构与 metadata.json
|
||||
|
||||
```
|
||||
📦 datasets/
|
||||
├── 📂 my_dataset/
|
||||
│ ├── 📂 train/
|
||||
│ │ ├── 📄 video001.mp4
|
||||
│ │ ├── 📄 video002.mp4
|
||||
│ │ └── 📄 ...
|
||||
│ └── 📄 metadata.json
|
||||
```
|
||||
|
||||
**相对路径格式**（示例）：
|
||||
```json
|
||||
[
|
||||
{
|
||||
"file_path": "train/video001.mp4",
|
||||
"text": "A beautiful sunset over the ocean, golden hour lighting",
|
||||
"type": "video",
|
||||
"width": 1024,
|
||||
"height": 1024
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
**关键字段说明**：
|
||||
- `file_path`：视频路径（相对路径或绝对路径）。
|
||||
- `text`：视频描述（英文 prompt）。
|
||||
- `type`：数据类型，固定为 `"video"`。
|
||||
- `width` / `height`：视频宽高（**强烈建议提供**，用于分桶训练）。可用 `scripts/process_json_add_width_and_height.py` 提取这两个字段。
|
||||
|
||||
**相对路径与绝对路径**：
|
||||
- 相对路径：`export DATASET_NAME="datasets/X-Fun-Videos-Demo/"` 与 `export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"`。
|
||||
- 绝对路径：`export DATASET_NAME=""` 与 `export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"`。
|
||||
|
||||
---
|
||||
|
||||
## 三、Reward LoRA 训练
|
||||
|
||||
### 3.1 下载预训练基座模型
|
||||
|
||||
```bash
|
||||
# 创建模型目录
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# 下载 Wan2.2 Fun 官方权重
|
||||
# A14B 模型（双 Transformer 架构）
|
||||
modelscope download --model PAI/Wan2.2-Fun-A14B-InP --local_dir models/Diffusion_Transformer/Wan2.2-Fun-A14B-InP
|
||||
# 或 5B 模型（单 Transformer 架构）
|
||||
# modelscope download --model PAI/Wan2.2-Fun-5B-InP --local_dir models/Diffusion_Transformer/Wan2.2-Fun-5B-InP
|
||||
```
|
||||
|
||||
### 3.2 快速开始（FSDP）
|
||||
|
||||
Reward LoRA 训练会把完整的去噪 rollout 与 reward 模型的前向计算图保持在显存中，因此推荐使用 FSDP（`FULL_SHARD`）。当 `--boundary_type="high"` 时，*另一个*（冻结的）专家在训练期已被分片并常驻 GPU，验证阶段会直接复用它，而不会从磁盘再加载一份副本。
|
||||
|
||||
**Wan2.2 Fun Reward LoRA 训练示例（单一 HPS v2.1 reward，FSDP）**：
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-Fun-A14B-InP"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# NCCL_IB_DISABLE=1 与 NCCL_P2P_DISABLE=1 用于没有 RDMA 的多机环境。
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
export NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False \
|
||||
scripts/wan2.2_fun/train_reward_lora.py \
|
||||
--config_path="config/wan2.2/wan_civitai_i2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=50 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_wan2.2_fun_reward_lora" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--enable_bucket \
|
||||
--boundary_type="low" \
|
||||
--rank=128 \
|
||||
--network_alpha=64 \
|
||||
--target_name="q,k,v,ffn.0,ffn.2" \
|
||||
--num_inference_steps=25 \
|
||||
--num_decoded_latents=3 \
|
||||
--reward_fn="HPSReward" \
|
||||
--reward_fn_kwargs='{"version": "v2.1"}' \
|
||||
--backprop \
|
||||
--backprop_strategy="tail" \
|
||||
--backprop_num_steps=1 \
|
||||
--train_mode="inpaint" \
|
||||
--low_vram
|
||||
```
|
||||
|
||||
> 本目录下的 `train_reward_lora.sh` 提供了一份可直接运行的模板。请根据你的环境调整 `MODEL_NAME`、数据集路径、`--reward_fn` 与 `--boundary_type`。
|
||||
|
||||
### 3.3 Reward 专用参数解析
|
||||
|
||||
除标准 LoRA 参数外（见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)），reward 训练脚本还新增以下参数：
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|-----------|-------------|---------------|
|
||||
| `--reward_fn` | Reward 模型类名，可用**逗号分隔的列表**启用多 reward 训练；各 reward 的可微 loss 以加权和合并。可选：`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`、`VideoAlignReward`。 | `"HPSv3Reward,VideoAlignReward"` |
|
||||
| `--reward_fn_kwargs` | Reward 模型的关键词参数。单 reward 用**扁平 JSON dict**；多 reward 用**以 reward 名为 key 的嵌套 JSON dict**。 | `'{"VideoAlignReward": {"max_reward": 2.0}, "HPSv3Reward": {"max_reward": 12}}'` |
|
||||
| `--multi_reward_weights` | 多 reward 各 loss 权重的 JSON dict。会归一化到和为 1；多 reward 时若省略则使用等权重；单 reward 时忽略。 | `'{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}'` |
|
||||
| `--num_inference_steps` | 可微 rollout 与验证中的去噪步数。 | 25 |
|
||||
| `--num_decoded_latents` | 解码回像素以供 reward 评估的潜帧数。分块因果 VAE 受梯度检查点限制，除非设置 `--use_1steps_vae`，否则要求 `<= 3`。 | 3 |
|
||||
| `--num_reward_frames` | 随机选取该数量的解码帧用于 reward 评估。`None` = 使用全部。仅对**图像** reward 生效（见下）。 | 2 |
|
||||
| `--use_1steps_vae` | 使用单段式 Wan2.2 VAE（`AutoencoderKLWan3_8_1Steps`）做可微解码；数值等价，一次前向解码整个 latent，干净地按解码块应用梯度检查点，并解除 `<= 3` 的限制。仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。 | - |
|
||||
| `--backprop` | 启用 reward 反向传播训练模式。 | - |
|
||||
| `--backprop_strategy` | rollout 中哪些步保留梯度图：`last`、`tail`、`uniform`、`random`。 | `tail` |
|
||||
| `--backprop_num_steps` | backprop 步数，仅当 `--backprop_strategy` 为 `tail`/`uniform`/`random` 时生效。 | 1 |
|
||||
| `--first_frame_mse_weight` | inpaint 模式下首帧 latent MSE 辅助 loss 的权重。 | 1.0 |
|
||||
| `--force_t2v` | 当 `train_mode` 为 `"inpaint"`（I2V）或 `"ti2v"` 时，按纯 T2V 训练（条件置零、使用纯 prompt 数据集）。 | - |
|
||||
|
||||
**图像 reward 与视频 reward（`is_image_reward`）**：
|
||||
- **图像 reward**（`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`）逐帧打分，因此帧抽稀（`--num_reward_frames`）是安全的。
|
||||
- **视频 reward**（`VideoAlignReward`，`is_image_reward=False`）需要完整时序序列来评判运动质量；若压缩帧会把视频退化为近乎静态的输入并破坏运动维度。因此这类 reward 保留全部解码帧（数量受 `--num_decoded_latents` 约束）。
|
||||
- 在 inpaint 模式下，所有 reward 在评估前都会丢弃首帧（`--force_t2v` 时除外）。
|
||||
|
||||
> **重要**：使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 这类视频 reward 必须在 `--reward_fn_kwargs` 中设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会彻底切断梯度链。
|
||||
|
||||
### 3.4 单 Reward 与多 Reward 训练
|
||||
|
||||
单 reward 使用 `--reward_fn="HPSReward"` 配一个扁平的 `--reward_fn_kwargs`。若要组合多个目标，传入逗号分隔的列表、嵌套的 kwargs dict，以及（可选的）权重。每个 reward 的 loss 与 reward 值会分开记录（如 `train/loss_HPSv3Reward`、`train/reward_VideoAlignReward`），而实际优化的 `loss = Σ wᵢ · lossᵢ`。
|
||||
|
||||
**多 reward 示例（HPSv3 + VideoAlign，T2V）**：
|
||||
|
||||
```bash
|
||||
--num_decoded_latents=13 \
|
||||
--num_reward_frames=2 \
|
||||
--use_1steps_vae \
|
||||
--reward_fn="VideoAlignReward,HPSv3Reward" \
|
||||
--reward_fn_kwargs='{"VideoAlignReward": {"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "fps": 16, "reward_dim": "TA+MQ", "differentiable": true, "use_norm": true, "max_reward": 2.0, "loss_scale": 1.0}, "HPSv3Reward": {"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}}' \
|
||||
--multi_reward_weights='{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}' \
|
||||
--backprop \
|
||||
--backprop_strategy="tail" \
|
||||
--backprop_num_steps=1
|
||||
```
|
||||
|
||||
### 3.5 可微解码与 `--use_1steps_vae`
|
||||
|
||||
默认的分块因果 VAE（`AutoencoderKLWan3_8`）按块应用梯度检查点，从而把可微解码限制在**前 3 个潜帧**（`--num_decoded_latents <= 3`）。不设置 `--use_1steps_vae` 而使用 `--num_decoded_latents > 3` 会抛出 `ValueError`。
|
||||
|
||||
启用 `--use_1steps_vae` 可切换到单段式变体，它在一次前向中解码整个 latent、按解码块干净地应用梯度检查点，并解除该限制——当视频 reward 需要更多时序上下文时很有用。它数值等价，且仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。
|
||||
|
||||
### 3.6 训练验证
|
||||
|
||||
你可以配置验证参数，在训练过程中周期性生成测试视频，以监控对齐进度。在 FSDP 下，`log_validation` 里的 pipeline 前向是一次跨 rank 的集合通信，因此**所有进程都必须进入验证调用**（它刻意**没有**包在 `is_main_process` 里）；可训练专家已带上 LoRA，冻结专家则被复用而非重新加载。
|
||||
|
||||
| 参数 | 说明 | 默认值 |
|
||||
|-----------|-------------|---------------|
|
||||
| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
|
||||
| `--validation_epochs` | 每 N 个 epoch 运行一次验证 | 5 |
|
||||
| `--validation_prompts` | 用于验证视频生成的 prompt | None |
|
||||
| `--validation_paths` | I2V 验证的参考图路径（仅 inpaint 模式） | None |
|
||||
|
||||
**验证示例**：
|
||||
|
||||
```bash
|
||||
--validation_prompts="A brown dog shaking its head, sitting on a light-colored sofa in a cozy room." \
|
||||
--validation_steps=50 \
|
||||
--validation_epochs=500
|
||||
```
|
||||
|
||||
> 验证视频会保存到 `output_dir` 目录。若做 T2V 验证（无参考图），请设置 `train_mode="normal"`（或加上 `--force_t2v`）并省略 `--validation_paths`。
|
||||
|
||||
---
|
||||
|
||||
## 四、推理测试
|
||||
|
||||
训练好的 reward LoRA 就是一个普通 LoRA，加载方式与 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 完全一致。对于 A14B 双 Transformer 模型，在 [examples/wan2.2_fun/predict_t2v.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.2_fun/predict_t2v.py)（或 `predict_i2v.py`）中，为 low-noise reward LoRA 设置 `lora_path` / `lora_weight`，为 high-noise reward LoRA 设置 `lora_high_path` / `lora_high_weight`：
|
||||
|
||||
```python
|
||||
# 训练好的 reward LoRA 权重路径，例如
|
||||
# "output_dir_wan2.2_fun_reward_lora/checkpoint-xxx/diffusion_pytorch_model.safetensors"
|
||||
lora_path = None
|
||||
lora_high_path = None
|
||||
lora_weight = 0.55
|
||||
lora_high_weight = 0.55
|
||||
```
|
||||
|
||||
单卡推理运行：
|
||||
|
||||
```bash
|
||||
python examples/wan2.2_fun/predict_t2v.py
|
||||
```
|
||||
|
||||
显存管理模式与多卡并行推理的详细说明，请参见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 的推理测试章节。
|
||||
|
||||
---
|
||||
|
||||
## 五、局限性
|
||||
|
||||
1. 训练到一定程度后，reward 持续上升，但生成视频的质量不再进一步提升。模型会取巧地学到一些捷径（在背景里添加伪影，即对抗补丁）来抬高 reward。
|
||||
2. 目前仍缺乏适合视频生成的偏好模型。直接使用图像偏好模型无法沿时序维度评估偏好（如运动性与一致性），并且会降低生成视频的运动性。可以仅用解码视频的第一帧来计算 reward 来部分缓解，但影响依然存在——这也正是对于运动敏感的对齐，推荐使用视频 reward（如 `VideoAlignReward`）并配合 `--use_1steps_vae` 与更大的 `--num_decoded_latents` 的原因。
|
||||
|
||||
---
|
||||
|
||||
## 六、更多资源
|
||||
|
||||
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
|
||||
- **标准 LoRA 训练**：[README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)
|
||||
- **英文版**：[README_TRAIN_REWARD.md](./README_TRAIN_REWARD.md)
|
||||
|
||||
---
|
||||
|
||||
## 参考文献
|
||||
<ol>
|
||||
<li id="ref1">Clark, Kevin, et al. "Directly fine-tuning diffusion models on differentiable rewards.". In ICLR 2024.</li>
|
||||
<li id="ref2">Prabhudesai, Mihir, et al. "Aligning text-to-image diffusion models with reward backpropagation." arXiv preprint arXiv:2310.03739 (2023).</li>
|
||||
|
||||
@@ -0,0 +1,374 @@
|
||||
# Wan2.2 Reward LoRA 训练指南
|
||||
|
||||
本文档提供 Wan2.2 **Reward 反向传播（Reward Backpropagation）** 微调的完整流程，即用 LoRA 优化生成视频，使其更符合人类偏好。内容涵盖环境配置、数据准备、Reward 专用训练参数（单/多 reward、可微 VAE 解码、backprop 策略）、训练验证与推理测试。
|
||||
|
||||
> **说明**：Wan2.2 是一个支持文本到视频（T2V）、图像到视频（I2V）、文本-图像到视频（TI2V）的视频生成模型。A14B 系列采用双 Transformer 架构（高噪声/低噪声模型），TI2V-5B 使用单 Transformer 架构。Reward LoRA 训练会跑一段可微的 rollout（一条简短的去噪轨迹），并把 reward 模型的 loss 反传到可训练的专家上，因此比标准 LoRA 微调更耗显存。本指南覆盖为任意 Wan2.2 变体训练自有 reward LoRA 的流程。
|
||||
|
||||
我们采用 Reward Backpropagation 技术 <sup>[1](#ref1) [2](#ref2)</sup> 来优化 [Wan2.2](https://github.com/aigc-apps/VideoX-Fun) 生成的视频，使其更好地对齐人类偏好。训练脚本为 [train_reward_lora.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/scripts/wan2.2/train_reward_lora.py)；你可以训练自己的 reward LoRA，并作为插件加载到对应的基座模型上。
|
||||
|
||||
> [!NOTE]
|
||||
> Reward 反向传播容易出现 reward-hacking：reward 持续上升，但感知质量可能停滞甚至退化（背景出现对抗伪影）。相比纯图像 reward，更推荐使用**视频 reward**（如 `VideoAlignReward`）或**多 reward 组合**（如 `HPSv3Reward` + `VideoAlignReward`），以缓解时序运动性的下降——见 [3.4](#34-单-reward-与多-reward-训练) 与 [五、局限性](#五局限性)。
|
||||
|
||||
---
|
||||
|
||||
## 目录
|
||||
- [一、环境配置](#一环境配置)
|
||||
- [二、数据准备](#二数据准备)
|
||||
- [2.1 快速测试数据集](#21-快速测试数据集)
|
||||
- [2.2 Reward 模型权重](#22-reward-模型权重)
|
||||
- [2.3 数据集结构与 metadata.json](#23-数据集结构与-metadatajson)
|
||||
- [三、Reward LoRA 训练](#三reward-lora-训练)
|
||||
- [3.1 下载预训练基座模型](#31-下载预训练基座模型)
|
||||
- [3.2 快速开始（FSDP）](#32-快速开始fsdp)
|
||||
- [3.3 Reward 专用参数解析](#33-reward-专用参数解析)
|
||||
- [3.4 单 Reward 与多 Reward 训练](#34-单-reward-与多-reward-训练)
|
||||
- [3.5 可微解码与 `--use_1steps_vae`](#35-可微解码与---use_1steps_vae)
|
||||
- [3.6 训练验证](#36-训练验证)
|
||||
- [四、推理测试](#四推理测试)
|
||||
- [五、局限性](#五局限性)
|
||||
- [六、更多资源](#六更多资源)
|
||||
- [参考文献](#参考文献)
|
||||
|
||||
---
|
||||
|
||||
## 一、环境配置
|
||||
|
||||
**方式 1：使用 requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**方式 2：手动安装依赖**
|
||||
|
||||
```bash
|
||||
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
|
||||
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
|
||||
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
|
||||
pip install yunchang xfuser modelscope openpyxl
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
pip install deepspeed==0.17.0 numpy==1.26.4
|
||||
```
|
||||
|
||||
> 部分视频 reward 模型（如 `VideoAlignReward`、`HPSv3Reward`）基于 Qwen2-VL 骨干，需要 `transformers==4.51.3`。使用这些 reward 时请安装对应版本。
|
||||
|
||||
**方式 3：使用 docker**
|
||||
|
||||
使用 docker 的情况下，请保证机器中已经正确安装显卡驱动与 CUDA 环境，然后以此执行以下命令：
|
||||
|
||||
```
|
||||
# pull image
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
|
||||
# enter image
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 二、数据准备
|
||||
|
||||
### 2.1 快速测试数据集
|
||||
|
||||
Reward LoRA 训练从文本 prompt 出发跑可微 rollout，因此标准的 T2V reward 路径（`train_mode="normal"`）只需一个**纯 prompt 数据集**即可，这也是推荐配置。对于 I2V/inpaint 路径，则使用视频数据集来提供参考条件。
|
||||
|
||||
```bash
|
||||
# 下载官方示例数据集
|
||||
modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
|
||||
```
|
||||
|
||||
### 2.2 Reward 模型权重
|
||||
|
||||
不同 reward 的骨干/checkpoint 来源不同：
|
||||
|
||||
- **`HPSReward`（HPS v2.1）与 `MPSReward`**：打分头很小，会在**首次运行时自动下载**到 torch hub 缓存（`~/.cache/torch/hub`）；其中 **HPS 还需给 `hpsv2` 内置的 open_clip 补两个特殊配置文件**（见下），MPS 无需额外准备；
|
||||
- **`VideoAlignReward`（VideoReward）与 `HPSv3Reward`**：基于 Qwen2-VL 骨干，体积较大，**建议提前下载到本地**，再用 `--reward_fn_kwargs` 指向本地路径。
|
||||
|
||||
各 reward 的依赖与来源汇总如下（`--reward_fn` 取值见 [3.3](#33-reward-专用参数解析)）：
|
||||
|
||||
| Reward | 骨干 / 依赖 | 打分头 / checkpoint | 手动下载 |
|
||||
|--|--|--|--|
|
||||
| `HPSReward`（`version="v2.1"`） | open_clip `ViT-H-14`，需 `pip install hpsv2` + `factory.py` 补丁与 BPE 词表 | `HPS_v2.1_compressed.pt`（自动下载） | 否，可用 `model_path` 覆盖 |
|
||||
| `MPSReward` | `laion/CLIP-ViT-H-14-laion2B-s32B-b79K`（config/tokenizer 自动下载），MPS 代码已内置于 `videox_fun/reward/MPS/` | `MPS_overall.pth`（自动下载） | 否，可用 `model_path` 覆盖 |
|
||||
| `VideoAlignReward` | `Qwen/Qwen2-VL-2B-Instruct` | `KlingTeam/VideoReward`（含 `model_config.json` 与权重） | **是**，`model_path` 必填 |
|
||||
| `HPSv3Reward` | `Qwen/Qwen2-VL-7B-Instruct` | `MizzenAI/HPSv3` 的 `HPSv3.safetensors` | 建议，缺省时自动从 HF 下载 |
|
||||
|
||||
**HPS v2.1（`HPSReward`）**——除了 `pip install hpsv2`，还要给 `hpsv2` 内置的 open_clip 补两个特殊配置文件：`create_model_and_transforms("ViT-H-14", ...)` 用到了 HPS 分支新增的参数，需要打过补丁的 `factory.py`；`get_tokenizer("ViT-H-14")` 需要 BPE 词表 `bpe_simple_vocab_16e6.txt.gz`。缺任一都会在加载时报错。`ViT-H-14` 骨干与 `HPS_v2.1_compressed.pt` 打分头会在首次运行时自动下载到 torch hub 缓存：
|
||||
|
||||
```bash
|
||||
pip install hpsv2
|
||||
# 覆盖 hpsv2 内置 open_clip 的 factory.py 补丁，并补上 BPE 词表
|
||||
site_packages=$(python -c "import site; print(site.getsitepackages()[0])")
|
||||
wget -O $site_packages/hpsv2/src/open_clip/factory.py https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/package/patches/hpsv2_src_open_clip_factory_patches.py
|
||||
wget -O $site_packages/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz https://github.com/tgxs002/HPSv2/raw/refs/heads/master/hpsv2/src/open_clip/bpe_simple_vocab_16e6.txt.gz
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="HPSReward" \
|
||||
--reward_fn_kwargs='{"version": "v2.1"}'
|
||||
```
|
||||
|
||||
**MPS（`MPSReward`）**——MPS 代码已内置于 `videox_fun/reward/MPS/`，无需打补丁；`MPS_overall.pth` 打分头会自动从 OSS 下载，`laion/CLIP-ViT-H-14-laion2B-s32B-b79K` 的 config 与 tokenizer 会自动从 HuggingFace 下载（若无法访问 HF，可用 `processor_name_or_path` 指向本地镜像）：
|
||||
|
||||
```bash
|
||||
--reward_fn="MPSReward" \
|
||||
--reward_fn_kwargs='{"max_reward": 1, "loss_scale": 1}'
|
||||
```
|
||||
|
||||
**VideoReward（`VideoAlignReward`）**——需 Qwen2-VL-2B 骨干 + VideoReward 打分头目录（内含 `model_config.json` 与权重），二者都要提前下载，且 `model_path` 必填：
|
||||
|
||||
```bash
|
||||
modelscope download --model Qwen/Qwen2-VL-2B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-2B-Instruct
|
||||
huggingface-cli download KlingTeam/VideoReward --local-dir models/Personalized_Model/VideoReward
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="VideoAlignReward" \
|
||||
--reward_fn_kwargs='{"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "reward_dim": "TA+MQ", "fps": 16, "use_norm": true, "differentiable": true, "max_reward": 2.0}'
|
||||
```
|
||||
|
||||
**HPSv3（`HPSv3Reward`）**——需 Qwen2-VL-7B 骨干 + `HPSv3.safetensors`；不指定 `checkpoint_path` 时会自动从 HF `MizzenAI/HPSv3` 下载，建议提前落盘：
|
||||
|
||||
```bash
|
||||
modelscope download --model Qwen/Qwen2-VL-7B-Instruct --local_dir models/Personalized_Model/Qwen2-VL-7B-Instruct
|
||||
modelscope download --model MizzenAI/HPSv3 --local_dir models/Personalized_Model/HPSv3
|
||||
```
|
||||
|
||||
```bash
|
||||
--reward_fn="HPSv3Reward" \
|
||||
--reward_fn_kwargs='{"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}'
|
||||
```
|
||||
|
||||
> 使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 必须设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会切断梯度链（见 [3.3](#33-reward-专用参数解析)）。
|
||||
|
||||
### 2.3 数据集结构与 metadata.json
|
||||
|
||||
```
|
||||
📦 datasets/
|
||||
├── 📂 my_dataset/
|
||||
│ ├── 📂 train/
|
||||
│ │ ├── 📄 video001.mp4
|
||||
│ │ ├── 📄 video002.mp4
|
||||
│ │ └── 📄 ...
|
||||
│ └── 📄 metadata.json
|
||||
```
|
||||
|
||||
**相对路径格式**（示例）：
|
||||
```json
|
||||
[
|
||||
{
|
||||
"file_path": "train/video001.mp4",
|
||||
"text": "A beautiful sunset over the ocean, golden hour lighting",
|
||||
"type": "video",
|
||||
"width": 1024,
|
||||
"height": 1024
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
**关键字段说明**：
|
||||
- `file_path`：视频路径（相对路径或绝对路径）。
|
||||
- `text`：视频描述（英文 prompt）。
|
||||
- `type`：数据类型，固定为 `"video"`。
|
||||
- `width` / `height`：视频宽高（**强烈建议提供**，用于分桶训练）。可用 `scripts/process_json_add_width_and_height.py` 提取这两个字段。
|
||||
|
||||
**相对路径与绝对路径**：
|
||||
- 相对路径：`export DATASET_NAME="datasets/X-Fun-Videos-Demo/"` 与 `export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"`。
|
||||
- 绝对路径：`export DATASET_NAME=""` 与 `export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"`。
|
||||
|
||||
---
|
||||
|
||||
## 三、Reward LoRA 训练
|
||||
|
||||
### 3.1 下载预训练基座模型
|
||||
|
||||
```bash
|
||||
# 创建模型目录
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# 下载 Wan2.2 官方权重
|
||||
# T2V 模型（文本到视频，双 Transformer）
|
||||
modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-T2V-A14B
|
||||
# 或 I2V 模型（图像到视频，双 Transformer）
|
||||
# modelscope download --model Wan-AI/Wan2.2-I2V-A14B --local_dir models/Diffusion_Transformer/Wan2.2-I2V-A14B
|
||||
# 或 TI2V 模型（文本-图像到视频，单 Transformer）
|
||||
# modelscope download --model Wan-AI/Wan2.2-TI2V-5B --local_dir models/Diffusion_Transformer/Wan2.2-TI2V-5B
|
||||
```
|
||||
|
||||
### 3.2 快速开始（FSDP）
|
||||
|
||||
Reward LoRA 训练会把完整的去噪 rollout 与 reward 模型的前向计算图保持在显存中，因此推荐使用 FSDP（`FULL_SHARD`）。当 `--boundary_type="high"` 时，*另一个*（冻结的）专家在训练期已被分片并常驻 GPU，验证阶段会直接复用它，而不会从磁盘再加载一份副本。
|
||||
|
||||
**Wan2.2 T2V Reward LoRA 训练示例（单一 HPS v2.1 reward，FSDP）**：
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# NCCL_IB_DISABLE=1 与 NCCL_P2P_DISABLE=1 用于没有 RDMA 的多机环境。
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
export NCCL_DEBUG=INFO
|
||||
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False \
|
||||
scripts/wan2.2/train_reward_lora.py \
|
||||
--config_path="config/wan2.2/wan_civitai_t2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--image_sample_size=640 \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=1 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=50 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_wan2.2_reward_lora" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--enable_bucket \
|
||||
--boundary_type="high" \
|
||||
--rank=128 \
|
||||
--network_alpha=64 \
|
||||
--target_name="q,k,v,ffn.0,ffn.2" \
|
||||
--train_mode="normal" \
|
||||
--guidance_scale=5.0 \
|
||||
--num_inference_steps=40 \
|
||||
--num_decoded_latents=3 \
|
||||
--reward_fn="HPSReward" \
|
||||
--reward_fn_kwargs='{"version": "v2.1"}' \
|
||||
--backprop \
|
||||
--backprop_strategy="tail" \
|
||||
--backprop_num_steps=1 \
|
||||
--low_vram
|
||||
```
|
||||
|
||||
> 本目录下的 `train_reward_lora.sh` 提供了一份可直接运行的模板。请根据你的环境调整 `MODEL_NAME`、数据集路径、`--reward_fn` 与 `--boundary_type`。
|
||||
|
||||
### 3.3 Reward 专用参数解析
|
||||
|
||||
除标准 LoRA 参数外（见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)），reward 训练脚本还新增以下参数：
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|-----------|-------------|---------------|
|
||||
| `--reward_fn` | Reward 模型类名，可用**逗号分隔的列表**启用多 reward 训练；各 reward 的可微 loss 以加权和合并。可选：`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`、`VideoAlignReward`。 | `"HPSv3Reward,VideoAlignReward"` |
|
||||
| `--reward_fn_kwargs` | Reward 模型的关键词参数。单 reward 用**扁平 JSON dict**；多 reward 用**以 reward 名为 key 的嵌套 JSON dict**。 | `'{"VideoAlignReward": {"max_reward": 2.0}, "HPSv3Reward": {"max_reward": 12}}'` |
|
||||
| `--multi_reward_weights` | 多 reward 各 loss 权重的 JSON dict。会归一化到和为 1；多 reward 时若省略则使用等权重；单 reward 时忽略。 | `'{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}'` |
|
||||
| `--num_inference_steps` | 可微 rollout 与验证中的去噪步数。 | 40 |
|
||||
| `--num_decoded_latents` | 解码回像素以供 reward 评估的潜帧数。分块因果 VAE 受梯度检查点限制，除非设置 `--use_1steps_vae`，否则要求 `<= 3`。 | 3 |
|
||||
| `--num_reward_frames` | 随机选取该数量的解码帧用于 reward 评估。`None` = 使用全部。仅对**图像** reward 生效（见下）。 | 2 |
|
||||
| `--use_1steps_vae` | 使用单段式 Wan2.2 VAE（`AutoencoderKLWan3_8_1Steps`）做可微解码；数值等价，一次前向解码整个 latent，干净地按解码块应用梯度检查点，并解除 `<= 3` 的限制。仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。 | - |
|
||||
| `--backprop` | 启用 reward 反向传播训练模式。 | - |
|
||||
| `--backprop_strategy` | rollout 中哪些步保留梯度图：`last`、`tail`、`uniform`、`random`。 | `tail` |
|
||||
| `--backprop_num_steps` | backprop 步数，仅当 `--backprop_strategy` 为 `tail`/`uniform`/`random` 时生效。 | 1 |
|
||||
| `--first_frame_mse_weight` | inpaint 模式下首帧 latent MSE 辅助 loss 的权重。 | 1.0 |
|
||||
|
||||
**图像 reward 与视频 reward（`is_image_reward`）**：
|
||||
- **图像 reward**（`AestheticReward`、`HPSReward`、`PickScoreReward`、`MPSReward`、`HPSv3Reward`）逐帧打分，因此帧抽稀（`--num_reward_frames`）是安全的。
|
||||
- **视频 reward**（`VideoAlignReward`，`is_image_reward=False`）需要完整时序序列来评判运动质量；若压缩帧会把视频退化为近乎静态的输入并破坏运动维度。因此这类 reward 保留全部解码帧（数量受 `--num_decoded_latents` 约束）。
|
||||
- 在 inpaint 模式（`train_mode="inpaint"`）下，所有 reward 在评估前都会丢弃首帧。
|
||||
|
||||
> **重要**：使用 `--backprop` 时，`VideoAlignReward` / `HPSv3Reward` 这类视频 reward 必须在 `--reward_fn_kwargs` 中设置 `"differentiable": true`，否则其 `@torch.no_grad()` 路径会彻底切断梯度链。
|
||||
|
||||
### 3.4 单 Reward 与多 Reward 训练
|
||||
|
||||
单 reward 使用 `--reward_fn="HPSv3Reward"` 配一个扁平的 `--reward_fn_kwargs`。若要组合多个目标，传入逗号分隔的列表、嵌套的 kwargs dict，以及（可选的）权重。每个 reward 的 loss 与 reward 值会分开记录（如 `train/loss_HPSv3Reward`、`train/reward_VideoAlignReward`），而实际优化的 `loss = Σ wᵢ · lossᵢ`。
|
||||
|
||||
**多 reward 示例（HPSv3 + VideoAlign，T2V）**——最强的对齐配置，因为视频 reward 能补上图像 reward 无法覆盖的时序（运动/一致性）信号：
|
||||
|
||||
```bash
|
||||
--num_decoded_latents=13 \
|
||||
--num_reward_frames=2 \
|
||||
--use_1steps_vae \
|
||||
--reward_fn="VideoAlignReward,HPSv3Reward" \
|
||||
--reward_fn_kwargs='{"VideoAlignReward": {"model_path": "models/Personalized_Model/VideoReward/", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-2B-Instruct", "fps": 16, "reward_dim": "TA+MQ", "differentiable": true, "use_norm": true, "max_reward": 2.0, "loss_scale": 1.0}, "HPSv3Reward": {"checkpoint_path": "models/Personalized_Model/HPSv3/HPSv3.safetensors", "model_name_or_path": "models/Personalized_Model/Qwen2-VL-7B-Instruct", "differentiable": true, "max_reward": 12}}' \
|
||||
--multi_reward_weights='{"VideoAlignReward": 0.5, "HPSv3Reward": 0.5}' \
|
||||
--backprop \
|
||||
--backprop_strategy="tail" \
|
||||
--backprop_num_steps=1
|
||||
```
|
||||
|
||||
### 3.5 可微解码与 `--use_1steps_vae`
|
||||
|
||||
默认的分块因果 VAE（`AutoencoderKLWan3_8`）按块应用梯度检查点，从而把可微解码限制在**前 3 个潜帧**（`--num_decoded_latents <= 3`）。不设置 `--use_1steps_vae` 而使用 `--num_decoded_latents > 3` 会抛出 `ValueError`。
|
||||
|
||||
启用 `--use_1steps_vae` 可切换到单段式变体，它在一次前向中解码整个 latent、按解码块干净地应用梯度检查点，并解除该限制——当视频 reward 需要更多时序上下文时很有用。它数值等价，且仅影响 `AutoencoderKLWan3_8`（2.2 VAE）路径。
|
||||
|
||||
### 3.6 训练验证
|
||||
|
||||
你可以配置验证参数，在训练过程中周期性生成测试视频，以监控对齐进度。在 FSDP 下，`log_validation` 里的 pipeline 前向是一次跨 rank 的集合通信，因此**所有进程都必须进入验证调用**（它刻意**没有**包在 `is_main_process` 里）；可训练专家已带上 LoRA，冻结专家则被复用而非重新加载。
|
||||
|
||||
| 参数 | 说明 | 默认值 |
|
||||
|-----------|-------------|---------------|
|
||||
| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
|
||||
| `--validation_epochs` | 每 N 个 epoch 运行一次验证 | 5 |
|
||||
| `--validation_prompts` | 用于验证视频生成的 prompt | None |
|
||||
| `--validation_paths` | I2V 验证的参考图路径（仅 inpaint 模式） | None |
|
||||
|
||||
**验证示例**：
|
||||
|
||||
```bash
|
||||
--validation_prompts="A stylish woman strolls down a bustling Tokyo street at night, warm neon glow and animated city signage reflecting on the wet pavement; she wears a black leather jacket over a red dress, confident casual demeanor, medium shot, cinematic." \
|
||||
--validation_steps=50 \
|
||||
--validation_epochs=500
|
||||
```
|
||||
|
||||
> 验证视频会保存到 `output_dir` 目录。若做 T2V 验证（无参考图），保持 `train_mode="normal"` 并省略 `--validation_paths`。
|
||||
|
||||
---
|
||||
|
||||
## 四、推理测试
|
||||
|
||||
训练好的 reward LoRA 就是一个普通 LoRA，加载方式与 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 完全一致。对于 A14B 双 Transformer 模型，在 [examples/wan2.2/predict_t2v.py](https://github.com/aigc-apps/VideoX-Fun/blob/main/examples/wan2.2/predict_t2v.py)（或 `predict_i2v.py`）中，为 low-noise reward LoRA 设置 `lora_path` / `lora_weight`，为 high-noise reward LoRA 设置 `lora_high_path` / `lora_high_weight`：
|
||||
|
||||
```python
|
||||
# 训练好的 reward LoRA 权重路径，例如
|
||||
# "output_dir_wan2.2_reward_lora/checkpoint-xxx/diffusion_pytorch_model.safetensors"
|
||||
lora_path = None
|
||||
lora_high_path = None
|
||||
lora_weight = 0.55
|
||||
lora_high_weight = 0.55
|
||||
```
|
||||
|
||||
单卡推理运行：
|
||||
|
||||
```bash
|
||||
python examples/wan2.2/predict_t2v.py
|
||||
```
|
||||
|
||||
显存管理模式与多卡并行推理的详细说明，请参见 [README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md) 的推理测试章节。
|
||||
|
||||
---
|
||||
|
||||
## 五、局限性
|
||||
|
||||
1. 训练到一定程度后，reward 持续上升，但生成视频的质量不再进一步提升。模型会取巧地学到一些捷径（在背景里添加伪影，即对抗补丁）来抬高 reward。
|
||||
2. 目前仍缺乏适合视频生成的偏好模型。直接使用图像偏好模型无法沿时序维度评估偏好（如运动性与一致性），并且会降低生成视频的运动性。可以仅用解码视频的第一帧来计算 reward 来部分缓解，但影响依然存在——这也正是对于运动敏感的对齐，推荐使用视频 reward（如 `VideoAlignReward`）并配合 `--use_1steps_vae` 与更大的 `--num_decoded_latents`、或采用多 reward 组合的原因。
|
||||
|
||||
---
|
||||
|
||||
## 六、更多资源
|
||||
|
||||
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
|
||||
- **标准 LoRA 训练**：[README_TRAIN_LORA_zh-CN.md](./README_TRAIN_LORA_zh-CN.md)
|
||||
- **英文版**：[README_TRAIN_REWARD.md](./README_TRAIN_REWARD.md)
|
||||
|
||||
---
|
||||
|
||||
## 参考文献
|
||||
<ol>
|
||||
<li id="ref1">Clark, Kevin, et al. "Directly fine-tuning diffusion models on differentiable rewards.". In ICLR 2024.</li>
|
||||
<li id="ref2">Prabhudesai, Mihir, et al. "Aligning text-to-image diffusion models with reward backpropagation." arXiv preprint arXiv:2310.03739 (2023).</li>
|
||||
</ol>
|
||||
+1111
-422
File diff suppressed because it is too large
Load Diff
@@ -1,65 +1,48 @@
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-Fun-A14B-InP"
|
||||
export TRAIN_PROMPT_PATH="MovieGenVideoBench_train.txt"
|
||||
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
|
||||
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
|
||||
# NCCL_IB_DISABLE=1 与 NCCL_P2P_DISABLE=1 用于没有 RDMA 的多机环境。
|
||||
# export NCCL_IB_DISABLE=1
|
||||
# export NCCL_P2P_DISABLE=1
|
||||
export NCCL_DEBUG=INFO
|
||||
|
||||
# Train HPSv2.1 reward LoRA for the low noise model of Wan2.2-Fun-A14B-InP
|
||||
accelerate launch --mixed_precision="bf16" --num-processes=8 --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json scripts/wan2.2_fun/train_reward_lora.py \
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False \
|
||||
scripts/wan2.2_fun/train_reward_lora.py \
|
||||
--config_path="config/wan2.2/wan_civitai_i2v.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_size=640 \
|
||||
--token_sample_size=640 \
|
||||
--video_sample_stride=2 \
|
||||
--video_sample_n_frames=81 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--max_train_steps=10000 \
|
||||
--checkpointing_steps=100 \
|
||||
--learning_rate=1e-05 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=50 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir" \
|
||||
--output_dir="output_dir_wan2.2_fun_reward_lora" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--max_grad_norm=0.3 \
|
||||
--vae_mini_batch=1 \
|
||||
--enable_bucket \
|
||||
--boundary_type="low" \
|
||||
--lora_skip_name="ffn" \
|
||||
--low_vram \
|
||||
--use_deepspeed \
|
||||
--prompt_path=$TRAIN_PROMPT_PATH \
|
||||
--train_sample_height=256 \
|
||||
--train_sample_width=256 \
|
||||
--num_inference_steps=40 \
|
||||
--video_length=81 \
|
||||
--num_decoded_latents=1 \
|
||||
--rank=128 \
|
||||
--network_alpha=64 \
|
||||
--target_name="q,k,v,ffn.0,ffn.2" \
|
||||
--num_inference_steps=25 \
|
||||
--num_decoded_latents=3 \
|
||||
--reward_fn="HPSReward" \
|
||||
--reward_fn_kwargs='{"version": "v2.1"}' \
|
||||
--backprop \
|
||||
--backprop_strategy="tail" \
|
||||
--backprop_num_steps=1 \
|
||||
--backprop
|
||||
|
||||
# Train MPS reward LoRA for the high noise model of Wan2.2-Fun-A14B-InP
|
||||
# accelerate launch --mixed_precision="bf16" --num-processes=8 --use_deepspeed --deepspeed_config_file config/zero_stage3_config_cpu_offload.json scripts/wan2.2_fun/train_reward_lora.py \
|
||||
# --config_path="config/wan2.2/wan_civitai_i2v.yaml" \
|
||||
# --pretrained_model_name_or_path=$MODEL_NAME \
|
||||
# --train_batch_size=1 \
|
||||
# --gradient_accumulation_steps=1 \
|
||||
# --max_train_steps=10000 \
|
||||
# --checkpointing_steps=100 \
|
||||
# --learning_rate=1e-05 \
|
||||
# --seed=42 \
|
||||
# --output_dir="output_dir" \
|
||||
# --gradient_checkpointing \
|
||||
# --mixed_precision="bf16" \
|
||||
# --adam_weight_decay=3e-2 \
|
||||
# --adam_epsilon=1e-10 \
|
||||
# --max_grad_norm=0.3 \
|
||||
# --boundary_type="high" \
|
||||
# --lora_skip_name="ffn" \
|
||||
# --low_vram \
|
||||
# --use_deepspeed \
|
||||
# --prompt_path=$TRAIN_PROMPT_PATH \
|
||||
# --train_sample_height=256 \
|
||||
# --train_sample_width=256 \
|
||||
# --num_inference_steps=40 \
|
||||
# --video_length=81 \
|
||||
# --num_decoded_latents=1 \
|
||||
# --reward_fn="MPSReward" \
|
||||
# --backprop_strategy="tail" \
|
||||
# --backprop_num_steps=1 \
|
||||
# --backprop
|
||||
--train_mode="inpaint" \
|
||||
--low_vram
|
||||
Vendored
+2
-1
@@ -25,5 +25,6 @@ from .minimax_h3_xfuser import MiniMaxH3MultiGPUsAttnProcessor
|
||||
from .qwen_xfuser import QwenImageMultiGPUsAttnProcessor2_0
|
||||
from .qwenimage21_xfuser import QwenImage21MultiGPUsAttnProcessor
|
||||
from .wan_xfuser import (usp_attn_forward, usp_attn_s2v_forward,
|
||||
usp_attn_self_forcing_forward)
|
||||
usp_attn_self_forcing_forward,
|
||||
usp_attn_sparkwan_forward)
|
||||
from .z_image_xfuser import ZMultiGPUsSingleStreamAttnProcessor
|
||||
|
||||
Vendored
+7
-6
@@ -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
|
||||
|
||||
|
||||
|
||||
Vendored
+106
-1
@@ -6,7 +6,7 @@ import torch.cuda.amp as amp
|
||||
from .fuser import (get_sequence_parallel_rank,
|
||||
get_sequence_parallel_world_size, get_sp_group,
|
||||
init_distributed_environment, initialize_model_parallel,
|
||||
xFuserLongContextAttention)
|
||||
ulysses_all_to_all, xFuserLongContextAttention)
|
||||
|
||||
|
||||
def pad_freqs(original_tensor, target_len):
|
||||
@@ -162,6 +162,111 @@ def usp_attn_forward(self,
|
||||
x = self.o(x)
|
||||
return x
|
||||
|
||||
def usp_attn_sparkwan_forward(self,
|
||||
x,
|
||||
seq_lens,
|
||||
grid_sizes,
|
||||
freqs,
|
||||
dtype=torch.bfloat16,
|
||||
t=0):
|
||||
"""Sequence-parallel self-attention that KEEPS the RoLa sparse kernel (``self.attn_op``).
|
||||
|
||||
The stock ``usp_attn_forward`` binds ``xFuserLongContextAttention``'s dense full-softmax
|
||||
and never reaches ``self.attn_op``, so using it on a ``SparkWanSelfAttention`` would
|
||||
silently drop the trained RoLa params (``attn_op.{proj_q,proj_k,gate_proj,gate_bias}``)
|
||||
and shift the numbers. This variant uses Ulysses head-parallelism instead: an all-to-all
|
||||
turns each rank's sequence shard ``[B, S/P, H, D]`` into a head shard ``[B, S, H/P, D]``,
|
||||
the RoLa branches (block-sparse exact softmax + rank-truncated low-rank linear + token
|
||||
gate) run over the FULL sequence on a subset of heads -- every branch is head-independent
|
||||
and topk is a per-query key-block ratio, so this reproduces the single-GPU result --
|
||||
then a second all-to-all returns to the sequence-shard layout.
|
||||
|
||||
Two RoLa-specific points under head-parallelism:
|
||||
* ``gate_bias`` is per-head ``[1, H, 1, 1]``; each rank holds the contiguous head block
|
||||
``[sp_rank * H/P, (sp_rank + 1) * H/P)`` produced by the ulysses all-to-all, so the
|
||||
gate slices ``gate_bias`` to that range.
|
||||
* the low-rank branch's rank RoPE runs on the FULL (post-all-to-all) sequence with the
|
||||
non-SP ``rope_apply_lowrank``, so positions are complete and need no sharding.
|
||||
|
||||
Requires a pure-Ulysses split: ``num_heads`` must divide the SP degree and ``ring_degree``
|
||||
should be 1 (the SLA kernel has no ring-attention path). ``topk >= 1.0`` (the dense
|
||||
fallback) defers to ``usp_attn_forward``.
|
||||
"""
|
||||
# Dense student (topk >= 1.0) or a plain Wan self-attention: keep the stock SP path.
|
||||
if getattr(self, "topk", 1.0) >= 1.0 or not hasattr(self, "attn_op"):
|
||||
return usp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=dtype, t=t)
|
||||
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
half_dtypes = (torch.float16, torch.bfloat16)
|
||||
|
||||
def half(t_):
|
||||
return t_ if t_.dtype in half_dtypes else t_.to(dtype)
|
||||
|
||||
# query, key, value function
|
||||
def qkv_fn(x):
|
||||
q = self.norm_q(self.q(x)).view(b, s, n, d)
|
||||
k = self.norm_k(self.k(x)).view(b, s, n, d)
|
||||
v = self.v(x).view(b, s, n, d)
|
||||
return q, k, v
|
||||
|
||||
q, k, v = qkv_fn(x)
|
||||
# PRE-RoPE qk-normed q/k feed the RoLa low-rank branch (RoPE is applied on the low-rank
|
||||
# bottleneck, outside the non-linear feature map). rope_apply_qk returns new tensors, so
|
||||
# q_pre/k_pre stay valid after the rebind below.
|
||||
q_pre, k_pre = q, k
|
||||
q, k = rope_apply_qk(q, k, grid_sizes, freqs)
|
||||
|
||||
sp_size = get_sequence_parallel_world_size()
|
||||
sp_rank = get_sequence_parallel_rank()
|
||||
|
||||
# PureSLA ablation: fused block-sparse + feature-map linear (proj_l), BSHD in/out.
|
||||
if getattr(self, "attn_variant", "rola") == "pure_sla":
|
||||
q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
|
||||
k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
|
||||
v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
|
||||
x = self.attn_op(q, k, v)
|
||||
x = ulysses_all_to_all(x, scatter_dim=1, gather_dim=2)
|
||||
x = x.flatten(2)
|
||||
x = self.o(x)
|
||||
return x
|
||||
|
||||
# ---- RoLa: seq-shard -> head-shard all-to-all for every branch input ----
|
||||
# post-RoPE q/k (sparse branch), pre-RoPE q/k (low-rank branch), v, and the raw hidden
|
||||
# states x (token gate) are all gathered to [B, S_full, H/P, D].
|
||||
q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
|
||||
k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
|
||||
v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
|
||||
q_pre = ulysses_all_to_all(half(q_pre), scatter_dim=2, gather_dim=1)
|
||||
k_pre = ulysses_all_to_all(half(k_pre), scatter_dim=2, gather_dim=1)
|
||||
x_full = ulysses_all_to_all(half(x).view(b, s, n, d), scatter_dim=2, gather_dim=1)
|
||||
|
||||
# BSHD -> BHLD over the head shard.
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2)
|
||||
v_bhld = v.transpose(1, 2)
|
||||
q_pre = q_pre.transpose(1, 2)
|
||||
k_pre = k_pre.transpose(1, 2)
|
||||
x_bhld = x_full.transpose(1, 2)
|
||||
|
||||
o_s = self.attn_op.sparse(q, k, v_bhld)
|
||||
o_l = self.attn_op.lowrank(q_pre, k_pre, v_bhld, grid_sizes, freqs)
|
||||
# gate_bias is per-head; slice to this rank's contiguous head block from the all-to-all.
|
||||
h_per = n // sp_size
|
||||
head_slice = slice(sp_rank * h_per, (sp_rank + 1) * h_per)
|
||||
gate = self.attn_op.gate(x_bhld, head_slice=head_slice)
|
||||
if o_s is not None:
|
||||
scale = (o_s.pow(2).mean(dim=-1, keepdim=True) + 1e-5).sqrt()
|
||||
out = ((o_s.float() / scale.float() + gate * o_l.float()) * scale.float()).to(o_l.dtype)
|
||||
else:
|
||||
out = gate * o_l
|
||||
|
||||
# BHLD -> BSHD, head-shard -> seq-shard, back to this rank's sequence chunk.
|
||||
out = out.transpose(1, 2)
|
||||
out = ulysses_all_to_all(out, scatter_dim=1, gather_dim=2)
|
||||
out = out.flatten(2)
|
||||
out = self.o(out)
|
||||
return out
|
||||
|
||||
@amp.autocast(enabled=False)
|
||||
@torch.compiler.disable()
|
||||
def s2v_rope_apply(x, grid_sizes, freqs):
|
||||
|
||||
@@ -90,8 +90,16 @@ from .qwenimage_transformer2d import QwenImageTransformer2DModel
|
||||
from .qwenimage_transformer2d_control import QwenImageControlTransformer2DModel
|
||||
from .qwenimage_transformer2d_instantx import QwenImageInstantXControlNetModel
|
||||
from .qwenimage_vae import AutoencoderKLQwenImage
|
||||
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
|
||||
@@ -108,5 +116,7 @@ from .wan_transformer3d_self_forcing import WanTransformer3DModel_SelfForcing
|
||||
from .wan_transformer3d_vace import VaceWanTransformer3DModel
|
||||
from .wan_vae import AutoencoderKLWan, AutoencoderKLWan_
|
||||
from .wan_vae3_8 import AutoencoderKLWan2_2_, AutoencoderKLWan3_8
|
||||
from .wan_vae3_8_1steps import AutoencoderKLWan3_8_1Steps
|
||||
from .wan_vae_1steps import AutoencoderKLWan_1Steps
|
||||
from .z_image_transformer2d import ZImageTransformer2DModel
|
||||
from .z_image_transformer2d_control import ZImageControlTransformer2DModel
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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):
|
||||
@@ -424,4 +427,624 @@ class SparseLinearAttention(nn.Module):
|
||||
if return_sparsity:
|
||||
return o, real_topk / sparse_map.shape[-1]
|
||||
else:
|
||||
return o
|
||||
return o
|
||||
|
||||
|
||||
class RoLASparseAttention(nn.Module):
|
||||
R'''
|
||||
RoLa sparse + low-rank self-attention branch (the official SparkDiffusion default).
|
||||
|
||||
Mirrors ``sparkdiffusion/networks/wan_rola_attention.py::WanSelfAttentionRoLa``:
|
||||
- sparse branch : block-sparse exact softmax via the in-repo SLA Triton kernel
|
||||
(same ``get_block_map`` + ``_sparse_linear_attention`` used by
|
||||
``SparseLinearAttention``; runs on post-RoPE q/k).
|
||||
- low-rank branch : ``proj_q``/``proj_k`` (head_dim -> rank) -> SiLU -> rank-truncated
|
||||
3D RoPE -> linear attention ``Q_lr (K_lr^T V)`` -> RMSNorm
|
||||
(runs on PRE-RoPE, qk-normed q/k; RoPE is applied OUTSIDE the
|
||||
non-linear feature map, on the low-rank bottleneck).
|
||||
- token gate : ``sigmoid(gate_proj(x) + gate_bias)`` scales the low-rank output
|
||||
magnitude per token/head; ``gate`` is computed from the raw hidden
|
||||
states ``x`` (not q/k/v).
|
||||
|
||||
Fusion (per official)::
|
||||
|
||||
scale = sqrt(mean(o_s^2) + 1e-5) # per-token/head RMS of the sparse output
|
||||
out = (o_s / scale + gate * o_l) * scale # == o_s + gate * o_l * scale
|
||||
|
||||
New learnable params per layer (all under ``attn_op.*``): ``proj_q``, ``proj_k``,
|
||||
``gate_proj``, ``gate_bias``. They are absent from a dense checkpoint (loaded with
|
||||
strict=False) and start at their init below. This replaces the earlier PureSLA path
|
||||
(single zero-init ``proj_l`` + softmax feature-map linear attention).
|
||||
|
||||
All forward helpers take/return BHLD (``[B, H, L, *]``) tensors.
|
||||
'''
|
||||
def __init__(self, head_dim, num_heads, topk=0.03, rola_rank=64,
|
||||
BLKQ=64, BLKK=64, use_bf16=True):
|
||||
R'''
|
||||
Args:
|
||||
head_dim: dimension of each head.
|
||||
num_heads: number of attention heads (sizes ``gate_bias``).
|
||||
topk: ratio of key blocks kept for the sparse branch (0.1 -> 90% sparse).
|
||||
rola_rank: low-rank bottleneck dimension of the linear branch.
|
||||
BLKQ/BLKK: query/key block sizes for the sparse kernel.
|
||||
use_bf16: compute the sparse kernel in bfloat16 (default) or float16.
|
||||
'''
|
||||
super().__init__()
|
||||
self.head_dim = head_dim
|
||||
self.num_heads = num_heads
|
||||
self.topk = topk
|
||||
self.rola_rank = rola_rank
|
||||
self.BLKQ = BLKQ
|
||||
self.BLKK = BLKK
|
||||
self.dtype = torch.bfloat16 if use_bf16 else torch.float16
|
||||
|
||||
r = rola_rank
|
||||
self.proj_q = nn.Linear(head_dim, r, bias=False)
|
||||
self.proj_k = nn.Linear(head_dim, r, bias=False)
|
||||
self.gate_proj = nn.Linear(head_dim, 1, bias=False)
|
||||
# -1.946 -> sigmoid ~= 0.125, so the sparse branch dominates at init.
|
||||
self.gate_bias = nn.Parameter(torch.full((1, num_heads, 1, 1), -1.946))
|
||||
|
||||
self.init_weights_()
|
||||
|
||||
def init_weights_(self):
|
||||
nn.init.kaiming_uniform_(self.proj_q.weight, a=math.sqrt(5))
|
||||
nn.init.kaiming_uniform_(self.proj_k.weight, a=math.sqrt(5))
|
||||
nn.init.zeros_(self.gate_proj.weight)
|
||||
with torch.no_grad():
|
||||
self.gate_bias.fill_(-1.946)
|
||||
|
||||
def sparse(self, q, k, v):
|
||||
R'''Block-sparse exact softmax attention.
|
||||
|
||||
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
|
||||
(the caller then falls back to ``gate * o_l``).
|
||||
'''
|
||||
L = q.shape[2]
|
||||
if L < max(self.BLKQ, self.BLKK):
|
||||
return None
|
||||
orig_dtype = q.dtype
|
||||
q = q.contiguous()
|
||||
k = k.contiguous()
|
||||
v = v.contiguous()
|
||||
sparse_map, lut, real_topk = get_block_map(
|
||||
q, k, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK)
|
||||
q = q.to(self.dtype)
|
||||
k = k.to(self.dtype)
|
||||
v = v.to(self.dtype)
|
||||
o_s = _sparse_linear_attention.apply(
|
||||
q, k, v, sparse_map, lut, real_topk, self.BLKQ, self.BLKK)
|
||||
return o_s.to(orig_dtype)
|
||||
|
||||
def lowrank(self, q_normed, k_normed, v, grid_sizes, freqs):
|
||||
R'''Low-rank linear attention branch with rank-truncated RoPE.
|
||||
|
||||
Args:
|
||||
q_normed, k_normed: [B, H, L, D] PRE-RoPE, qk-normed queries/keys.
|
||||
v: [B, H, L, D] values.
|
||||
grid_sizes: [B, 3] (F, H, W) for the rank RoPE expansion.
|
||||
freqs: complex rotary table [M, head_dim // 2].
|
||||
Returns:
|
||||
[B, H, L, D], RMS-normalized.
|
||||
'''
|
||||
from .spark_transformer3d import rope_apply_lowrank
|
||||
|
||||
q_lr = F.silu(self.proj_q(q_normed)) # [B, H, L, rank]
|
||||
k_lr = F.silu(self.proj_k(k_normed))
|
||||
# RoPE applied outside the non-linear feature map, on the rank bottleneck.
|
||||
q_lr = rope_apply_lowrank(q_lr, grid_sizes, freqs, self.rola_rank)
|
||||
k_lr = rope_apply_lowrank(k_lr, grid_sizes, freqs, self.rola_rank)
|
||||
c = k_lr.transpose(-2, -1) @ v # [B, H, rank, D]
|
||||
out = q_lr @ c # [B, H, L, D]
|
||||
return F.rms_norm(out, [out.shape[-1]])
|
||||
|
||||
def gate(self, x_bhld, head_slice=None):
|
||||
R'''Token-wise gate from the raw hidden states.
|
||||
|
||||
Args:
|
||||
x_bhld: [B, H, L, D] hidden states (attention input), head layout.
|
||||
head_slice: optional slice selecting this rank's head shard (multi-GPU
|
||||
Ulysses head-parallel), so the per-head ``gate_bias`` lines up.
|
||||
Returns:
|
||||
[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)
|
||||
|
||||
|
||||
# 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
|
||||
@@ -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]
|
||||
@@ -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]
|
||||
@@ -0,0 +1,420 @@
|
||||
# SparkDiffusion sparse student backbones for the Wan2.2 A14B MoE models and the
|
||||
# Wan2.1 single-expert models.
|
||||
import torch
|
||||
|
||||
from .attention_utils import (RoLASparseAttention, SparseLinearAttention,
|
||||
attention)
|
||||
from .wan_transformer3d import (Wan2_2Transformer3DModel, WanSelfAttention, WanTransformer3DModel,
|
||||
rope_apply_qk)
|
||||
|
||||
|
||||
@torch.compiler.disable()
|
||||
def rope_apply_lowrank(x, grid_sizes, freqs, rank):
|
||||
"""Rank-truncated 3D RoPE for the RoLa low-rank bottleneck (official convention).
|
||||
|
||||
Applies rotary positional embedding to a rank-``r`` bottleneck instead of the full
|
||||
head_dim, reusing a prefix of the pre-trained rotary schedule (per RoLA, arXiv
|
||||
2609.06712) rather than recomputing frequencies. A rank-``r`` feature carries
|
||||
``r // 2`` complex rotary pairs; following the upstream SparkDiffusion
|
||||
``_apply_rope_rank``, the rotated pairs take the CONTIGUOUS prefix of the
|
||||
[t|h|w]-ordered pre-trained table (``freqs[:L, :rope_dim]``): all temporal pairs
|
||||
first, then height pairs, then width pairs, in table order, with interleaved
|
||||
(adjacent-pair) rotary coupling; the remaining dims pass through unchanged.
|
||||
|
||||
Args:
|
||||
x: [B, H, L, rank] low-rank bottleneck (BHLD), real-valued.
|
||||
grid_sizes: [B, 3] (F, H, W) spatial grid per sample.
|
||||
freqs: pre-trained complex rotary table [M, head_dim // 2] (t|h|w concatenated).
|
||||
rank: bottleneck dimension ``r`` (== ``x.shape[-1]``).
|
||||
Returns:
|
||||
[B, H, L, rank] with RoPE applied (same layout/dtype as ``x``).
|
||||
"""
|
||||
c = freqs.shape[1] # head_dim // 2
|
||||
rope_dim = min(rank // 2, c) # rotary pairs the bottleneck rotates
|
||||
t_blk, h_blk, w_blk = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
|
||||
# Contiguous prefix budget in table order: temporal, then height, then width.
|
||||
n_t = min(rope_dim, t_blk.shape[1])
|
||||
n_h = min(rope_dim - n_t, h_blk.shape[1])
|
||||
n_w = rope_dim - n_t - n_h
|
||||
n_heads = x.size(1)
|
||||
dtype = x.dtype
|
||||
output = []
|
||||
for i, (f, h, w) in enumerate(grid_sizes.tolist()):
|
||||
seq_len = f * h * w
|
||||
# Per-token prefix table [seq_len, rope_dim], t-major then h then w ordering.
|
||||
prefix = torch.cat([
|
||||
t_blk[:f, :n_t].view(f, 1, 1, -1).expand(f, h, w, -1),
|
||||
h_blk[:h, :n_h].view(1, h, 1, -1).expand(f, h, w, -1),
|
||||
w_blk[:w, :n_w].view(1, 1, w, -1).expand(f, h, w, -1)
|
||||
],
|
||||
dim=-1).reshape(seq_len, rope_dim)
|
||||
# Rotate the first 2*rope_dim dims: x * exp(i*freq), interleaved pairs.
|
||||
x_i = torch.view_as_complex(x[i, :, :seq_len, :2 * rope_dim].to(
|
||||
torch.float32).reshape(n_heads, seq_len, rope_dim, 2))
|
||||
x_i = torch.view_as_real(x_i * prefix.unsqueeze(0)).flatten(-2)
|
||||
# Concatenate padding tokens (if any) and the untouched tail dims.
|
||||
x_i = torch.cat([x_i, x[i, :, seq_len:, :2 * rope_dim]], dim=1)
|
||||
x_i = torch.cat([x_i, x[i, :, :, 2 * rope_dim:]], dim=-1)
|
||||
output.append(x_i)
|
||||
return torch.stack(output).to(dtype)
|
||||
|
||||
|
||||
class SparkWanSelfAttention(WanSelfAttention):
|
||||
"""Wan self-attention whose softmax kernel is swapped for RoLa sparse+low-rank attention."""
|
||||
|
||||
def __init__(self,
|
||||
dim,
|
||||
num_heads,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
eps=1e-6,
|
||||
topk=0.03,
|
||||
attn_variant='rola',
|
||||
rola_rank=64):
|
||||
super().__init__(dim, num_heads, window_size, qk_norm, eps)
|
||||
# topk: ratio of key blocks kept per query (0.1 -> 90% sparse, 0.05 -> 95%, 0.03 -> 97%).
|
||||
# >= 1.0 falls back to the original dense attention.
|
||||
self.topk = topk
|
||||
self.attn_variant = attn_variant
|
||||
if topk < 1.0:
|
||||
if attn_variant == 'rola':
|
||||
self.attn_op = RoLASparseAttention(
|
||||
self.head_dim, num_heads, topk, rola_rank=rola_rank)
|
||||
elif attn_variant == 'pure_sla':
|
||||
self.attn_op = SparseLinearAttention(self.head_dim, topk)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"unknown attn_variant {attn_variant!r}; expected 'rola' or 'pure_sla'")
|
||||
|
||||
def forward(self, x, seq_lens, grid_sizes, freqs, dtype=torch.bfloat16, t=0):
|
||||
r"""
|
||||
Args:
|
||||
x(Tensor): Shape [B, L, num_heads, C / num_heads]
|
||||
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)
|
||||
|
||||
# Keep the PRE-RoPE qk-normed q/k for the RoLa low-rank branch (rope_apply_qk
|
||||
# returns new tensors, so q_pre/k_pre stay valid after the rebind below).
|
||||
q_pre, k_pre = q, k
|
||||
q, k = rope_apply_qk(q, k, grid_sizes, freqs)
|
||||
|
||||
# PureSLA ablation path: fused block-sparse + feature-map linear (proj_l).
|
||||
if self.attn_variant == 'pure_sla':
|
||||
x = self.attn_op(q.to(dtype), k.to(dtype), v.to(dtype))
|
||||
x = x.to(dtype).flatten(2)
|
||||
return self.o(x)
|
||||
|
||||
# ---- RoLa: block-sparse branch + gated rank-truncated low-rank branch ----
|
||||
v_bhld = v.transpose(1, 2) # [B,H,L,D]
|
||||
o_s = self.attn_op.sparse(q.transpose(1, 2), k.transpose(1, 2), v_bhld)
|
||||
o_l = self.attn_op.lowrank(
|
||||
q_pre.transpose(1, 2), k_pre.transpose(1, 2), v_bhld, grid_sizes, freqs)
|
||||
# Gate is computed from the raw hidden states (attention input), per official RoLa.
|
||||
x_bhld = x.view(b, s, n, d).transpose(1, 2) # [B,H,L,D]
|
||||
gate = self.attn_op.gate(x_bhld) # [B,H,L,1]
|
||||
if o_s is not None:
|
||||
scale = (o_s.pow(2).mean(dim=-1, keepdim=True) + 1e-5).sqrt()
|
||||
out = ((o_s.float() / scale.float() + gate * o_l.float()) * scale.float()).to(o_l.dtype)
|
||||
else:
|
||||
out = gate * o_l
|
||||
x = out.transpose(1, 2).flatten(2) # BHLD -> [B,L,H*D]
|
||||
return self.o(x)
|
||||
|
||||
|
||||
class SparkWan2_2Transformer3DModel(Wan2_2Transformer3DModel):
|
||||
"""Wan2.2 transformer (T2V / I2V, single expert per checkpoint) with a sparse
|
||||
student 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,
|
||||
student_topk=0.03,
|
||||
attn_variant='rola',
|
||||
rola_rank=64,
|
||||
):
|
||||
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,
|
||||
)
|
||||
self.student_topk = student_topk
|
||||
self.attn_variant = attn_variant
|
||||
self.rola_rank = rola_rank
|
||||
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 = SparkWanSelfAttention(
|
||||
dense_attn.dim, dense_attn.num_heads, dense_attn.window_size,
|
||||
dense_attn.qk_norm, dense_attn.eps, topk=topk,
|
||||
attn_variant=self.attn_variant, rola_rank=self.rola_rank,
|
||||
)
|
||||
# Under low_cpu_mem_usage from_pretrained the replaced params are meta tensors
|
||||
# and get filled by the checkpoint load afterwards; only upgrade in place (copy
|
||||
# the real tensors, the freshly-init RoLa params always survive) when 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 RoLa 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):
|
||||
"""(Re)build the sparse kernels at the given top-k ratio. Always recreates
|
||||
attn_op, so the RoLa params (proj_q/proj_k/gate_proj/gate_bias) are reset to their
|
||||
init: call right after from_pretrained and BEFORE loading a stage-1 sparse checkpoint
|
||||
(its trained RoLa params would be dropped otherwise)."""
|
||||
self.student_topk = topk
|
||||
self._apply_sparse_self_attention(topk)
|
||||
|
||||
def enable_multi_gpus_inference(self,):
|
||||
"""Bind the sparse-aware USP forward instead of the parent's dense one.
|
||||
|
||||
``Wan2_2Transformer3DModel.enable_multi_gpus_inference`` monkey-patches every
|
||||
``blocks[i].self_attn.forward`` to ``usp_attn_forward`` (dense full-softmax via
|
||||
``xFuserLongContextAttention``), which never reaches ``SparkWanSelfAttention.attn_op``
|
||||
and would silently drop the trained RoLa params under multi-GPU. Rebind to
|
||||
``usp_attn_sparkwan_forward`` (ulysses head-parallel + RoLa kernel) so the sparse branch
|
||||
stays live; a plain ``shard_model`` FSDP wrap does not help here -- it only shards
|
||||
params and shares the same SP branch, so the patch still applies.
|
||||
"""
|
||||
import types
|
||||
|
||||
from ..dist import (get_sequence_parallel_rank,
|
||||
get_sequence_parallel_world_size, get_sp_group,
|
||||
usp_attn_sparkwan_forward)
|
||||
self.sp_world_size = get_sequence_parallel_world_size()
|
||||
self.sp_world_rank = get_sequence_parallel_rank()
|
||||
self.all_gather = get_sp_group().all_gather
|
||||
|
||||
# Ulysses head-split needs num_heads divisible by the SP degree.
|
||||
n_heads = self.blocks[0].self_attn.num_heads
|
||||
if self.sp_world_size > 1 and n_heads % self.sp_world_size != 0:
|
||||
raise ValueError(
|
||||
f"sparse multi-GPU (ulysses) needs num_heads ({n_heads}) divisible by the "
|
||||
f"sequence-parallel degree ({self.sp_world_size}); set ulysses_degree accordingly "
|
||||
f"and keep ring_degree=1 (the SLA kernel has no ring-attention path).")
|
||||
|
||||
for block in self.blocks:
|
||||
block.self_attn.forward = types.MethodType(
|
||||
usp_attn_sparkwan_forward, block.self_attn)
|
||||
|
||||
if hasattr(self, 'vace_blocks'):
|
||||
for block in self.vace_blocks:
|
||||
block.self_attn.forward = types.MethodType(
|
||||
usp_attn_sparkwan_forward, block.self_attn)
|
||||
|
||||
def sparse_trainable_params(self):
|
||||
"""Names of the extra sparse-branch parameters, for logging / stage-A freezing."""
|
||||
return [name for name, _ in self.named_parameters() if "attn_op." in name]
|
||||
|
||||
|
||||
class SparkWanTransformer3DModel(WanTransformer3DModel):
|
||||
"""Wan2.1 transformer (T2V / I2V, single expert) with a sparse student
|
||||
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.03,
|
||||
attn_variant='rola',
|
||||
rola_rank=64,
|
||||
):
|
||||
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.attn_variant = attn_variant
|
||||
self.rola_rank = rola_rank
|
||||
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 = SparkWanSelfAttention(
|
||||
dense_attn.dim, dense_attn.num_heads, dense_attn.window_size,
|
||||
dense_attn.qk_norm, dense_attn.eps, topk=topk,
|
||||
attn_variant=self.attn_variant, rola_rank=self.rola_rank,
|
||||
)
|
||||
# Under low_cpu_mem_usage from_pretrained the replaced params are meta tensors
|
||||
# and get filled by the checkpoint load afterwards; only upgrade in place (copy
|
||||
# the real tensors, the freshly-init RoLa params always survive) when 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 RoLa 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):
|
||||
"""(Re)build the sparse kernels at the given top-k ratio. Always recreates
|
||||
attn_op, so the RoLa params (proj_q/proj_k/gate_proj/gate_bias) are reset to their
|
||||
init: call right after from_pretrained and BEFORE loading a stage-1 sparse checkpoint
|
||||
(its trained RoLa params would be dropped otherwise)."""
|
||||
self.student_topk = topk
|
||||
self._apply_sparse_self_attention(topk)
|
||||
|
||||
def enable_multi_gpus_inference(self,):
|
||||
"""Bind the sparse-aware USP forward instead of the parent's dense one.
|
||||
|
||||
``WanTransformer3DModel.enable_multi_gpus_inference`` monkey-patches every
|
||||
``blocks[i].self_attn.forward`` to ``usp_attn_forward`` (dense full-softmax via
|
||||
``xFuserLongContextAttention``), which never reaches ``SparkWanSelfAttention.attn_op``
|
||||
and would silently drop the trained RoLa params under multi-GPU. Rebind to
|
||||
``usp_attn_sparkwan_forward`` (ulysses head-parallel + RoLa kernel) so the sparse branch
|
||||
stays live; a plain ``shard_model`` FSDP wrap does not help here -- it only shards
|
||||
params and shares the same SP branch, so the patch still applies.
|
||||
"""
|
||||
import types
|
||||
|
||||
from ..dist import (get_sequence_parallel_rank,
|
||||
get_sequence_parallel_world_size, get_sp_group,
|
||||
usp_attn_sparkwan_forward)
|
||||
self.sp_world_size = get_sequence_parallel_world_size()
|
||||
self.sp_world_rank = get_sequence_parallel_rank()
|
||||
self.all_gather = get_sp_group().all_gather
|
||||
|
||||
# Ulysses head-split needs num_heads divisible by the SP degree.
|
||||
n_heads = self.blocks[0].self_attn.num_heads
|
||||
if self.sp_world_size > 1 and n_heads % self.sp_world_size != 0:
|
||||
raise ValueError(
|
||||
f"sparse multi-GPU (ulysses) needs num_heads ({n_heads}) divisible by the "
|
||||
f"sequence-parallel degree ({self.sp_world_size}); set ulysses_degree accordingly "
|
||||
f"and keep ring_degree=1 (the SLA kernel has no ring-attention path).")
|
||||
|
||||
for block in self.blocks:
|
||||
block.self_attn.forward = types.MethodType(
|
||||
usp_attn_sparkwan_forward, block.self_attn)
|
||||
|
||||
if hasattr(self, 'vace_blocks'):
|
||||
for block in self.vace_blocks:
|
||||
block.self_attn.forward = types.MethodType(
|
||||
usp_attn_sparkwan_forward, block.self_attn)
|
||||
|
||||
def sparse_trainable_params(self):
|
||||
"""Names of the extra sparse-branch parameters, for logging / stage-A freezing."""
|
||||
return [name for name, _ in self.named_parameters() if "attn_op." in name]
|
||||
@@ -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)
|
||||
+302
-66
@@ -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):
|
||||
@@ -335,24 +397,142 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
# 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__(
|
||||
@@ -366,6 +546,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:
|
||||
@@ -373,7 +554,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)
|
||||
@@ -392,6 +575,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
|
||||
|
||||
@@ -411,12 +596,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
|
||||
@@ -437,7 +622,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)
|
||||
@@ -459,48 +644,99 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
|
||||
return self.decode(z, return_dict=return_dict)
|
||||
|
||||
@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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -703,6 +703,24 @@ class WanTransformer3DModel_FlexForcing(WanTransformer3DModel_SelfForcing):
|
||||
optionally clipped to ``local_attn_size`` frames.
|
||||
"""
|
||||
chunk_sizes = self._require_exact_partition(chunk_sizes, num_frames)
|
||||
# Read the value the mask below actually uses - not the config mirror, so
|
||||
# the check cannot pass against a window the builder then ignores.
|
||||
local_attn_size = self.local_attn_size
|
||||
# One chunk is one bidirectional unit, so every frame inside it has to see
|
||||
# every other. The clipped window below is anchored at the chunk *end*
|
||||
# (`kv_idx >= ends[q_idx] - local_attn_size * frame_seqlen`), so a chunk
|
||||
# wider than the window quietly loses its own head-to-tail connectivity:
|
||||
# the layout would be trained as if it were partially causal while every
|
||||
# caller describes it as bidirectional. Reject it loudly instead. The same
|
||||
# rule is enforced on the KV-cache read side in the pipelines.
|
||||
if local_attn_size != -1 and max(chunk_sizes) > local_attn_size:
|
||||
raise ValueError(
|
||||
f"A chunk of {max(chunk_sizes)} latent frames exceeds "
|
||||
f"`local_attn_size={local_attn_size}`: the window is anchored at "
|
||||
f"the chunk end, so that chunk's first frames could not attend to "
|
||||
f"its last ones and its intra-chunk bidirectionality (§3.1) would "
|
||||
f"be lost. Use `local_attn_size=-1` or keep every chunk within the "
|
||||
f"window.")
|
||||
total_length = num_frames * frame_seqlen
|
||||
|
||||
# Padded query rows keep ends=0 exactly like the inherited uniform
|
||||
@@ -712,8 +730,6 @@ class WanTransformer3DModel_FlexForcing(WanTransformer3DModel_SelfForcing):
|
||||
device=device, dtype=torch.long)
|
||||
ends[:total_length] = chunk_ends_tensor(chunk_sizes, frame_seqlen, device=device)
|
||||
|
||||
local_attn_size = self.local_attn_size
|
||||
|
||||
def attention_mask(b, h, q_idx, kv_idx):
|
||||
if local_attn_size == -1:
|
||||
# Global chunk-wise causal: bidirectional inside a chunk, all
|
||||
@@ -723,6 +739,17 @@ class WanTransformer3DModel_FlexForcing(WanTransformer3DModel_SelfForcing):
|
||||
return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx)
|
||||
|
||||
self._build_block_mask(attention_mask, total_length, device)
|
||||
# Keep the parent's cached-mask descriptor in sync with the mask just
|
||||
# installed: the guard reuses `self.block_mask` whenever
|
||||
# `_block_mask_expected_len`/`_block_mask_is_teacher_forcing` match the
|
||||
# next call, so a builder that replaced the mask without updating them
|
||||
# would leave them describing a mask that no longer exists. The parent
|
||||
# forward re-records the same two values right after dispatching here,
|
||||
# so this is the invariant being stated at its source rather than a
|
||||
# rebuild saving - `set_flex_chunk_sizes` is what forces one rebuild per
|
||||
# new partition by clearing the length.
|
||||
self._block_mask_expected_len = total_length
|
||||
self._block_mask_is_teacher_forcing = False
|
||||
|
||||
# Consumers that still expect the scalar Self-Forcing knobs (Forcing-KV
|
||||
# AR stride / rolling-cache budget) get the conservative largest chunk.
|
||||
@@ -777,6 +804,11 @@ class WanTransformer3DModel_FlexForcing(WanTransformer3DModel_SelfForcing):
|
||||
return eye_mask | clean_mask | noise_mask
|
||||
|
||||
self._build_block_mask(attention_mask, total_length, device)
|
||||
# Same descriptor sync as the non-teacher-forcing builder above; here
|
||||
# `total_length` already counts clean + noisy, which is exactly what the
|
||||
# parent guard derives when `clean_x is not None`.
|
||||
self._block_mask_expected_len = total_length
|
||||
self._block_mask_is_teacher_forcing = True
|
||||
|
||||
self.num_frame_per_block = max(chunk_sizes)
|
||||
|
||||
|
||||
@@ -0,0 +1,945 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
import logging
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.cuda.amp as amp
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
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.modeling_outputs import AutoencoderKLOutput
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.utils.accelerate_utils import apply_forward_hook
|
||||
from einops import rearrange
|
||||
from torch.utils.checkpoint import checkpoint as torch_checkpoint
|
||||
|
||||
CACHE_T = 2
|
||||
|
||||
|
||||
class CausalConv3d(nn.Conv3d):
|
||||
"""
|
||||
Causal 3d convolusion.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._padding = (
|
||||
self.padding[2],
|
||||
self.padding[2],
|
||||
self.padding[1],
|
||||
self.padding[1],
|
||||
2 * self.padding[0],
|
||||
0,
|
||||
)
|
||||
self.padding = (0, 0, 0)
|
||||
|
||||
def forward(self, x):
|
||||
# Single-pass causal conv: the full temporal sequence is available at
|
||||
# once, so the past is provided purely by the front temporal pad in
|
||||
# ``self._padding`` (no feat_cache threading).
|
||||
x = F.pad(x, self._padding)
|
||||
x = super().forward(x)
|
||||
return x
|
||||
|
||||
|
||||
class RMS_norm(nn.Module):
|
||||
|
||||
def __init__(self, dim, channel_first=True, images=True, bias=False):
|
||||
super().__init__()
|
||||
broadcastable_dims = (1, 1, 1) if not images else (1, 1)
|
||||
shape = (dim, *broadcastable_dims) if channel_first else (dim,)
|
||||
|
||||
self.channel_first = channel_first
|
||||
self.scale = dim**0.5
|
||||
self.gamma = nn.Parameter(torch.ones(shape))
|
||||
self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
|
||||
|
||||
def forward(self, x):
|
||||
return (F.normalize(x, dim=(1 if self.channel_first else -1)) *
|
||||
self.scale * self.gamma + self.bias)
|
||||
|
||||
|
||||
class Upsample(nn.Upsample):
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Fix bfloat16 support for nearest neighbor interpolation.
|
||||
"""
|
||||
return super().forward(x.float()).type_as(x)
|
||||
|
||||
|
||||
class Resample(nn.Module):
|
||||
|
||||
def __init__(self, dim, mode):
|
||||
assert mode in (
|
||||
"none",
|
||||
"upsample2d",
|
||||
"upsample3d",
|
||||
"downsample2d",
|
||||
"downsample3d",
|
||||
)
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.mode = mode
|
||||
|
||||
# layers
|
||||
if mode == "upsample2d":
|
||||
self.resample = nn.Sequential(
|
||||
Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
|
||||
nn.Conv2d(dim, dim, 3, padding=1),
|
||||
)
|
||||
elif mode == "upsample3d":
|
||||
self.resample = nn.Sequential(
|
||||
Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
|
||||
nn.Conv2d(dim, dim, 3, padding=1),
|
||||
# nn.Conv2d(dim, dim//2, 3, padding=1)
|
||||
)
|
||||
self.time_conv = CausalConv3d(
|
||||
dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
|
||||
elif mode == "downsample2d":
|
||||
self.resample = nn.Sequential(
|
||||
nn.ZeroPad2d((0, 1, 0, 1)),
|
||||
nn.Conv2d(dim, dim, 3, stride=(2, 2)))
|
||||
elif mode == "downsample3d":
|
||||
self.resample = nn.Sequential(
|
||||
nn.ZeroPad2d((0, 1, 0, 1)),
|
||||
nn.Conv2d(dim, dim, 3, stride=(2, 2)))
|
||||
self.time_conv = CausalConv3d(
|
||||
dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
|
||||
else:
|
||||
self.resample = nn.Identity()
|
||||
|
||||
def forward(self, x):
|
||||
b, c, t, h, w = x.size()
|
||||
|
||||
if self.mode == 'upsample3d':
|
||||
# Causal temporal upsample without cache: the first frame is kept
|
||||
# as-is and only the remaining frames are temporally doubled. This
|
||||
# reproduces the chunked path (frame 0 emitted with zero left
|
||||
# context, later frames with 2 frames of context).
|
||||
first_frame = x[:, :, :1, :, :]
|
||||
rest_frames = x[:, :, 1:, :, :] if t > 1 else None
|
||||
if rest_frames is not None:
|
||||
rest_frames = self.time_conv(rest_frames)
|
||||
rest_frames = rest_frames.reshape(b, 2, c, t - 1, h, w)
|
||||
rest_frames = torch.stack((rest_frames[:, 0], rest_frames[:, 1]),
|
||||
3)
|
||||
rest_frames = rest_frames.reshape(b, c, (t - 1) * 2, h, w)
|
||||
x = torch.cat([first_frame, rest_frames], dim=2)
|
||||
else:
|
||||
x = first_frame
|
||||
|
||||
t = x.shape[2]
|
||||
x = rearrange(x, 'b c t h w -> (b t) c h w')
|
||||
x = self.resample(x)
|
||||
x = rearrange(x, '(b t) c h w -> b c t h w', t=t)
|
||||
|
||||
if self.mode == 'downsample3d':
|
||||
# Causal temporal downsample without cache: first frame passes
|
||||
# through unchanged and the stride-2 causal conv runs over the full
|
||||
# sequence (its own windows supply the previous-frame left context),
|
||||
# exactly reproducing the chunked cache-threaded path.
|
||||
if x.shape[2] > 1:
|
||||
first_frame = x[:, :, :1, :, :]
|
||||
rest_frames = self.time_conv(x)
|
||||
x = torch.cat([first_frame, rest_frames], dim=2)
|
||||
|
||||
return x
|
||||
|
||||
def init_weight(self, conv):
|
||||
conv_weight = conv.weight.detach().clone()
|
||||
nn.init.zeros_(conv_weight)
|
||||
c1, c2, t, h, w = conv_weight.size()
|
||||
one_matrix = torch.eye(c1, c2)
|
||||
init_matrix = one_matrix
|
||||
nn.init.zeros_(conv_weight)
|
||||
conv_weight.data[:, :, 1, 0, 0] = init_matrix # * 0.5
|
||||
conv.weight = nn.Parameter(conv_weight)
|
||||
nn.init.zeros_(conv.bias.data)
|
||||
|
||||
def init_weight2(self, conv):
|
||||
conv_weight = conv.weight.data.detach().clone()
|
||||
nn.init.zeros_(conv_weight)
|
||||
c1, c2, t, h, w = conv_weight.size()
|
||||
init_matrix = torch.eye(c1 // 2, c2)
|
||||
conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
|
||||
conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
|
||||
conv.weight = nn.Parameter(conv_weight)
|
||||
nn.init.zeros_(conv.bias.data)
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
|
||||
def __init__(self, in_dim, out_dim, dropout=0.0):
|
||||
super().__init__()
|
||||
self.in_dim = in_dim
|
||||
self.out_dim = out_dim
|
||||
|
||||
# layers
|
||||
self.residual = nn.Sequential(
|
||||
RMS_norm(in_dim, images=False),
|
||||
nn.SiLU(),
|
||||
CausalConv3d(in_dim, out_dim, 3, padding=1),
|
||||
RMS_norm(out_dim, images=False),
|
||||
nn.SiLU(),
|
||||
nn.Dropout(dropout),
|
||||
CausalConv3d(out_dim, out_dim, 3, padding=1),
|
||||
)
|
||||
self.shortcut = (
|
||||
CausalConv3d(in_dim, out_dim, 1)
|
||||
if in_dim != out_dim else nn.Identity())
|
||||
|
||||
def forward(self, x):
|
||||
h = self.shortcut(x)
|
||||
x = self.residual(x)
|
||||
return x + h
|
||||
|
||||
|
||||
class AttentionBlock(nn.Module):
|
||||
"""
|
||||
Causal self-attention with a single head.
|
||||
"""
|
||||
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
|
||||
# layers
|
||||
self.norm = RMS_norm(dim)
|
||||
self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
|
||||
self.proj = nn.Conv2d(dim, dim, 1)
|
||||
|
||||
# zero out the last layer params
|
||||
nn.init.zeros_(self.proj.weight)
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
b, c, t, h, w = x.size()
|
||||
x = rearrange(x, "b c t h w -> (b t) c h w")
|
||||
x = self.norm(x)
|
||||
# compute query, key, value
|
||||
q, k, v = (
|
||||
self.to_qkv(x).reshape(b * t, 1, c * 3,
|
||||
-1).permute(0, 1, 3,
|
||||
2).contiguous().chunk(3, dim=-1))
|
||||
|
||||
# apply attention
|
||||
x = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
)
|
||||
x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
|
||||
|
||||
# output
|
||||
x = self.proj(x)
|
||||
x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
|
||||
return x + identity
|
||||
|
||||
|
||||
def patchify(x, patch_size):
|
||||
if patch_size == 1:
|
||||
return x
|
||||
if x.dim() == 4:
|
||||
x = rearrange(
|
||||
x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
|
||||
elif x.dim() == 5:
|
||||
x = rearrange(
|
||||
x,
|
||||
"b c f (h q) (w r) -> b (c r q) f h w",
|
||||
q=patch_size,
|
||||
r=patch_size,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid input shape: {x.shape}")
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def unpatchify(x, patch_size):
|
||||
if patch_size == 1:
|
||||
return x
|
||||
|
||||
if x.dim() == 4:
|
||||
x = rearrange(
|
||||
x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
|
||||
elif x.dim() == 5:
|
||||
x = rearrange(
|
||||
x,
|
||||
"b (c r q) f h w -> b c f (h q) (w r)",
|
||||
q=patch_size,
|
||||
r=patch_size,
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
class AvgDown3D(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
factor_t,
|
||||
factor_s=1,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.factor_t = factor_t
|
||||
self.factor_s = factor_s
|
||||
self.factor = self.factor_t * self.factor_s * self.factor_s
|
||||
|
||||
assert in_channels * self.factor % out_channels == 0
|
||||
self.group_size = in_channels * self.factor // out_channels
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
|
||||
pad = (0, 0, 0, 0, pad_t, 0)
|
||||
x = F.pad(x, pad)
|
||||
B, C, T, H, W = x.shape
|
||||
x = x.view(
|
||||
B,
|
||||
C,
|
||||
T // self.factor_t,
|
||||
self.factor_t,
|
||||
H // self.factor_s,
|
||||
self.factor_s,
|
||||
W // self.factor_s,
|
||||
self.factor_s,
|
||||
)
|
||||
x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
|
||||
x = x.view(
|
||||
B,
|
||||
C * self.factor,
|
||||
T // self.factor_t,
|
||||
H // self.factor_s,
|
||||
W // self.factor_s,
|
||||
)
|
||||
x = x.view(
|
||||
B,
|
||||
self.out_channels,
|
||||
self.group_size,
|
||||
T // self.factor_t,
|
||||
H // self.factor_s,
|
||||
W // self.factor_s,
|
||||
)
|
||||
x = x.mean(dim=2)
|
||||
return x
|
||||
|
||||
|
||||
class DupUp3D(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
factor_t,
|
||||
factor_s=1,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.factor_t = factor_t
|
||||
self.factor_s = factor_s
|
||||
self.factor = self.factor_t * self.factor_s * self.factor_s
|
||||
|
||||
assert out_channels * self.factor % in_channels == 0
|
||||
self.repeats = out_channels * self.factor // in_channels
|
||||
|
||||
def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
|
||||
x = x.repeat_interleave(self.repeats, dim=1)
|
||||
x = x.view(
|
||||
x.size(0),
|
||||
self.out_channels,
|
||||
self.factor_t,
|
||||
self.factor_s,
|
||||
self.factor_s,
|
||||
x.size(2),
|
||||
x.size(3),
|
||||
x.size(4),
|
||||
)
|
||||
x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
|
||||
x = x.view(
|
||||
x.size(0),
|
||||
self.out_channels,
|
||||
x.size(2) * self.factor_t,
|
||||
x.size(4) * self.factor_s,
|
||||
x.size(6) * self.factor_s,
|
||||
)
|
||||
if first_chunk:
|
||||
x = x[:, :, self.factor_t - 1:, :, :]
|
||||
return x
|
||||
|
||||
|
||||
class Down_ResidualBlock(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
in_dim,
|
||||
out_dim,
|
||||
dropout,
|
||||
mult,
|
||||
temperal_downsample=False,
|
||||
down_flag=False):
|
||||
super().__init__()
|
||||
|
||||
# Shortcut path with downsample
|
||||
self.avg_shortcut = AvgDown3D(
|
||||
in_dim,
|
||||
out_dim,
|
||||
factor_t=2 if temperal_downsample else 1,
|
||||
factor_s=2 if down_flag else 1,
|
||||
)
|
||||
|
||||
# Main path with residual blocks and downsample
|
||||
downsamples = []
|
||||
for _ in range(mult):
|
||||
downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
|
||||
in_dim = out_dim
|
||||
|
||||
# Add the final downsample block
|
||||
if down_flag:
|
||||
mode = "downsample3d" if temperal_downsample else "downsample2d"
|
||||
downsamples.append(Resample(out_dim, mode=mode))
|
||||
|
||||
self.downsamples = nn.Sequential(*downsamples)
|
||||
|
||||
def forward(self, x):
|
||||
x_copy = x
|
||||
for module in self.downsamples:
|
||||
x = module(x)
|
||||
|
||||
return x + self.avg_shortcut(x_copy)
|
||||
|
||||
|
||||
class Up_ResidualBlock(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
in_dim,
|
||||
out_dim,
|
||||
dropout,
|
||||
mult,
|
||||
temperal_upsample=False,
|
||||
up_flag=False):
|
||||
super().__init__()
|
||||
# Shortcut path with upsample
|
||||
if up_flag:
|
||||
self.avg_shortcut = DupUp3D(
|
||||
in_dim,
|
||||
out_dim,
|
||||
factor_t=2 if temperal_upsample else 1,
|
||||
factor_s=2 if up_flag else 1,
|
||||
)
|
||||
else:
|
||||
self.avg_shortcut = None
|
||||
|
||||
# Main path with residual blocks and upsample
|
||||
upsamples = []
|
||||
for _ in range(mult):
|
||||
upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
|
||||
in_dim = out_dim
|
||||
|
||||
# Add the final upsample block
|
||||
if up_flag:
|
||||
mode = "upsample3d" if temperal_upsample else "upsample2d"
|
||||
upsamples.append(Resample(out_dim, mode=mode))
|
||||
|
||||
self.upsamples = nn.Sequential(*upsamples)
|
||||
|
||||
def forward(self, x):
|
||||
# Single-pass: the whole latent sequence starts at frame 0, so the
|
||||
# DupUp3D shortcut always applies the ``first_chunk`` front trim.
|
||||
x_main = x
|
||||
for module in self.upsamples:
|
||||
x_main = module(x_main)
|
||||
if self.avg_shortcut is not None:
|
||||
x_shortcut = self.avg_shortcut(x, first_chunk=True)
|
||||
return x_main + x_shortcut
|
||||
else:
|
||||
return x_main
|
||||
|
||||
|
||||
class Encoder3d(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim=128,
|
||||
z_dim=4,
|
||||
dim_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
attn_scales=[],
|
||||
temperal_downsample=[True, True, False],
|
||||
dropout=0.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.z_dim = z_dim
|
||||
self.dim_mult = dim_mult
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.attn_scales = attn_scales
|
||||
self.temperal_downsample = temperal_downsample
|
||||
|
||||
# dimensions
|
||||
dims = [dim * u for u in [1] + dim_mult]
|
||||
scale = 1.0
|
||||
|
||||
# init block
|
||||
self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
|
||||
|
||||
# downsample blocks
|
||||
downsamples = []
|
||||
for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
|
||||
t_down_flag = (
|
||||
temperal_downsample[i]
|
||||
if i < len(temperal_downsample) else False)
|
||||
downsamples.append(
|
||||
Down_ResidualBlock(
|
||||
in_dim=in_dim,
|
||||
out_dim=out_dim,
|
||||
dropout=dropout,
|
||||
mult=num_res_blocks,
|
||||
temperal_downsample=t_down_flag,
|
||||
down_flag=i != len(dim_mult) - 1,
|
||||
))
|
||||
scale /= 2.0
|
||||
self.downsamples = nn.Sequential(*downsamples)
|
||||
|
||||
# middle blocks
|
||||
self.middle = nn.Sequential(
|
||||
ResidualBlock(out_dim, out_dim, dropout),
|
||||
AttentionBlock(out_dim),
|
||||
ResidualBlock(out_dim, out_dim, dropout),
|
||||
)
|
||||
|
||||
# # output blocks
|
||||
self.head = nn.Sequential(
|
||||
RMS_norm(out_dim, images=False),
|
||||
nn.SiLU(),
|
||||
CausalConv3d(out_dim, z_dim, 3, padding=1),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
# Single-pass encode: the whole (patchified) video goes through the
|
||||
# encoder in one forward, no feat_cache / chunk loop.
|
||||
x = self.conv1(x)
|
||||
x = self.downsamples(x)
|
||||
x = self.middle(x)
|
||||
x = self.head(x)
|
||||
return x
|
||||
|
||||
|
||||
class Decoder3d(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim=128,
|
||||
z_dim=4,
|
||||
dim_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
attn_scales=[],
|
||||
temperal_upsample=[False, True, True],
|
||||
dropout=0.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.z_dim = z_dim
|
||||
self.dim_mult = dim_mult
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.attn_scales = attn_scales
|
||||
self.temperal_upsample = temperal_upsample
|
||||
|
||||
# dimensions
|
||||
dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
|
||||
scale = 1.0 / 2**(len(dim_mult) - 2)
|
||||
# init block
|
||||
self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
|
||||
|
||||
# middle blocks
|
||||
self.middle = nn.Sequential(
|
||||
ResidualBlock(dims[0], dims[0], dropout),
|
||||
AttentionBlock(dims[0]),
|
||||
ResidualBlock(dims[0], dims[0], dropout),
|
||||
)
|
||||
|
||||
# upsample blocks
|
||||
upsamples = []
|
||||
for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
|
||||
t_up_flag = temperal_upsample[i] if i < len(
|
||||
temperal_upsample) else False
|
||||
upsamples.append(
|
||||
Up_ResidualBlock(
|
||||
in_dim=in_dim,
|
||||
out_dim=out_dim,
|
||||
dropout=dropout,
|
||||
mult=num_res_blocks + 1,
|
||||
temperal_upsample=t_up_flag,
|
||||
up_flag=i != len(dim_mult) - 1,
|
||||
))
|
||||
self.upsamples = nn.Sequential(*upsamples)
|
||||
|
||||
# output blocks
|
||||
self.head = nn.Sequential(
|
||||
RMS_norm(out_dim, images=False),
|
||||
nn.SiLU(),
|
||||
CausalConv3d(out_dim, 12, 3, padding=1),
|
||||
)
|
||||
|
||||
# Toggle for single-pass per-block gradient checkpointing, propagated
|
||||
# from AutoencoderKLWan2_2_.decode.
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(self, x):
|
||||
# Single-pass decode. When gradient checkpointing is enabled each
|
||||
# sub-block is recomputed during backward so only one block's
|
||||
# activations stay alive at a time (vs. the whole upsampled video).
|
||||
def _run(module, t):
|
||||
if self.gradient_checkpointing and torch.is_grad_enabled():
|
||||
return torch_checkpoint(module, t, use_reentrant=False)
|
||||
return module(t)
|
||||
|
||||
x = self.conv1(x)
|
||||
for layer in self.middle:
|
||||
x = _run(layer, x)
|
||||
for layer in self.upsamples:
|
||||
x = _run(layer, x)
|
||||
for layer in self.head:
|
||||
x = _run(layer, x)
|
||||
return x
|
||||
|
||||
|
||||
def count_conv3d(model):
|
||||
count = 0
|
||||
for m in model.modules():
|
||||
if isinstance(m, CausalConv3d):
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
class AutoencoderKLWan2_2_(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim=160,
|
||||
dec_dim=256,
|
||||
z_dim=16,
|
||||
dim_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
attn_scales=[],
|
||||
temperal_downsample=[True, True, False],
|
||||
dropout=0.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.z_dim = z_dim
|
||||
self.dim_mult = dim_mult
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.attn_scales = attn_scales
|
||||
self.temperal_downsample = temperal_downsample
|
||||
self.temperal_upsample = temperal_downsample[::-1]
|
||||
|
||||
# modules
|
||||
self.encoder = Encoder3d(
|
||||
dim,
|
||||
z_dim * 2,
|
||||
dim_mult,
|
||||
num_res_blocks,
|
||||
attn_scales,
|
||||
self.temperal_downsample,
|
||||
dropout,
|
||||
)
|
||||
self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
|
||||
self.conv2 = CausalConv3d(z_dim, z_dim, 1)
|
||||
self.decoder = Decoder3d(
|
||||
dec_dim,
|
||||
z_dim,
|
||||
dim_mult,
|
||||
num_res_blocks,
|
||||
attn_scales,
|
||||
self.temperal_upsample,
|
||||
dropout,
|
||||
)
|
||||
|
||||
# When True the decoder runs with per-block gradient checkpointing.
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(self, x, scale=[0, 1]):
|
||||
mu = self.encode(x, scale)
|
||||
x_recon = self.decode(mu, scale)
|
||||
return x_recon, mu
|
||||
|
||||
def encode(self, x, scale):
|
||||
# Single-pass encode: patchify the full video and run the encoder once
|
||||
# (no chunk loop, no feat_cache).
|
||||
scale = [item.to(x.device, x.dtype) for item in scale]
|
||||
x = patchify(x, patch_size=2)
|
||||
out = self.encoder(x)
|
||||
mu, log_var = self.conv1(out).chunk(2, dim=1)
|
||||
if isinstance(scale[0], torch.Tensor):
|
||||
mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
|
||||
1, self.z_dim, 1, 1, 1)
|
||||
else:
|
||||
mu = (mu - scale[0]) * scale[1]
|
||||
x = torch.cat([mu, log_var], dim=1)
|
||||
return x
|
||||
|
||||
def decode(self, z, scale):
|
||||
# Single-pass decode: run the whole latent sequence through the decoder
|
||||
# in one forward. Because there is no feat_cache mutation, gradient
|
||||
# checkpointing can be applied per decoder block (see Decoder3d.forward)
|
||||
# instead of the previous per-frame cache-threading workaround.
|
||||
scale = [item.to(z.device, z.dtype) for item in scale]
|
||||
if isinstance(scale[0], torch.Tensor):
|
||||
z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
|
||||
1, self.z_dim, 1, 1, 1)
|
||||
else:
|
||||
z = z / scale[1] + scale[0]
|
||||
self.decoder.gradient_checkpointing = self.gradient_checkpointing
|
||||
x = self.conv2(z)
|
||||
out = self.decoder(x)
|
||||
out = unpatchify(out, patch_size=2)
|
||||
return out
|
||||
|
||||
def reparameterize(self, mu, log_var):
|
||||
std = torch.exp(0.5 * log_var)
|
||||
eps = torch.randn_like(std)
|
||||
return eps * std + mu
|
||||
|
||||
def sample(self, imgs, deterministic=False):
|
||||
mu, log_var = self.encode(imgs)
|
||||
if deterministic:
|
||||
return mu
|
||||
std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
|
||||
return mu + std * torch.randn_like(std)
|
||||
|
||||
|
||||
def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs):
|
||||
# params
|
||||
cfg = dict(
|
||||
dim=dim,
|
||||
z_dim=z_dim,
|
||||
dim_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
attn_scales=[],
|
||||
temperal_downsample=[True, True, True],
|
||||
dropout=0.0,
|
||||
)
|
||||
cfg.update(**kwargs)
|
||||
|
||||
# init model
|
||||
model = AutoencoderKLWan2_2_(**cfg)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
class AutoencoderKLWan3_8_1Steps(ModelMixin, ConfigMixin, FromOriginalModelMixin):
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
latent_channels=48,
|
||||
c_dim=160,
|
||||
vae_pth=None,
|
||||
dim_mult=[1, 2, 4, 4],
|
||||
temperal_downsample=[False, True, True],
|
||||
temporal_compression_ratio=4,
|
||||
spatial_compression_ratio=16
|
||||
):
|
||||
super().__init__()
|
||||
mean = torch.tensor(
|
||||
[
|
||||
-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,
|
||||
], dtype=torch.float32
|
||||
)
|
||||
std = torch.tensor(
|
||||
[
|
||||
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,
|
||||
], dtype=torch.float32
|
||||
)
|
||||
self.scale = [mean, 1.0 / std]
|
||||
|
||||
# init model
|
||||
self.model = _video_vae(
|
||||
pretrained_path=vae_pth,
|
||||
z_dim=latent_channels,
|
||||
dim=c_dim,
|
||||
dim_mult=dim_mult,
|
||||
temperal_downsample=temperal_downsample,
|
||||
).eval().requires_grad_(False)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def _set_gradient_checkpointing(self, *args, **kwargs):
|
||||
if "value" in kwargs:
|
||||
self.gradient_checkpointing = kwargs["value"]
|
||||
elif "enable" in kwargs:
|
||||
self.gradient_checkpointing = kwargs["enable"]
|
||||
else:
|
||||
raise ValueError("Invalid set gradient checkpointing")
|
||||
self.model.gradient_checkpointing = self.gradient_checkpointing
|
||||
|
||||
def _encode(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = [
|
||||
self.model.encode(u.unsqueeze(0), self.scale).squeeze(0)
|
||||
for u in x
|
||||
]
|
||||
x = torch.stack(x)
|
||||
return x
|
||||
|
||||
@apply_forward_hook
|
||||
def encode(
|
||||
self, x: torch.Tensor, return_dict: bool = True
|
||||
) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
|
||||
h = self._encode(x)
|
||||
|
||||
posterior = DiagonalGaussianDistribution(h)
|
||||
|
||||
if not return_dict:
|
||||
return (posterior,)
|
||||
return AutoencoderKLOutput(latent_dist=posterior)
|
||||
|
||||
def _decode(self, zs):
|
||||
# Single-pass decode with per-block gradient checkpointing inside the
|
||||
# decoder (see Decoder3d.forward), which lowers the backward peak to a
|
||||
# single decoder block's activations while keeping the whole video in
|
||||
# one forward (no feat_cache chunk loop).
|
||||
self.model.gradient_checkpointing = self.gradient_checkpointing
|
||||
dec = [
|
||||
self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0)
|
||||
for u in zs
|
||||
]
|
||||
dec = torch.stack(dec)
|
||||
|
||||
return DecoderOutput(sample=dec)
|
||||
|
||||
@apply_forward_hook
|
||||
def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
|
||||
decoded = self._decode(z).sample
|
||||
|
||||
if not return_dict:
|
||||
return (decoded,)
|
||||
return DecoderOutput(sample=decoded)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_path, additional_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
|
||||
|
||||
model = cls(**filter_kwargs(cls, additional_kwargs))
|
||||
if pretrained_model_path.endswith(".safetensors"):
|
||||
from safetensors.torch import load_file, safe_open
|
||||
state_dict = load_file(pretrained_model_path)
|
||||
else:
|
||||
state_dict = torch.load(pretrained_model_path, map_location="cpu", weights_only=True)
|
||||
tmp_state_dict = {}
|
||||
for key in state_dict:
|
||||
if not key.startswith("model."):
|
||||
tmp_state_dict["model." + key] = state_dict[key]
|
||||
else:
|
||||
tmp_state_dict[key] = state_dict[key]
|
||||
state_dict = tmp_state_dict
|
||||
m, u = model.load_state_dict(state_dict, strict=False)
|
||||
print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
|
||||
print(m, u)
|
||||
return model
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1721,6 +1721,30 @@ class MiniMaxH3Pipeline(DiffusionPipeline):
|
||||
audio_scheduler=audio_scheduler,
|
||||
)
|
||||
|
||||
@property
|
||||
def _execution_device(self):
|
||||
r"""The device the pipeline executes on.
|
||||
|
||||
diffusers' default resolves this through `self.device`, which returns the `.device` of the first `nn.Module`
|
||||
found while iterating an *unordered* `set` of the component names. That is fragile whenever the components do
|
||||
not all sit on one device: under `--fsdp_offload_params` the student `transformer` keeps its flat param on the
|
||||
CPU at rest and reports `cpu`, while the VAEs and the sharded conditioner are on the GPU, so the resolved device
|
||||
depends on the per-process hash order and can nondeterministically come back `cpu` — which then feeds CPU
|
||||
tensors to GPU-resident FSDP params and raises a device-mismatch on whichever rank drew the short straw. Pin the
|
||||
result to a CUDA component by scanning `self.components` in registration order; a genuinely all-CPU assembly
|
||||
still resolves to the CPU.
|
||||
"""
|
||||
device = super()._execution_device
|
||||
if device.type != "cpu":
|
||||
return device
|
||||
for model in self.components.values():
|
||||
if not isinstance(model, torch.nn.Module):
|
||||
continue
|
||||
model_device = getattr(model, "device", None)
|
||||
if model_device is not None and model_device.type != "cpu":
|
||||
return model_device
|
||||
return device
|
||||
|
||||
@property
|
||||
def vae_spatial_compression_ratio(self) -> int:
|
||||
if getattr(self, "vae", None) is not None:
|
||||
|
||||
@@ -3,8 +3,7 @@
|
||||
import math
|
||||
from contextlib import nullcontext
|
||||
from dataclasses import dataclass
|
||||
from typing import (Any, Callable, Dict, List, Optional, Sequence, Tuple,
|
||||
Union)
|
||||
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
|
||||
|
||||
import torch
|
||||
from diffusers import FlowMatchEulerDiscreteScheduler
|
||||
@@ -13,14 +12,13 @@ from diffusers.utils import BaseOutput, logging, replace_example_docstring
|
||||
|
||||
from ..models import (AutoencoderKLWan, AutoTokenizer, WanT5EncoderModel,
|
||||
WanTransformer3DModel_FlexForcing)
|
||||
from ..utils.flex_chunking import (build_full_then_blocks_partitions,
|
||||
build_pyramid_partitions, chunk_boundaries,
|
||||
normalize_chunk_spec, uniform_chunks,
|
||||
validate_nested_partitions)
|
||||
from ..utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
|
||||
get_sampling_sigmas)
|
||||
from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
|
||||
from ..utils.flex_chunking import (build_full_then_blocks_partitions,
|
||||
build_pyramid_partitions,
|
||||
chunk_boundaries, normalize_chunk_spec,
|
||||
uniform_chunks,
|
||||
validate_nested_partitions)
|
||||
from .pipeline_wan_self_forcing import (WanSelfForcingPipeline,
|
||||
WanSelfForcingPipelineOutput,
|
||||
retrieve_timesteps,
|
||||
@@ -168,6 +166,11 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
(``[1] * F``), ``"bidir"`` (``[F]``) or an int for uniform
|
||||
chunks. ``None`` keeps the inherited ``num_frame_per_block`` /
|
||||
``independent_first_frame`` schedule untouched.
|
||||
The uniform forms (``None``, an int, ``"uniform[:N]"``) honour
|
||||
``independent_first_frame`` - which defaults to ``True`` here -
|
||||
so ``"uniform:3"`` yields ``[1, 3, 3, ...]``, matching the layout
|
||||
an IFF-trained checkpoint expects. Explicit size lists are taken
|
||||
verbatim: write the leading ``1`` yourself if you want it.
|
||||
denoise_mode: Does the frame partition change with the noise level?
|
||||
``"fixed"`` (default) holds one partition across every denoising
|
||||
step, i.e. the block-major Self-Forcing / 3.1 schedule.
|
||||
@@ -302,7 +305,8 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
"generates and currently require T2V mode "
|
||||
"(`initial_latent=None`).")
|
||||
|
||||
base = normalize_chunk_spec(chunk_spec, latent_frames) \
|
||||
base = normalize_chunk_spec(chunk_spec, latent_frames,
|
||||
independent_first_frame=independent_first_frame) \
|
||||
if chunk_spec is not None else None
|
||||
if depth > 1:
|
||||
if ladder_mode in FULL_THEN_BLOCKS_MODES:
|
||||
@@ -310,6 +314,22 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
# two-level ladder whose fine level is the uniform
|
||||
# `num_frame_per_block` partition rather than a binary split, so
|
||||
# it does not go through `build_pyramid_partitions`.
|
||||
# [NOTE] This ladder is fully determined by `num_frame_per_block`:
|
||||
# it ignores both `chunk_spec` (level 0 is always the whole clip)
|
||||
# and `min_num_frame_per_block` (the fine level is the uniform
|
||||
# grid, not a binary split), so say so instead of dropping the
|
||||
# caller's explicit arguments silently.
|
||||
if chunk_spec is not None:
|
||||
logger.warning(
|
||||
"`denoise_mode=%r` ignores `chunk_spec=%r`: its level 0 is "
|
||||
"always one chunk over the whole clip and its fine level "
|
||||
"is the uniform `num_frame_per_block` partition.",
|
||||
ladder_mode, chunk_spec)
|
||||
if min_num_frame_per_block != 1:
|
||||
logger.warning(
|
||||
"`denoise_mode=%r` ignores `min_num_frame_per_block=%r`: "
|
||||
"the finest level is fixed by `num_frame_per_block`.",
|
||||
ladder_mode, min_num_frame_per_block)
|
||||
partitions = build_full_then_blocks_partitions(
|
||||
latent_frames, num_frame_per_block,
|
||||
independent_first_frame=independent_first_frame)
|
||||
@@ -324,25 +344,71 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
chunk_sizes = partitions[0]
|
||||
else:
|
||||
# `chunk_spec` cannot be None here: that case returned above.
|
||||
chunk_sizes = normalize_chunk_spec(chunk_spec, latent_frames)
|
||||
chunk_sizes = normalize_chunk_spec(
|
||||
chunk_spec, latent_frames,
|
||||
independent_first_frame=independent_first_frame)
|
||||
partitions = None
|
||||
|
||||
local_attn_size = getattr(self.transformer.config, 'local_attn_size', -1)
|
||||
if partitions is not None:
|
||||
# The pyramid re-noises between levels, which only the stochastic
|
||||
# schedule defines, and it moves `current_start` backwards whenever
|
||||
# a chunk splits. Both are safe against a full KV cache and neither
|
||||
# is against a rolling one.
|
||||
if not stochastic_sampling:
|
||||
# A level change hands the next level a buffer that sits exactly at
|
||||
# `timesteps[step_idx + 1]`, and the finer level restarts at frame
|
||||
# 0, so `current_start` moves backwards. Both are safe against a
|
||||
# full KV cache and neither is against a rolling one, which evicts
|
||||
# by `current_start` deltas.
|
||||
#
|
||||
# The handoff itself is not tied to `stochastic_sampling`. The
|
||||
# re-noising branch defines the buffer as
|
||||
# `(1 - t_i_1) * x0 + t_i_1 * eps`; the deterministic branch
|
||||
# defines it as the Euler increment
|
||||
# `x + (sigma_next - sigma_t) * flow_pred`, which lands on the
|
||||
# very same sigma because `sigma_next` is read off the grid at
|
||||
# `global_idx + 1` rather than off a scheduler counter. That
|
||||
# increment is what `--flow_euler_rollout` trains, and the
|
||||
# trainer's own KV-cache walk pops the same `(span, step_idx,
|
||||
# commit)` stack, so a ladder under a deterministic schedule is
|
||||
# legitimate - and is the rollout that actually matches a
|
||||
# checkpoint distilled with `--flow_euler_rollout`.
|
||||
if local_attn_size != -1:
|
||||
raise ValueError(
|
||||
"Pyramid timestep chunking (`denoise_mode='pyramid'`) requires "
|
||||
"`stochastic_sampling=True`: the buffered step between two "
|
||||
"levels is the schedule's own re-noising.")
|
||||
if getattr(self.transformer.config, 'local_attn_size', -1) != -1:
|
||||
"A per-step partition ladder (`denoise_mode='pyramid'` or "
|
||||
"'full_then_blocks') requires `local_attn_size=-1` (full KV "
|
||||
"cache). A rolling window evicts by `current_start` deltas, "
|
||||
"which a level change moves backwards.")
|
||||
# Forcing-KV (arXiv 2605.09681) rewrites the cache by grouping heads
|
||||
# into static/dynamic histories on a *clean* pass, and it anchors its
|
||||
# `ar_step` on a fixed block size. A ladder breaks both premises: the
|
||||
# partition changes between steps, so neither the per-block clean
|
||||
# commit nor a fixed stride survives it.
|
||||
if self._forcing_kv_enabled():
|
||||
raise ValueError(
|
||||
"Pyramid timestep chunking (`denoise_mode='pyramid'`) requires "
|
||||
"`local_attn_size=-1` (full KV cache). A rolling window "
|
||||
"evicts by `current_start` deltas, which a splitting chunk "
|
||||
"moves backwards.")
|
||||
"A per-step partition ladder (`denoise_mode='pyramid'` or "
|
||||
"'full_then_blocks') cannot be combined with Forcing-KV head "
|
||||
"compression: compression commits a clean pass per block and "
|
||||
"derives its step from a fixed block size, neither of which "
|
||||
"survives a partition that changes between steps.")
|
||||
# Applies to every mode, not just the pyramid. The partition defines one
|
||||
# bidirectional chunk, so a frame must see the whole of its own chunk. A
|
||||
# rolling window is anchored at the chunk *end* (`window_start =
|
||||
# local_end_index - local_attn_size * frame_seqlen`), so once a chunk is
|
||||
# wider than the window its earliest frames cannot attend to its latest
|
||||
# ones and the chunk silently degrades to partial causality. `sink_size`
|
||||
# does not rescue this: it pins the clip start, not the neighbourhood of
|
||||
# the chunk being denoised.
|
||||
# `partitions` is the ladder (a list of levels, each a partition) when
|
||||
# one is on; otherwise the single `chunk_sizes` partition is the only
|
||||
# level. Widest chunk over every level a rollout can reach.
|
||||
levels = partitions if partitions is not None else [chunk_sizes]
|
||||
widest = max(max(level) for level in levels)
|
||||
if local_attn_size != -1 and widest > local_attn_size:
|
||||
raise ValueError(
|
||||
f"Partition breaks the local attention window: the widest chunk "
|
||||
f"is {widest} latent frames but `local_attn_size="
|
||||
f"{local_attn_size}`, so a chunk's first frames cannot attend to "
|
||||
f"its last ones and the intra-chunk bidirectionality that "
|
||||
f"Flex-Forcing 3.1 relies on is lost. Either use "
|
||||
f"`local_attn_size=-1`, or keep every chunk at or below "
|
||||
f"`local_attn_size` frames.")
|
||||
|
||||
# Below is the Flex-Forcing rollout proper. It performs the same setup
|
||||
# the inherited `__call__` does - prompt encoding, timesteps, latents,
|
||||
@@ -490,6 +556,11 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
# copy of the initial noise in which every frame sits at its current
|
||||
# noise level, and which the walk both reads from and writes back into.
|
||||
state = noise.clone()
|
||||
# The stochastic grid carries a trailing element that the re-noising
|
||||
# reads as `timesteps[global_idx + 1]`, so it is not a step and gets
|
||||
# sliced off. The deterministic grid needs no tail of its own: its
|
||||
# transfer reads `scheduler.sigmas`, which `set_timesteps` already
|
||||
# terminates with a 0.
|
||||
denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps
|
||||
last_step = len(denoise_timesteps) - 1
|
||||
block_idx = 0
|
||||
@@ -531,8 +602,10 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
# A splitting span takes a single buffered step; a leaf runs
|
||||
# whatever is left of the schedule to completion. `step_idx` is
|
||||
# where the slice starts inside `timesteps`, which is what lets the
|
||||
# re-noising below read the right (t_i, t_i+1) pair when the walk
|
||||
# resumes a span mid-schedule.
|
||||
# transfer below pick up the right (t_i, t_i+1) pair when the walk
|
||||
# resumes a span mid-schedule - indexed into `timesteps` directly
|
||||
# under a stochastic rollout, resolved by the scheduler from the `t`
|
||||
# it is handed under a deterministic one.
|
||||
schedule = (denoise_timesteps[step_idx:step_idx + 1] if splits
|
||||
else denoise_timesteps[step_idx:])
|
||||
|
||||
@@ -650,6 +723,26 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
# Delegate to the scheduler's own step() so each sampler
|
||||
# (Flow Euler, UniPC, DPM++) applies its multi-step
|
||||
# formula, exactly as the inherited rollout does.
|
||||
#
|
||||
# This holds under a ladder as well, which is not
|
||||
# obvious: the walk revisits the schedule out of order,
|
||||
# but the per-span `_step_index = None` reset above
|
||||
# makes step() re-derive its position by looking the
|
||||
# handed-over `t` up in `scheduler.timesteps` instead of
|
||||
# counting calls, and inside one span the counter then
|
||||
# advances in step with `schedule`. `sigmas` carries the
|
||||
# trailing 0 `set_timesteps` appends, so the leaf's last
|
||||
# step lands exactly on x0.
|
||||
#
|
||||
# Known and accepted: this grid is
|
||||
# `set_timesteps(num_inference_steps)`, which at shift=5
|
||||
# gives [1000, 909.7, 717.3, 24.4], whereas the trainer
|
||||
# distils on [1000, 937.5, 833.3, 625] - the grid
|
||||
# `stochastic_sampling_timesteps` reproduces entry for
|
||||
# entry. A deterministic ladder therefore samples noise
|
||||
# levels this checkpoint was not distilled at. Use
|
||||
# `stochastic_sampling=True` when the samples have to
|
||||
# sit on the trained trajectory.
|
||||
noisy_input = self.scheduler.step(
|
||||
flow_pred, t, noisy_input, return_dict=False)[0]
|
||||
denoised_pred = noisy_input
|
||||
@@ -787,6 +880,7 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
edit_steps: int = 1,
|
||||
shift: float = 5.0,
|
||||
context_noise: int = 0,
|
||||
stochastic_sampling: bool = True,
|
||||
num_frame_per_block: Optional[int] = None,
|
||||
max_sequence_length: int = 512,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
@@ -818,6 +912,12 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
edit_steps: How many *trailing* steps of that schedule to run. Keep
|
||||
it small - editing at a planning timestep would restructure the
|
||||
clip instead of refining it.
|
||||
stochastic_sampling: Which rollout the refinement runs on - the same
|
||||
knob :meth:`__call__` takes. ``True`` re-noises between steps on
|
||||
the hardcoded ``stochastic_sampling_timesteps`` grid. ``False``
|
||||
delegates each step to ``self.scheduler`` (Flow Euler, UniPC,
|
||||
DPM++), which is what a checkpoint distilled with
|
||||
``--flow_euler_rollout`` was trained on.
|
||||
num_frame_per_block: Granularity of the clean-context commit - the
|
||||
same uniform block size the inherited rollout uses, doing the
|
||||
same job. ``None`` commits the whole clip in one bidirectional
|
||||
@@ -875,8 +975,28 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
* (latents.shape[4] // patch_size[2]))
|
||||
seq_len = frame_seq_length * num_latent_frames
|
||||
|
||||
timesteps = stochastic_sampling_timesteps(
|
||||
num_inference_steps, shift, device)
|
||||
# The same schedule dispatch `__call__` makes: the stochastic grid is
|
||||
# the hardcoded few-step one, the deterministic grid comes from the
|
||||
# scheduler so that whichever sampler is configured supplies its own
|
||||
# multi-step formula. `num_inference_steps` may be rewritten by
|
||||
# `retrieve_timesteps` and `step_offset` below is derived from it, so
|
||||
# this has to run first.
|
||||
if stochastic_sampling:
|
||||
timesteps = stochastic_sampling_timesteps(
|
||||
num_inference_steps, shift, device)
|
||||
elif isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
|
||||
timesteps, num_inference_steps = retrieve_timesteps(
|
||||
self.scheduler, num_inference_steps, device, None)
|
||||
elif isinstance(self.scheduler, FlowUniPCMultistepScheduler):
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
|
||||
timesteps = self.scheduler.timesteps
|
||||
elif isinstance(self.scheduler, FlowDPMSolverMultistepScheduler):
|
||||
sampling_sigmas = get_sampling_sigmas(num_inference_steps, shift)
|
||||
timesteps, _ = retrieve_timesteps(
|
||||
self.scheduler, device=device, sigmas=sampling_sigmas)
|
||||
else:
|
||||
timesteps, num_inference_steps = retrieve_timesteps(
|
||||
self.scheduler, num_inference_steps, device, None)
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
self._initialize_kv_cache(
|
||||
@@ -957,18 +1077,27 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
# 3. Denoise the span at the refinement levels only, with the whole clip
|
||||
# visible so it conditions on past *and* future clean tokens - the
|
||||
# any-order half of 4.2, which is what `flex_state` opens here and
|
||||
# step 1 deliberately left closed. This call is always stochastic, so
|
||||
# the schedule is the trailing `edit_steps` of `timesteps[:-1]` and
|
||||
# the step is the schedule's own re-noising. `self.scheduler` takes no
|
||||
# part in this method at all - `timesteps` above comes from
|
||||
# `stochastic_sampling_timesteps`, and neither loop calls
|
||||
# `scheduler.step()` - so there is no scheduler state to reset here.
|
||||
# step 1 deliberately left closed. Either rollout runs the trailing
|
||||
# `edit_steps` of the schedule; they differ in whether the grid
|
||||
# carries a tail. The stochastic one does - the re-noising below reads
|
||||
# it as `timesteps[global_idx + 1]` - so it is sliced off here. The
|
||||
# deterministic one needs no tail of its own, because its transfer
|
||||
# reads `scheduler.sigmas`, which `set_timesteps` already terminates
|
||||
# with a 0.
|
||||
flex_state = {"attn_window": (0, seq_len)}
|
||||
schedule = timesteps[:-1][step_offset:]
|
||||
denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps
|
||||
schedule = denoise_timesteps[step_offset:]
|
||||
span_frames = edit_end - edit_start
|
||||
current_start = edit_start * frame_seq_length
|
||||
noisy_input = noisy_span
|
||||
edited = noisy_span
|
||||
# Only the deterministic branch touches `self.scheduler`, but that one
|
||||
# keeps its own position there and, for the multi-step samplers, a
|
||||
# history of model outputs. Start both fresh, the way `__call__` does at
|
||||
# the head of every span.
|
||||
self.scheduler._step_index = None
|
||||
if hasattr(self.scheduler, "model_outputs"):
|
||||
self.scheduler.model_outputs = []
|
||||
with self.progress_bar(total=len(schedule)) as progress_bar:
|
||||
for local_idx, t in enumerate(schedule):
|
||||
timestep = torch.ones([batch_size, span_frames], device=device,
|
||||
@@ -1042,13 +1171,24 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
|
||||
f"Unexpected flow_pred dim: {flow_pred.dim()}, "
|
||||
f"shape: {flow_pred.shape}")
|
||||
|
||||
global_idx = step_offset + local_idx
|
||||
t_i = (timesteps[global_idx] / 1000).to(weight_dtype)
|
||||
t_i_1 = (timesteps[global_idx + 1] / 1000).to(weight_dtype)
|
||||
edited = noisy_input - flow_pred * t_i
|
||||
noisy_input = (1 - t_i_1) * edited + t_i_1 * torch.randn(
|
||||
edited.shape, dtype=edited.dtype, device=device,
|
||||
generator=generator)
|
||||
if stochastic_sampling:
|
||||
global_idx = step_offset + local_idx
|
||||
t_i = (timesteps[global_idx] / 1000).to(weight_dtype)
|
||||
t_i_1 = (timesteps[global_idx + 1] / 1000).to(weight_dtype)
|
||||
edited = noisy_input - flow_pred * t_i
|
||||
noisy_input = (1 - t_i_1) * edited + t_i_1 * torch.randn(
|
||||
edited.shape, dtype=edited.dtype, device=device,
|
||||
generator=generator)
|
||||
else:
|
||||
# Delegate to the scheduler's own step(), exactly as
|
||||
# `__call__` does. This refinement is a single linear span
|
||||
# rather than a walk, so the position `step()` re-derives
|
||||
# from the `t` it is handed is the one the schedule is at,
|
||||
# and the trailing sigma of 0 lands the last step on x0 -
|
||||
# which is what `output` below writes back into the clip.
|
||||
noisy_input = self.scheduler.step(
|
||||
flow_pred, t, noisy_input, return_dict=False)[0]
|
||||
edited = noisy_input
|
||||
progress_bar.update()
|
||||
|
||||
output = latents.clone()
|
||||
|
||||
@@ -868,13 +868,14 @@ class WanSelfForcingPipeline(DiffusionPipeline):
|
||||
end_idx = start_idx + current_num_frames
|
||||
noisy_input = noise[:, :, start_idx:end_idx]
|
||||
|
||||
# Denoising loop for current block
|
||||
# Reset scheduler state for each block (required for causal generation)
|
||||
# For Euler scheduler, resetting _step_index is sufficient.
|
||||
# For multi-step schedulers (UniPC, DPM++), also clear accumulated model outputs.
|
||||
self.scheduler._step_index = None
|
||||
if hasattr(self.scheduler, 'model_outputs'):
|
||||
self.scheduler.model_outputs = []
|
||||
if not stochastic_sampling and isinstance(self.scheduler, FlowUniPCMultistepScheduler):
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
|
||||
timesteps = self.scheduler.timesteps
|
||||
elif not stochastic_sampling and isinstance(self.scheduler, FlowDPMSolverMultistepScheduler):
|
||||
self.scheduler.set_timesteps(device=device, sigmas=sampling_sigmas)
|
||||
timesteps = self.scheduler.timesteps
|
||||
else:
|
||||
self.scheduler._step_index = None
|
||||
|
||||
denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps
|
||||
with self.progress_bar(total=len(denoise_timesteps)) as progress_bar:
|
||||
|
||||
@@ -17,6 +17,7 @@ from .fp8_optimization import (autocast_model_forward,
|
||||
convert_weight_dtype_wrapper,
|
||||
replace_parameters_by_name)
|
||||
from .fsdp_ema import FSDPEMA, LORAFSDPEMA
|
||||
from .fsdp_utils import checkpoint_fsdp_forward
|
||||
from .gpu_memory_mode import (GPU_MEMORY_MODES, QUANT_MODES,
|
||||
apply_gpu_memory_mode, register_quant_mode,
|
||||
split_gpu_memory_mode)
|
||||
|
||||
@@ -26,7 +26,7 @@ Conventions used across this module:
|
||||
* ``num_frames`` always denotes **latent** frames (post-VAE), matching
|
||||
``target_shape[1]`` in the trainers and ``num_latent_frames`` in the pipeline.
|
||||
"""
|
||||
from typing import List, Optional, Sequence, Tuple, Union
|
||||
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
|
||||
|
||||
import torch
|
||||
|
||||
@@ -45,6 +45,8 @@ __all__ = [
|
||||
"chunk_ends_tensor",
|
||||
"broadcast_chunk_sizes",
|
||||
"chunk_sizes_to_block_kwargs",
|
||||
"FLEX_RANDOM_PROB",
|
||||
"sample_flex_level0",
|
||||
]
|
||||
|
||||
|
||||
@@ -89,7 +91,8 @@ def uniform_chunks(num_frames: int,
|
||||
|
||||
def normalize_chunk_spec(spec: Union[None, int, str, Sequence[int]],
|
||||
num_frames: int,
|
||||
chunk_size: Optional[int] = None) -> ChunkSizes:
|
||||
chunk_size: Optional[int] = None,
|
||||
independent_first_frame: bool = False) -> ChunkSizes:
|
||||
"""Parse a user-facing chunk specification into a validated partition.
|
||||
|
||||
Accepted forms (``num_frames`` = latent frames the partition must cover):
|
||||
@@ -102,18 +105,26 @@ def normalize_chunk_spec(spec: Union[None, int, str, Sequence[int]],
|
||||
* ``"ar"`` -> fully autoregressive ``[1] * num_frames``
|
||||
* ``"bidir"`` / ``"full"`` -> fully bidirectional ``[num_frames]``
|
||||
|
||||
``independent_first_frame`` needs no separate flag: it is exactly the
|
||||
partition whose first chunk has size 1 (read it back with
|
||||
``sizes[0] == 1``).
|
||||
``independent_first_frame`` is applied to the *uniform* forms only (``None``,
|
||||
``int``, ``"uniform[:N]"``), where it inserts the leading size-1 chunk that
|
||||
makes the layout ``[1, N, N, ...]``. Without it those forms produce a
|
||||
partition the ``independent_first_frame`` trainer never trained on. The
|
||||
keyword forms leave it alone on purpose: ``"ar"`` is already one frame per
|
||||
chunk, and prefixing ``"bidir"`` would break the single bidirectional chunk.
|
||||
An explicit size list also leaves it alone - a leading 1 there is written by
|
||||
the caller, and is read back with ``sizes[0] == 1``.
|
||||
"""
|
||||
if spec is None:
|
||||
sizes = uniform_chunks(num_frames, chunk_size if chunk_size is not None else 1)
|
||||
sizes = uniform_chunks(num_frames, chunk_size if chunk_size is not None else 1,
|
||||
independent_first_frame=independent_first_frame)
|
||||
elif isinstance(spec, bool):
|
||||
raise TypeError(f"chunk spec must not be a bool, got {spec!r}")
|
||||
elif isinstance(spec, int):
|
||||
sizes = uniform_chunks(num_frames, spec)
|
||||
sizes = uniform_chunks(num_frames, spec,
|
||||
independent_first_frame=independent_first_frame)
|
||||
elif isinstance(spec, str):
|
||||
sizes = _parse_chunk_string(spec, num_frames)
|
||||
sizes = _parse_chunk_string(spec, num_frames,
|
||||
independent_first_frame=independent_first_frame)
|
||||
elif isinstance(spec, (list, tuple)):
|
||||
sizes = [int(s) for s in spec]
|
||||
else:
|
||||
@@ -124,7 +135,8 @@ def normalize_chunk_spec(spec: Union[None, int, str, Sequence[int]],
|
||||
return _validate_sizes(sizes, num_frames, spec)
|
||||
|
||||
|
||||
def _parse_chunk_string(spec: str, num_frames: int) -> ChunkSizes:
|
||||
def _parse_chunk_string(spec: str, num_frames: int,
|
||||
independent_first_frame: bool = False) -> ChunkSizes:
|
||||
text = spec.strip()
|
||||
lowered = text.lower()
|
||||
if lowered in ("ar", "autoregressive", "causal"):
|
||||
@@ -132,9 +144,11 @@ def _parse_chunk_string(spec: str, num_frames: int) -> ChunkSizes:
|
||||
if lowered in ("bidir", "bidirectional", "full"):
|
||||
return [int(num_frames)]
|
||||
if lowered.startswith("uniform:"):
|
||||
return uniform_chunks(num_frames, int(lowered.split(":", 1)[1]))
|
||||
return uniform_chunks(num_frames, int(lowered.split(":", 1)[1]),
|
||||
independent_first_frame=independent_first_frame)
|
||||
if lowered.startswith("uniform"):
|
||||
return uniform_chunks(num_frames, 1)
|
||||
return uniform_chunks(num_frames, 1,
|
||||
independent_first_frame=independent_first_frame)
|
||||
body = text.strip("[]() ")
|
||||
sep = "," if "," in body else ("-" if "-" in body else None)
|
||||
if sep is None:
|
||||
@@ -159,18 +173,41 @@ def _validate_sizes(sizes: ChunkSizes, num_frames: int, spec) -> ChunkSizes:
|
||||
return [int(s) for s in sizes]
|
||||
|
||||
|
||||
# Share of iterations that pin the launcher's own uniform ``num_frame_per_block``
|
||||
# layout instead of drawing a random partition. ``sample_flexible_chunks`` can
|
||||
# reach every partition whose chunks are individually legal, but it reaches any
|
||||
# *particular* one only by chance: the uniform ``[3] * 7`` at 21 latent frames
|
||||
# has probability 1/183708, i.e. 0.11 expected hits in 20000 iterations. Without
|
||||
# this band ``--num_frame_per_block`` would name a layout the model never trains
|
||||
# on, and the validation pass - which pins exactly that layout - would render
|
||||
# out of distribution. Raising the sampler's randomness makes this *worse*, not
|
||||
# better: more reachable partitions means each one is drawn less often.
|
||||
# Deliberately not a CLI flag, and shared by both trainers so the two stages
|
||||
# cannot drift apart.
|
||||
UNIFORM_BLOCK_PROB = 0.1
|
||||
# Share of iterations the level-0 mixture's *random* arm takes; the other arms
|
||||
# split the remainder equally, so with a pyramid this is 30% coarse / 30%
|
||||
# uniform / 30% full_then_blocks / 10% random. ``sample_flex_level0`` below is
|
||||
# the single source of truth for the coarsest layout that **both** trainers
|
||||
# draw: stage 2 (``train_distill.py``) nests it into the §3.2 ladder, stage 1
|
||||
# (``train_ode.py``) regresses against it directly because it predicts x0 in one
|
||||
# forward per chunk and has no ladder to build. Sharing the draw is what keeps
|
||||
# the two stages on the same layout distribution - stage 1's checkpoint is stage
|
||||
# 2's initialisation (``--ode_transformer_path``), so a layout stage 2 spends
|
||||
# most of its iterations on must not be one stage 1 never regressed.
|
||||
#
|
||||
# The layouts ``--num_frame_per_block`` names are arms, not something the random
|
||||
# draw is expected to stumble onto: ``sample_flexible_chunks`` reaches every
|
||||
# individually legal partition but any *particular* one only by chance. At 21
|
||||
# latent frames the uniform ``[3] * 7`` has probability 1/183708, and the fully
|
||||
# causal ``[1] * 21`` is 1/21! - unreachable in practice at any ``max_chunk``,
|
||||
# which is why no ``--flex_chunk_min/max`` setting can substitute for these arms.
|
||||
# Pinning them is what makes ``--num_frame_per_block`` name a layout the model
|
||||
# genuinely trains on, and keeps ``log_validation`` - which pins exactly that
|
||||
# layout - in distribution. Deliberately not a CLI flag.
|
||||
#
|
||||
# Deliberate deviation from the paper, and the reason this constant exists at
|
||||
# all: arXiv 2607.03509 §5 says "we randomly pick the chunk size in training
|
||||
# from 2 to 10" and pins no arm of its own, so its level 0 is drawn entirely
|
||||
# from that random family. That family spans thousands of distinct layouts, so
|
||||
# the random arm buys exactly what §3.3 asks of it - coverage of "attention
|
||||
# configurations that range from strictly causal to fully non-causal" - but it
|
||||
# cannot train any *particular* layout, and the layouts inference actually
|
||||
# selects (``PAPER_CHUNK_CONFIGS``, plus the one ``log_validation`` pins) are
|
||||
# the deterministic ones. Holding it to a minority share is what leaves those
|
||||
# arms enough iterations to be trained at all. Setting this to 1.0 reproduces
|
||||
# the paper's mixture exactly: every threshold below collapses to 0, so the
|
||||
# three deterministic arms become unreachable and the uniform draw always lands
|
||||
# in the random one.
|
||||
FLEX_RANDOM_PROB = 0.10
|
||||
|
||||
|
||||
def sample_flexible_chunks(num_frames: int,
|
||||
@@ -205,21 +242,13 @@ def sample_flexible_chunks(num_frames: int,
|
||||
# Two legal moves, offered everywhere along the clip including the tail:
|
||||
# take the whole remainder, or leave at least min_chunk frames behind.
|
||||
# Anything in between would strand a sub-min_chunk sliver mid-video.
|
||||
#
|
||||
# The tail used to be closed unconditionally (`if remaining <=
|
||||
# max_chunk: append(remaining); break`), which silently removed most of
|
||||
# the sampler's randomness: at 21 frames with 2..10 a run of 3s stopped
|
||||
# at `[3, 3, 3, 3, 9]` after four draws, so the uniform `[3] * 7` was
|
||||
# unreachable outright - exhaustive enumeration found 679 partitions and
|
||||
# none of them was it. Uniform layouts also bias short: the old rule
|
||||
# topped out at 7 chunks with 65% of draws at just 3.
|
||||
upper = min(max_chunk, remaining - min_chunk)
|
||||
choices = list(range(min_chunk, upper + 1)) if upper >= min_chunk else []
|
||||
if remaining <= max_chunk:
|
||||
choices.append(remaining)
|
||||
if not choices:
|
||||
# Degenerate only: min_chunk == max_chunk and the remainder is
|
||||
# strictly between them, so nothing legal fits. Same as before.
|
||||
# strictly between them, so nothing legal fits.
|
||||
sizes.append(remaining)
|
||||
break
|
||||
step = choices[int(torch.randint(
|
||||
@@ -414,3 +443,172 @@ def chunk_sizes_to_block_kwargs(chunk_sizes: Sequence[int]) -> dict:
|
||||
"num_frame_per_block": max(sizes) if sizes else 1,
|
||||
"independent_first_frame": bool(sizes and sizes[0] == 1),
|
||||
}
|
||||
|
||||
|
||||
def sample_flex_level0(num_frames: int,
|
||||
num_frame_per_block: int,
|
||||
chunk_min: int,
|
||||
chunk_max: int,
|
||||
pyramid: bool,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
device: Union[str, torch.device] = "cpu",
|
||||
independent_first_frame: bool = False) -> Dict[str, Any]:
|
||||
"""Draw one iteration's coarsest Flex-Forcing layout (the four-arm mixture).
|
||||
|
||||
Single source of truth for the level-0 partition **both** trainers draw:
|
||||
``train_distill.py`` nests the result into the §3.2 ladder, ``train_ode.py``
|
||||
regresses against it directly. Keeping the draw here is what stops the two
|
||||
stages from drifting apart - stage 1's checkpoint is stage 2's
|
||||
initialisation, so every layout stage 2 spends iterations on has to be one
|
||||
stage 1 actually regressed.
|
||||
|
||||
Four arms, decided by a single uniform draw. ``"random"`` takes
|
||||
:data:`FLEX_RANDOM_PROB` (10%) and the other three split what is left
|
||||
equally, i.e. 30% each:
|
||||
|
||||
* ``"coarse"`` - one chunk over the whole clip (inference's planning step),
|
||||
i.e. fully bidirectional. Only drawn when ``pyramid`` is on: without a
|
||||
ladder a ``[F]`` level 0 would be a degenerate single-chunk rollout, so
|
||||
the two surviving arms widen to 45% each and ``"random"`` stays at 10%.
|
||||
* ``"uniform"`` - a pure Self-Forcing grid (no binary refinement) in one of
|
||||
three sub-layouts drawn from ``num_frame_per_block`` and fully causal:
|
||||
``[1]*F`` alone, ``[nfpb]*k`` alone, or ``[nfpb]*k`` planning followed by
|
||||
``[1]*F`` for every remaining step - the coarse-block -> per-frame descent
|
||||
the pyramid never samples, because it halves through ``[2, 1x14]`` first.
|
||||
The three are equiprobable *within* the arm, so 10% of iterations each.
|
||||
* ``"full_then_blocks"`` - ``denoise_mode="full_then_blocks"``: a whole-clip
|
||||
planning chunk, then block-major refinement at a width drawn from
|
||||
``{1, num_frame_per_block}``.
|
||||
* ``"random"`` - the §3.3 / §5 spectrum via :func:`sample_flexible_chunks`.
|
||||
|
||||
Distributed contract: the arm selector is drawn once and broadcast from
|
||||
rank 0, then **every** arm issues exactly one :func:`broadcast_chunk_sizes`.
|
||||
All ranks therefore end up with a byte-identical layout *and* the same
|
||||
collective count. A per-rank arm draw would leave ranks in different arms
|
||||
with different walk forward counts -> FSDP all-gather desync (NCCL hang).
|
||||
|
||||
Args:
|
||||
pyramid: whether the caller ladders at all, i.e.
|
||||
``args.flex_pyramid_levels > 1``. Stage 1 has no ladder but must
|
||||
still pass stage 2's setting, because this flag decides whether the
|
||||
coarse arm exists and therefore what the other three arms weigh.
|
||||
|
||||
Returns:
|
||||
``"sizes"`` - the level-0 partition, already rank-reconciled.
|
||||
``"arm"`` - ``"coarse"`` / ``"uniform_causal"`` / ``"uniform_blocks"`` /
|
||||
``"uniform_plan_then_causal"`` / ``"full_then_blocks"`` / ``"random"``.
|
||||
The uniform arm is resolved to its sub-layout here so the caller can
|
||||
rebuild the ladder depth without re-deriving the band arithmetic.
|
||||
``"arm_text"`` - human-readable description for the layout log.
|
||||
``"causal"`` - ``[1]*F``, present on ``"uniform_plan_then_causal"``.
|
||||
``"block_level"`` - the block-major level, present on
|
||||
``"full_then_blocks"``.
|
||||
"""
|
||||
num_frames = int(num_frames)
|
||||
nfpb = max(1, int(num_frame_per_block))
|
||||
# The arm selector must be IDENTICAL on every rank, not just its broadcast
|
||||
# count: the arms build their ladder from different tensors, so a per-rank
|
||||
# `u` would desync the group. Broadcast rank 0's draw; the arm's own single
|
||||
# broadcast then reconciles its within-arm randomness (the random base, the
|
||||
# block width) to rank 0.
|
||||
u = torch.rand((), generator=generator, device=device)
|
||||
try:
|
||||
import torch.distributed as dist
|
||||
except ImportError: # pragma: no cover - torch always ships distributed
|
||||
dist = None
|
||||
if dist is not None and dist.is_available() and dist.is_initialized():
|
||||
dist.broadcast(u.reshape(1), src=0)
|
||||
u = u.item()
|
||||
|
||||
# `random` holds a fixed slice and the deterministic arms divide the rest,
|
||||
# so switching the coarse arm off widens the two survivors instead of
|
||||
# silently fattening `random`. `ftb_hi` is both the end of the last
|
||||
# deterministic arm and the start of the random one, and it is the final
|
||||
# threshold tested, so no rounding remainder can fall through unhandled.
|
||||
det_arms = 3 if pyramid else 2
|
||||
det_w = (1.0 - FLEX_RANDOM_PROB) / det_arms
|
||||
coarse_hi = det_w if pyramid else 0.0
|
||||
uniform_hi = coarse_hi + det_w
|
||||
ftb_hi = uniform_hi + det_w
|
||||
|
||||
if u < coarse_hi:
|
||||
# Coarse end of §3.1: reuse the ladder builder's own level-0 rule so the
|
||||
# `independent_first_frame` handling cannot drift from inference's.
|
||||
sizes = build_pyramid_partitions(
|
||||
num_frames, num_levels=1, base_chunks=None,
|
||||
independent_first_frame=independent_first_frame)[0]
|
||||
return {
|
||||
"sizes": broadcast_chunk_sizes(sizes, device=device),
|
||||
"arm": "coarse",
|
||||
"arm_text": "whole clip, the 3.2 coarse planning layout",
|
||||
}
|
||||
|
||||
if u < uniform_hi:
|
||||
# The three sub-layouts change the ladder DEPTH (hence stage 2's walk
|
||||
# forward count), so the choice must be rank-consistent. It is free here:
|
||||
# `u` was already broadcast, so slicing its position within the uniform
|
||||
# band gives every rank the same sub-arm with no extra collective. Like
|
||||
# every arm this issues EXACTLY ONE broadcast. `uniform_chunks` is called
|
||||
# directly rather than through `normalize_chunk_spec`, which only applies
|
||||
# `independent_first_frame` to its uniform forms via a keyword - the
|
||||
# direct call keeps this arm visibly identical to the bands around it.
|
||||
band = (u - coarse_hi) / det_w # [0, 1) in the uniform arm
|
||||
sub = min(int(band * 3), 2) # 0 / 1 / 2, equiprobable
|
||||
causal = uniform_chunks(num_frames, 1,
|
||||
independent_first_frame=independent_first_frame)
|
||||
if sub == 0:
|
||||
return {
|
||||
"sizes": broadcast_chunk_sizes(causal, device=device),
|
||||
"arm": "uniform_causal",
|
||||
"arm_text": "uniform 1-frame blocks, single level (fully causal)",
|
||||
}
|
||||
sizes = broadcast_chunk_sizes(
|
||||
uniform_chunks(num_frames, nfpb,
|
||||
independent_first_frame=independent_first_frame),
|
||||
device=device)
|
||||
if sub == 1:
|
||||
return {"sizes": sizes, "arm": "uniform_blocks",
|
||||
"arm_text": f"uniform {nfpb}-frame blocks, single level"}
|
||||
return {
|
||||
"sizes": sizes, "arm": "uniform_plan_then_causal",
|
||||
"causal": list(causal),
|
||||
"arm_text": (f"uniform {nfpb}-frame plan -> fully causal [1x"
|
||||
f"{num_frames}] for the remaining steps"),
|
||||
}
|
||||
|
||||
if u < ftb_hi:
|
||||
# "First step full, every later step block-major": a fixed two-level
|
||||
# ladder that does NOT go through the binary pyramid refinement, so the
|
||||
# block-major level is reached at high noise (the 2nd step) instead of as
|
||||
# a deep refinement - the one trajectory the pyramid arms never sample.
|
||||
# The block width is drawn between fully causal (1 frame / block) and the
|
||||
# launcher's `num_frame_per_block`, so both the tight-AR and the
|
||||
# coarse-block refinement of the whole-clip plan get trained. Level 0 is
|
||||
# just `[F]`, identical on every rank and needing no sync, so the ONE
|
||||
# broadcast this arm issues covers the block-major level; rank 0's draw
|
||||
# wins and every rank rebuilds the same two-level ladder locally.
|
||||
block_choices = sorted({1, nfpb})
|
||||
pick = int(torch.rand((), generator=generator, device=device).item()
|
||||
* len(block_choices))
|
||||
block = block_choices[min(pick, len(block_choices) - 1)]
|
||||
block = min(block, num_frames)
|
||||
ftb = build_full_then_blocks_partitions(
|
||||
num_frames, block, independent_first_frame=independent_first_frame)
|
||||
ftb[1] = broadcast_chunk_sizes(ftb[1], device=device)
|
||||
return {
|
||||
"sizes": list(ftb[0]), "arm": "full_then_blocks",
|
||||
"block_level": list(ftb[1]),
|
||||
"arm_text": (f"full clip -> uniform {max(ftb[1])}-frame blocks, "
|
||||
"the full_then_blocks ladder"),
|
||||
}
|
||||
|
||||
sizes = sample_flexible_chunks(
|
||||
num_frames, min_chunk=chunk_min, max_chunk=chunk_max,
|
||||
generator=generator, device=device,
|
||||
independent_first_frame=independent_first_frame)
|
||||
return {
|
||||
"sizes": broadcast_chunk_sizes(sizes, device=device),
|
||||
"arm": "random",
|
||||
"arm_text": (f"random {chunk_min}..{chunk_max}-frame blocks, "
|
||||
"the 3.3 stochastic-chunk spectrum"),
|
||||
}
|
||||
|
||||
@@ -37,7 +37,7 @@ class FSDPEMA:
|
||||
at the cost of one full fp32 copy in host RAM per rank).
|
||||
"""
|
||||
|
||||
def __init__(self, module, source, accelerator, fsdp_plugin, decay=0.9999):
|
||||
def __init__(self, module, source, accelerator, fsdp_plugin, decay=0.99):
|
||||
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
||||
|
||||
self.accelerator = accelerator
|
||||
@@ -56,18 +56,13 @@ class FSDPEMA:
|
||||
with torch.no_grad():
|
||||
module.load_state_dict(source.state_dict(), strict=True)
|
||||
|
||||
# Wrap with the same kwargs as the live model goes through in
|
||||
# `accelerator.prepare_model`, minus the mixed-precision policy: the
|
||||
# EMA never runs forward/backward, and a `param_dtype` cast would
|
||||
# quantize the shadow copy away from fp32 (see the note in the class
|
||||
# docstring), so no casting is wanted here.
|
||||
fsdp_plugin.set_auto_wrap_policy(module)
|
||||
self.module = FSDP(
|
||||
module,
|
||||
sharding_strategy=fsdp_plugin.sharding_strategy or fsdp_plugin.reshard_after_forward,
|
||||
cpu_offload=fsdp_plugin.cpu_offload,
|
||||
auto_wrap_policy=fsdp_plugin.auto_wrap_policy,
|
||||
mixed_precision=None,
|
||||
mixed_precision=fsdp_plugin.mixed_precision_policy,
|
||||
sync_module_states=fsdp_plugin.sync_module_states,
|
||||
backward_prefetch=fsdp_plugin.backward_prefetch,
|
||||
forward_prefetch=fsdp_plugin.forward_prefetch,
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
"""Blockwise autograd replay for memory-bounded FSDP training."""
|
||||
|
||||
from contextlib import nullcontext
|
||||
|
||||
import torch
|
||||
|
||||
__all__ = ["checkpoint_fsdp_forward"]
|
||||
|
||||
|
||||
def checkpoint_fsdp_forward(model, accumulated, *, context_fn=None, **kwargs):
|
||||
"""Replay independent model forwards and accumulate parameter gradients.
|
||||
|
||||
Share a fresh ``accumulated`` dictionary across all blocks of one rollout.
|
||||
Call ``loss.backward()`` once, then clip gradients and update the optimizer.
|
||||
Existing model gradients are cleared during replay; do not mix this helper
|
||||
with other gradient accumulation or direct parameter paths in the same loss.
|
||||
|
||||
Distributed use requires FSDP1 FULL_SHARD with CPUOffload,
|
||||
``use_orig_params=True``, and gradient synchronization enabled. FSDP2,
|
||||
``no_sync``, higher-order gradients, and repeated backward are unsupported.
|
||||
CPU modules can use the same replay semantics without distributed wrapping.
|
||||
|
||||
Inputs are keyword arguments containing tensors or immutable metadata.
|
||||
Tensors must be passed directly, not inside containers. Outputs must be a
|
||||
tensor or a tuple of tensors. All CUDA inputs must use one device.
|
||||
``context_fn`` optionally snapshots mutable forward state and returns
|
||||
``(forward_context, recompute_context)``, as in non-reentrant checkpointing.
|
||||
The model must not mutate parameters or buffers during replay.
|
||||
"""
|
||||
keys = tuple(kwargs)
|
||||
values = tuple(kwargs.values())
|
||||
tensors = [value for value in values if isinstance(value, torch.Tensor)]
|
||||
if not tensors:
|
||||
raise ValueError("FSDP replay requires at least one tensor input.")
|
||||
device = next((value.device for value in tensors if value.is_cuda), tensors[0].device)
|
||||
if device.type not in ("cpu", "cuda"):
|
||||
raise ValueError(f"FSDP replay supports CPU and CUDA inputs, got {device}.")
|
||||
if any(value.is_cuda and value.device != device for value in tensors):
|
||||
raise ValueError("FSDP replay requires all CUDA inputs on one device.")
|
||||
cuda_devices = [device.index] if device.type == "cuda" else []
|
||||
|
||||
class Replay(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, anchor, *args):
|
||||
ctx.tensor_indices = [index for index, value in enumerate(args) if isinstance(value, torch.Tensor)]
|
||||
ctx.other_args = {index: value for index, value in enumerate(args) if not isinstance(value, torch.Tensor)}
|
||||
tensors = [args[index] for index in ctx.tensor_indices]
|
||||
ctx.input_requires_grad = [value.requires_grad for value in tensors]
|
||||
ctx.save_for_backward(*tensors)
|
||||
forward_context, ctx.replay_context = context_fn() if context_fn is not None else (nullcontext(), nullcontext())
|
||||
ctx.cpu_rng = torch.get_rng_state()
|
||||
ctx.cuda_rng = torch.cuda.get_rng_state(device) if cuda_devices else None
|
||||
ctx.autocast = torch.is_autocast_enabled(device.type)
|
||||
ctx.autocast_dtype = torch.get_autocast_dtype(device.type)
|
||||
ctx.set_materialize_grads(False)
|
||||
with forward_context:
|
||||
outputs = model(**dict(zip(keys, args)))
|
||||
if not isinstance(outputs, torch.Tensor) and not (
|
||||
isinstance(outputs, tuple) and all(isinstance(output, torch.Tensor) for output in outputs)
|
||||
):
|
||||
raise TypeError("FSDP replay requires a tensor or a tuple of tensors as output.")
|
||||
return outputs
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *output_grads):
|
||||
parameters = list(model.parameters())
|
||||
# Clear gradient slots; the separate accumulator preserves earlier block gradients.
|
||||
previous_grads = [accumulated.get(parameter) for parameter in parameters]
|
||||
model.zero_grad(set_to_none=True)
|
||||
tensors = [value.detach().requires_grad_(requires_grad)
|
||||
for value, requires_grad in zip(ctx.saved_tensors, ctx.input_requires_grad)]
|
||||
args = dict(ctx.other_args)
|
||||
args.update(zip(ctx.tensor_indices, tensors))
|
||||
with torch.random.fork_rng(devices=cuda_devices), torch.enable_grad(), \
|
||||
torch.autocast(device.type, enabled=ctx.autocast, dtype=ctx.autocast_dtype):
|
||||
torch.set_rng_state(ctx.cpu_rng)
|
||||
if cuda_devices:
|
||||
torch.cuda.set_rng_state(ctx.cuda_rng, device)
|
||||
with ctx.replay_context:
|
||||
outputs = model(**{key: args[index] for index, key in enumerate(keys)})
|
||||
outputs = (outputs,) if isinstance(outputs, torch.Tensor) else outputs
|
||||
active = [(output, grad) for output, grad in zip(outputs, output_grads)
|
||||
if grad is not None and output.requires_grad]
|
||||
if active:
|
||||
torch.autograd.backward([output for output, _ in active], [grad for _, grad in active])
|
||||
# FSDP completes reduce-scatter and offload before accumulating local shards.
|
||||
for parameter, previous in zip(parameters, previous_grads):
|
||||
if previous is not None:
|
||||
if parameter.grad is None:
|
||||
parameter.grad = previous
|
||||
else:
|
||||
parameter.grad.add_(previous)
|
||||
if parameter.grad is not None:
|
||||
accumulated[parameter] = parameter.grad.detach().clone()
|
||||
result = [None] * len(keys)
|
||||
for index, tensor in zip(ctx.tensor_indices, tensors):
|
||||
result[index] = tensor.grad
|
||||
return (None, *result)
|
||||
|
||||
with torch.autograd.graph.save_on_cpu():
|
||||
return Replay.apply(torch.empty(0, device=device, requires_grad=True), *values)
|
||||
Reference in New Issue
Block a user