Update one forcing

Update multigpus in vsa, sla, spark
Update one-forcing and causal forcing params
2026-11-02 00:11:06 +08:00 · 2026-11-01 10:00:44 +08:00 · 2026-10-30 16:33:15 +08:00 · 2026-10-29 10:59:15 +08:00 · 2026-10-28 11:34:50 +08:00 · 2026-10-27 23:16:10 +08:00
113 changed files with 37667 additions and 1429 deletions
@@ -97,6 +97,12 @@ class LoadZImageTransformerModel:

        model_path = folder_paths.get_full_path("diffusion_models", model_name)
        transformer_state_dict = load_torch_file(model_path, safe_load=True)
+        # Some ComfyUI-native checkpoints store every diffusion-model tensor under a
+        # `model.diffusion_model.` prefix (e.g. CivitAI Z-Image fine-tunes). Strip it so
+        # the conversion gate below and `convert_state_dict` receive bare keys.
+        _prefix = "model.diffusion_model."
+        if transformer_state_dict and all(k.startswith(_prefix) for k in transformer_state_dict):
+            transformer_state_dict = {k[len(_prefix):]: v for k, v in transformer_state_dict.items()}
        def convert_state_dict(old_state_dict):
            new_state_dict = {}
            for key, value in old_state_dict.items():
@@ -0,0 +1,338 @@
+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
+# The FastVideo-H3 VSA student supports sequence-parallel inference through ulysses head-parallel
+# only (the VSA kernel has no ring-attention path): put every GPU into ulysses_degree, keep
+# ring_degree = 1, and ulysses_degree must divide the 56 attention heads (2 / 4 / 7 / 8 / ...).
+# The defaults stay single-GPU because the ~19B stack still needs an offloading mode below.
+ulysses_degree      = 1
+ring_degree         = 1
+# Use FSDP to save more GPU memory in multi gpus. Unused at ulysses_degree = ring_degree = 1.
+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:
+    # The VSA sparse kernel only supports ulysses head-parallel (no ring-attention path); a VSA
+    # student under ring_degree>1 would silently fold ring into a larger ulysses degree and give
+    # none of ring's per-rank memory saving, so reject it instead of mis-attending.
+    if ring_degree > 1:
+        raise ValueError(
+            "FastVideo-H3 VSA multi-GPU requires ring_degree=1; put every GPU into ulysses_degree "
+            "instead (ulysses_degree must divide the 56 attention heads).")
+    from functools import partial
+    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,319 @@
+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:
+    # The VSA kernel has no ring-attention path; a sparse student under ring_degree>1 would
+    # silently mis-attend. student_topk >= 1.0 is the dense fallback, which ring does support.
+    if ring_degree > 1 and student_topk < 1.0:
+        raise ValueError(
+            "FastVideo VSA multi-GPU requires ring_degree=1; put every GPU into ulysses_degree "
+            "instead (num_heads must divide ulysses_degree).")
+    from functools import partial
+    transformer.enable_multi_gpus_inference()
+    if fsdp_dit:
+        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,305 @@
+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)
+        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()
@@ -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()
@@ -213,6 +213,12 @@ pipeline = Wan2_2I2VPipeline(
    scheduler=scheduler,
 )
 if ulysses_degree > 1 or ring_degree > 1:
+    # The SLA sparse kernel only supports ulysses head-parallel (no ring-attention path); a
+    # sparse expert under ring_degree>1 would silently mis-attend.
+    if ring_degree > 1:
+        raise ValueError(
+            "TurboDiffusion SLA multi-GPU requires ring_degree=1; put every GPU into "
+            "ulysses_degree instead (num_heads must divide ulysses_degree).")
    from functools import partial
    transformer.enable_multi_gpus_inference()
    transformer_2.enable_multi_gpus_inference()
@@ -289,7 +295,6 @@ def save_results():
        image = (image * 255).numpy().astype(np.uint8)
        image = Image.fromarray(image)
        image.save(video_path)
-        print(f"Saved image to: {video_path}")
    else:
        video_path = os.path.join(save_path, prefix + ".mp4")
        save_videos_grid(sample, video_path, fps=fps)
@@ -67,7 +67,7 @@ sampler_name        = "Flow"
 # Used when the sampler is in "Flow_Unipc", "Flow_DPM++".
 # If you want to generate a 480p video, it is recommended to set the shift value to 3.0.
 # If you want to generate a 720p video, it is recommended to set the shift value to 5.0.
-shift               = 3 
+shift               = 5

 # Load pretrained model if need
 # Download from https://huggingface.co/TurboDiffusion/TurboWan2.1-T2V-1.3B-480P/
@@ -172,6 +172,12 @@ pipeline = WanPipeline(
    scheduler=scheduler,
 )
 if ulysses_degree > 1 or ring_degree > 1:
+    # The SLA sparse kernel only supports ulysses head-parallel (no ring-attention path); a
+    # sparse student under ring_degree>1 would silently mis-attend.
+    if ring_degree > 1:
+        raise ValueError(
+            "TurboDiffusion SLA multi-GPU requires ring_degree=1; put every GPU into "
+            "ulysses_degree instead (num_heads must divide ulysses_degree).")
    from functools import partial
    transformer.enable_multi_gpus_inference()
    if fsdp_dit:
@@ -234,7 +240,6 @@ def save_results():
        image = (image * 255).numpy().astype(np.uint8)
        image = Image.fromarray(image)
        image.save(video_path)
-        print(f"Saved image to: {video_path}")
    else:
        video_path = os.path.join(save_path, prefix + ".mp4")
        save_videos_grid(sample, video_path, fps=fps)
@@ -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,303 @@
+import os
+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,
+                               WanT5EncoderModel,
+                               WanTransformer3DModel_SelfForcing)
+from videox_fun.pipeline import WanSelfForcingPipeline
+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_full_load"
+# 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.
+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
+
+# Config and model path
+config_path         = "config/wan2.1/wan_civitai.yaml"
+# One-Forcing distills a Wan2.1-T2V-1.3B student (the 14B model is only the DMD teacher at training time).
+model_name          = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+
+# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++".
+sampler_name        = "Flow"
+# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics.
+# Used when the sampler is in "Flow_Unipc", "Flow_DPM++".
+shift               = 5.0
+# One-Forcing is a DMD-distilled student, so it uses stochastic sampling.
+stochastic_sampling = True
+
+# One-Forcing checkpoint to overlay on top of the Wan2.1 base model. Point at the trainer
+# output checkpoint dir (its live/EMA weights are auto-resolved below) or a plain `.pt`/`.safetensors` file.
+# Default: the official One-Forcing checkpoint, downloaded via
+#   hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+# The GAN discriminator head is a training-only submodule of fake_score, so it never appears in
+# the generator checkpoint and inference needs no discriminator wiring.
+transformer_path    = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
+# Load the EMA generator (transformer_ema/) when True, otherwise the live generator weights.
+# Only applies to trainer-output checkpoint dirs; a plain `.pt`/`.safetensors` file is loaded as-is.
+use_ema             = False
+vae_path            = None
+lora_path           = None
+
+# Other params
+sample_size         = [480, 832]
+video_length        = 81
+fps                 = 16
+
+# One-Forcing causal inference config
+# `num_frame_per_block`: 1 = frame-wise (One-Forcing default), 3 = chunk-wise.
+num_frame_per_block     = 1
+# Local attention window size (-1 for global attention)
+local_attn_size         = -1
+# Others
+independent_first_frame = False
+context_noise           = 0.0
+
+# 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 stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about."
+negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+# -------- One-Forcing step selector --------
+# One-Forcing bakes CFG into the student weights, so inference MUST use guidance_scale=1.0.
+# The official recipe generates each frame in a single step (denoising_step_list=[1000]):
+#   1-step (One-Forcing default): num_inference_steps = 1
+#   4-step (more stable, slower): num_inference_steps = 4 -> [1000, 750, 500, 250]
+# With stochastic_sampling=True, num_inference_steps=1 yields the warped single timestep [1000],
+# matching the official framewise 1-step generation.
+guidance_scale      = 1.0
+num_inference_steps = 1
+seed                = 43
+lora_weight         = 0.55
+save_path           = "samples/wan-videos-one-forcing"
+
+device = set_multi_gpus_devices(ulysses_degree, ring_degree)
+config = OmegaConf.load(config_path)
+
+# Load transformer with causal inference support if enabled
+transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs'])
+transformer_additional_kwargs['local_attn_size'] = local_attn_size
+
+transformer = WanTransformer3DModel_SelfForcing.from_pretrained(
+    os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+    transformer_additional_kwargs=transformer_additional_kwargs,
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+
+if transformer_path is not None:
+    def _resolve_transformer_path(raw_path: str, prefer_ema: bool) -> str:
+        """Resolve a checkpoint path to the actual weights file to load.
+
+        File paths are returned unchanged so external `.pt` ckpts keep working. For trainer-output
+        dirs, prefer the EMA generator when asked and available, otherwise fall back to the live
+        generator. The FSDP trainer writes the EMA under `transformer_ema/`; `ema_transformer/` is
+        kept as a compatibility alias.
+        """
+        if os.path.isfile(raw_path):
+            return raw_path
+        if os.path.isdir(raw_path):
+            candidates = []
+            if prefer_ema:
+                candidates.append(os.path.join(raw_path, "transformer_ema", "diffusion_pytorch_model.safetensors"))
+                candidates.append(os.path.join(raw_path, "ema_transformer", "diffusion_pytorch_model.safetensors"))
+            candidates.append(os.path.join(raw_path, "transformer", "diffusion_pytorch_model.safetensors"))
+            candidates.append(os.path.join(raw_path, "diffusion_pytorch_model.safetensors"))
+            for c in candidates:
+                if os.path.isfile(c):
+                    return c
+        raise FileNotFoundError(
+            f"transformer_path={raw_path!r} is neither a file nor a checkpoint dir "
+            f"with a known safetensors layout (transformer_ema/, transformer/ or root)."
+        )
+
+    _raw_transformer_path = transformer_path
+    transformer_path = _resolve_transformer_path(transformer_path, prefer_ema=use_ema)
+    if transformer_path != _raw_transformer_path:
+        print(f"use_ema={use_ema}: resolved {_raw_transformer_path} -> {transformer_path}")
+    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 "state_dict" in state_dict else state_dict
+    state_dict = state_dict["generator_ema"] if "generator_ema" in state_dict else state_dict
+    state_dict = state_dict["generator"] if "generator" in state_dict else state_dict
+    # FSDP-saved ckpts keep the `model._fsdp_wrapped_module.` prefix; strip it before the generic
+    # `model.` strip so both FSDP and plain ckpts land at bare parameter names.
+    if any("._fsdp_wrapped_module." in k for k in state_dict.keys()):
+        state_dict = {k.replace("model._fsdp_wrapped_module.", "model.", 1) if k.startswith("model._fsdp_wrapped_module.") else k: v for k, v in state_dict.items()}
+    if any(k.startswith("model.") for k in state_dict.keys()):
+        state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()}
+
+    m, u = transformer.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# 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 = WanSelfForcingPipeline(
+    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.
+# The order lives inside the helper: quantization has to happen before the offload hooks
+# are installed.
+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
+
+    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,
+        num_frame_per_block     = num_frame_per_block,
+        independent_first_frame = independent_first_frame,
+        context_noise           = context_noise,
+        stochastic_sampling     = stochastic_sampling,
+    ).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,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
@@ -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`
@@ -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"
@@ -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"
@@ -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"
@@ -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"
@@ -128,7 +128,7 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun

 ```bash
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
-export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 ```

 **Absolute Path**:
@@ -173,42 +173,50 @@ wget https://github.com/madebyollin/taehv/raw/main/taew2_2.safetensors
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
-export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # Optional: warm-start from released TAE weights instead of training from scratch.
 export TAE_PATH="models/Diffusion_Transformer/taew2_2.safetensors"
 NCCL_DEBUG=INFO

 accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
  --config_path="config/wan2.2/wan_civitai_5b.yaml" \
-  --pretrained_model_name_or_path=$MODEL_NAME \
-  --tae_path=$TAE_PATH \
-  --train_data_dir=$DATASET_NAME \
-  --train_data_meta=$DATASET_META_NAME \
-  --video_sample_size=512 \
-  --video_sample_stride=1 \
-  --video_sample_n_frames=33 \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --tae_path="$TAE_PATH" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --fix_sample_size 720 1280 \
+  --video_sample_strides 1 2 3 \
+  --video_sample_n_frames=25 \
+  --validation_n_frames=49 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
-  --dataloader_num_workers=4 \
-  --num_train_epochs=100 \
-  --checkpointing_steps=500 \
+  --dataloader_num_workers=8 \
+  --max_train_steps=10000 \
+  --checkpointing_steps=1000 \
+  --validation_steps=1000 \
+  --validation_paths "asset/inpaint_video.mp4" \
  --learning_rate=1e-04 \
  --lr_scheduler="constant_with_warmup" \
-  --lr_warmup_steps=100 \
+  --lr_warmup_steps=200 \
  --latent_loss_weight=1.0 \
+  --use_ema \
+  --ema_decay=0.999 \
  --seed=42 \
  --output_dir="output_dir_taehv_w2.2" \
  --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
  --adam_weight_decay=1e-2 \
  --adam_epsilon=1e-08 \
  --vae_mini_batch=1 \
  --max_grad_norm=1.0 \
-  --random_hw_adapt \
  --enable_bucket \
  --trainable_modules "." \
  --resume_from_checkpoint=latest
 ```

+These settings are also provided in `scripts/taehv/train_taehv.sh`. Edit its training arguments and run `bash scripts/taehv/train_taehv.sh`.
+
 **Wan2.1 / Wan2.2 14B (taew2_1) Example**:

 Same as above, with these changes:
@@ -237,15 +245,17 @@ export TAE_PATH="taew2_1.safetensors"
 | `--use_taehv_sequential` | Run the TAE in O(1)-memory sequential mode instead of parallel | - |
 | `--latent_loss_weight` | Weight of the latent MSE (TAE latent vs. full VAE latent) relative to pixel L1 | 1.0 |
 | `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` |
-| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata.json` |
+| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` |
 | `--train_batch_size` | Batch size (per device) | 1 |
 | `--video_sample_size` | Training resolution | 512 |
 | `--video_sample_stride` | Video sample stride | 1 |
+| `--video_sample_strides` | Space-separated stride candidates; one is drawn at random per clip (overrides `--video_sample_stride`) | `None` |
 | `--video_sample_n_frames` | Number of frames to sample. **Must be 4k+1** (33, 49, 81, ...) | 33 |
 | `--vae_mini_batch` | Mini batch size for teacher VAE encoding | 1 |
 | `--gradient_accumulation_steps` | Gradient accumulation steps | 1 |
 | `--dataloader_num_workers` | Number of DataLoader workers | 4 |
 | `--num_train_epochs` | Number of training epochs | 100 |
+| `--max_train_steps` | Total training steps; overrides `--num_train_epochs` when set | 10000 |
 | `--checkpointing_steps` | Save a checkpoint every N steps | 500 |
 | `--checkpoints_total_limit` | Max number of checkpoints to store | `None` |
 | `--learning_rate` | Initial learning rate | 1e-4 |
@@ -253,22 +263,26 @@ export TAE_PATH="taew2_1.safetensors"
 | `--lr_warmup_steps` | Learning rate warmup steps | 100 |
 | `--use_8bit_adam` / `--use_came` | Alternative optimizers | - |
 | `--use_ema` | Keep an EMA copy of the TAE (used for validation and final save) | - |
-| `--seed` | Random seed | 42 |
+| `--ema_decay` | EMA decay; also disables the EMA warmup ramp. Defaults to the `EMAModel` default (0.9999) | `None` |
+| `--seed` | Random seed; defaults to 42 so a bare CLI run stays reproducible | 42 |
 | `--output_dir` | Output directory | `output_dir_taehv_w2.2` |
 | `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
 | `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
 | `--enable_bucket` | Enable bucket training without cropping, groups by resolution | - |
-| `--random_hw_adapt` | Randomly scale videos to a range of resolutions | - |
+| `--fix_sample_size` | Fixed `[height, width]` canvas overriding the square `--video_sample_size` (aspect-preserving cover crop) | `None` |
 | `--low_vram` | Keep the teacher VAE on CPU and move it to GPU only when encoding | - |
 | `--trainable_modules` | Trainable modules (`"."` means all modules) | `"."` |
 | `--trainable_modules_low_learning_rate` | Trainable modules with lr/2 | `[]` |
 | `--resume_from_checkpoint` | Resume training from checkpoint, use `"latest"` to auto-select | `latest` |
 | `--validation_steps` / `--validation_epochs` | Run validation every N steps / epochs | 2000 / 5 |
-| `--validation_paths` | Video paths for validation reconstruction comparison | `"asset/1.mp4"` |
+| `--validation_paths` | Video paths for validation reconstruction comparison | `"asset/inpaint_video.mp4"` |
+| `--validation_n_frames` | Validation clip length (4k+1); defaults to `--video_sample_n_frames` | `None` |

 **Sample Size Configuration Guide**:
- `video_sample_size` represents the training resolution; when `random_hw_adapt` is enabled, it represents the minimum resolution.
- `video_sample_n_frames` must satisfy `4k+1` (e.g. 33, 49, 81) because both the full VAE and the TAE are causal 4x temporal compressors.
+- `video_sample_size` sets a square training canvas; pair it with `--enable_bucket` to train on native aspect ratios.
+- `--fix_sample_size 720 1280` pins a fixed non-square canvas (aspect-preserving cover resize + center crop); with `--enable_bucket` it is applied per batch and `video_sample_size` is bumped to `max(fix_sample_size)`.
+- `--video_sample_n_frames` must satisfy `4k+1` (e.g. 33, 49, 81) because both the full VAE and the TAE are causal 4x temporal compressors.
+- `--video_sample_strides 1 2 3` draws a temporal stride at random for every clip (multi-scale temporal augmentation); the frame count requested shrinks as the stride grows, so the time span a clip covers stays roughly the same.

 ### 3.4 Training Validation

@@ -278,16 +292,19 @@ You can configure validation parameters to periodically reconstruct test videos
 |-----------|-------------|-------------------|
 | `--validation_steps` | Run validation every N steps | 2000 |
 | `--validation_epochs` | Run validation every N epochs | 5 |
-| `--validation_paths` | Validation video paths | `"asset/1.mp4"` |
+| `--validation_paths` | Validation video paths | `"asset/inpaint_video.mp4"` |
+| `--validation_n_frames` | Validation clip length, **4k+1**; omit to reuse `--video_sample_n_frames` | 81 |

 ```bash
-  --validation_paths "asset/1.mp4" \
+  --validation_paths "asset/inpaint_video.mp4" \
  --validation_steps=2000 \
-  --validation_epochs=5
+  --validation_epochs=5 \
+  --validation_n_frames=81
 ```

 **Notes**:
 - Validation saves three videos per sample into `output_dir/sample/`: `*_input.mp4` (resized input), `*_taehv.mp4` (TAE reconstruction), `*_fullvae.mp4` (full VAE reconstruction for reference).
+- Validation clips are truncated to `min(video length, --validation_n_frames)` and then aligned down to `4k+1`, so a longer setting never fails on short assets.
 - When `--use_ema` is enabled, validation runs with the EMA weights.

 ### 3.5 Training Tips
@@ -297,7 +314,7 @@ You can configure validation parameters to periodically reconstruct test videos
 - **Decoder-only distillation**: Add `--freeze_tae_encoder` to only improve decoding quality.
 - **Memory**: The teacher full VAE is the main memory consumer. Use `--low_vram` to keep it on CPU between encoding steps, or reduce `--video_sample_n_frames` / `--video_sample_size`.
 - **Sequential TAE**: `--use_taehv_sequential` trades speed for O(1) activation memory w.r.t. video length.
- **EMA**: `--use_ema` is recommended for the final deliverable; the final saved `taehv` directory contains the EMA weights.
+- **EMA**: `--use_ema` is recommended for the final deliverable; the final saved `taehv` directory contains the EMA weights. Pass `--ema_decay` (e.g. `0.999`, as both `train_taehv.sh` and `train_taehv_super.sh` do) when short runs need the average to catch up faster than the 0.9999 default.

 ### 3.6 Multi-Node Distributed Training

@@ -0,0 +1,433 @@
+# TAE Super (Tiny AutoEncoder) Training Guide
+
+This document provides a complete workflow for training / fine-tuning the Wan2.2
+Super decoder (`AutoencoderTinyWan`), including environment setup, data
+preparation, training, and using the trained checkpoint for inference.
+
+> **Note**: Training uses the frozen official Wan2.2 encoder and updates only the
+> Super decoder's approximately 19M parameters. The tiny encoder is frozen and is
+> not used in the training forward pass.
+>
+> ```text
+> x (video, [-1, 1])
+>   -> official Wan2.2 encoder (frozen) : z = vae.encode(x).latent_dist.mode()
+>   -> Super decoder                  : x_hat = tae(z).sample
+> loss = 0.1 L1 + 0.1 LPIPS + 0.02 DINO + 1.0 grid
+> ```
+>
+> The losses compare reconstruction and input after mapping RGB to [0,1]. The
+> decoder learns to consume the official Wan latent space directly.
+
+The script supports the Wan2.2 5B Super model:
+
+| Family | Latent | Official full VAE | `--config_path` | Models |
+|--------|--------|-------------------|-----------------|--------|
+| taew2_2_super | 48ch, patch_size=2 | `AutoencoderKLWan3_8` | `config/wan2.2/wan_civitai_5b.yaml` | Wan2.2 TI2V-5B / Fun-2.2VAE |
+
+**Differences from [train_taehv.py](train_taehv.py)**:
+
+| Item | `train_taehv.py` | `train_taehv_super.py` |
+|------|-----------------|-----------------------|
+| Reconstruction path | Tiny encoder → tiny decoder; official encoder supplies latent targets | Frozen official encoder → Super decoder |
+| Trainable parameters | Tiny encoder and decoder by default, controlled by `--trainable_modules` | Super decoder only; official and tiny encoders are frozen |
+| Losses | Pixel L1 + latent MSE | Pixel L1 + LPIPS + DINO + grid |
+
+---
+
+## Table of Contents
+
+- [1. Environment Setup](#1-environment-setup)
+- [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. TAE Super Training](#3-tae-super-training)
+  - [3.1 Download Pretrained Model](#31-download-pretrained-model)
+  - [3.2 Quick Start](#32-quick-start)
+  - [3.3 Training Parameter Reference](#33-training-parameter-reference)
+  - [3.4 Training Validation](#34-training-validation)
+  - [3.5 Training Tips](#35-training-tips)
+  - [3.6 Further Tuning](#36-further-tuning)
+- [4. Inference Testing](#4-inference-testing)
+  - [4.1 Checkpoint Layout](#41-checkpoint-layout)
+  - [4.2 Use the Trained TAE in Predict Scripts](#42-use-the-trained-tae-in-predict-scripts)
+- [5. Additional Resources](#5-additional-resources)
+
+---
+
+## 1. Environment Setup
+
+Run the commands from the VideoX-Fun repository root. Install dependencies into
+your Python environment with pip.
+
+**Option 1: Using requirements.txt**
+
+```bash
+pip install -r requirements.txt
+pip install lpips torchvision av modelscope opencv-python
+```
+
+**Option 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 lpips torchvision av modelscope opencv-python
+```
+
+LPIPS's AlexNet backbone and DINOv2 download automatically when missing and reuse
+the PyTorch cache. LPIPS's calibration weights ship with the package. For offline
+DINOv2 loading, use `--dino_repo` and `--dino_weights`.
+
+> The script supports a single device and ordinary DDP through Accelerate. The
+> official Wan VAE and the frozen LPIPS/DINOv2 feature networks remain on the GPU.
+
+---
+
+## 2. Data Preparation
+
+### 2.1 Quick Test Dataset
+
+We provide a test dataset containing several training samples.
+
+```bash
+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 format):
+```json
+[
+  {
+    "file_path": "train/video001.mp4",
+    "text": "A beautiful sunset over the ocean, golden hour lighting",
+    "type": "video",
+    "width": 1024,
+    "height": 1024
+  }
+]
+```
+
+**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 (not used by the TAE loss, kept for meta format compatibility)
+- `type`: Data type, should be `"video"`
+- `width` / `height`: Video width and height (recommended metadata). The Super
+  training entry uses a fixed resize/crop canvas.
+
+### 2.4 Relative vs Absolute Path Usage
+
+**Relative Path**:
+
+```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**:
+
+```bash
+export DATASET_NAME=""
+export DATASET_META_NAME="/mnt/data/metadata.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. TAE Super Training
+
+### 3.1 Download Pretrained Model
+
+The training script needs the official Wan2.2 VAE weights, which ship inside the
+model directory. Download the released Super weights for fine-tuning:
+
+```bash
+mkdir -p models/Diffusion_Transformer
+
+modelscope download --model Wan-AI/Wan2.2-TI2V-5B --local_dir models/Diffusion_Transformer/Wan2.2-TI2V-5B
+
+wget -O models/Diffusion_Transformer/taew2_2_super.pth \
+  https://github.com/madebyollin/taehv/raw/main/taew2_2_super.pth
+```
+
+LPIPS and DINOv2 weights are loaded separately by the loss implementation.
+
+### 3.2 Quick Start
+
+**Wan2.2 TI2V-5B / Fun-2.2VAE (taew2_2_super) Example**:
+
+The launcher uses a 720x1280 canvas, 25 training frames, and EMA, with no
+`--max_train_steps`: training runs for `--num_train_epochs` (100 by default)
+epochs until you kill the processes. Set the GPU/process count with
+`accelerate config`; four GPUs with batch size 1 and accumulation 1 give an
+effective batch of 4.
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth"
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
+  --config_path="config/wan2.2/wan_civitai_5b.yaml" \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --tae_path="$TAE_PATH" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --fix_sample_size 720 1280 \
+  --video_sample_strides 1 2 3 \
+  --video_sample_n_frames=25 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --checkpointing_steps=1000 \
+  --learning_rate=5e-06 \
+  --lr_scheduler="cosine" \
+  --lr_warmup_steps=200 \
+  --lr_min_ratio=0.2 \
+  --seed=42 \
+  --output_dir="output_dir_taehv_w2.2_super" \
+  --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
+  --adam_weight_decay=1e-2 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=1.0 \
+  --l1_loss_weight=0.1 \
+  --lpips_loss_weight=0.1 \
+  --dino_loss_weight=0.02 \
+  --grid_loss_weight=1.0 \
+  --grid_period=16 \
+  --feature_batch_size=13 \
+  --use_ema \
+  --ema_decay=0.999 \
+  --selection_lpips_ratio=1.02 \
+  --selection_grid_ratio=1.05
+```
+
+These settings are also provided in `scripts/taehv/train_taehv_super.sh`. Edit its
+training arguments and run `bash scripts/taehv/train_taehv_super.sh`.
+
+**Train from scratch**:
+
+```bash
+accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --from_scratch \
+  --fix_sample_size 720 1280 \
+  --video_sample_n_frames=25 \
+  --video_sample_strides 1 2 3 \
+  --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
+  --use_ema --ema_decay=0.999 \
+  --dataloader_num_workers=8 --checkpointing_steps=1000 \
+  --output_dir="output_dir_taehv_w2.2_super_scratch"
+```
+
+Omit `--tae_path` when using `--from_scratch`. Only the decoder is trained; the
+saved tiny encoder remains randomly initialized in a scratch run.
+
+**Resume interrupted training**:
+
+```bash
+accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --fix_sample_size 720 1280 --video_sample_n_frames=25 \
+  --video_sample_strides 1 2 3 \
+  --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
+  --use_ema --ema_decay=0.999 \
+  --dataloader_num_workers=8 --checkpointing_steps=1000 \
+  --output_dir="output_dir_taehv_w2.2_super" \
+  --resume_from_checkpoint=latest
+```
+
+Restore with the same settings used to create the checkpoint. Without
+`--max_train_steps`, the schedule horizon derives from `--num_train_epochs`;
+keep the original value when resuming. The checkpoint restores optimizer,
+scheduler, EMA, RNG, step and batch cursor; changing workers or sampling
+settings can change the resumed sample sequence. Use a new output directory
+for a new run.
+
+### 3.3 Training Parameter Reference
+
+Values below describe the launcher configuration; they are not all Python CLI defaults.
+
+| Parameter | Description | Example Value |
+|-----------|-------------|---------------|
+| `--config_path` | Official Wan VAE component configuration | `config/wan2.2/wan_civitai_5b.yaml` |
+| `--pretrained_model_name_or_path` | Model directory containing the official VAE weights | `models/Diffusion_Transformer/Wan2.2-TI2V-5B` |
+| `--tae_path` | Super weights, file or Diffusers directory; omit to initialize randomly | `models/Diffusion_Transformer/taew2_2_super.pth` |
+| `--from_scratch` | Randomly initialize the Super decoder; excludes weight initialization and resume | - |
+| `--train_data_dir` / `--train_data_meta` | Video root and JSON metadata array | `datasets/X-Fun-Videos-Demo` |
+| `--fix_sample_size` | Fixed `[height, width]`, overriding square `--video_sample_size` | `720 1280` |
+| `--video_sample_strides` | Choose one sampling stride per training clip | `1 2 3` |
+| `--video_sample_n_frames` | Training clip length, must be 4k+1 | 25 |
+| `--train_batch_size` / `--gradient_accumulation_steps` | Batch per device / gradient accumulation | 1 / 1 |
+| `--vae_mini_batch` | Official VAE encoding mini-batch size | 1 |
+| `--dataloader_num_workers` | DataLoader workers | 8 |
+| `--num_train_epochs` | Epoch budget used when `--max_train_steps` is not set | 100 |
+| `--checkpointing_steps` / `--checkpoints_total_limit` | Save interval / checkpoint count limit | 1000 / `None` |
+| `--learning_rate` / `--lr_scheduler` | Peak learning rate / scheduler | `5e-6` / `cosine` |
+| `--lr_warmup_steps` / `--lr_min_ratio` | Warmup updates / final-to-peak ratio for cosine decay | 200 / 0.2 |
+| `--l1_loss_weight` / `--lpips_loss_weight` | Pixel L1 / LPIPS loss weights | 0.1 / 0.1 |
+| `--dino_loss_weight` / `--grid_loss_weight` | DINO / periodic grid loss weights | 0.02 / 1.0 |
+| `--feature_batch_size` / `--feature_checkpointing` | Perceptual frame batch / activation checkpointing | 13 / off |
+| `--use_ema` / `--ema_decay` | Maintain decoder EMA for validation and selection | enabled / 0.999 |
+| `--seed` / `--mixed_precision` | Random seed / mixed precision | 42 / `bf16` |
+| `--output_dir` | Output directory | `output_dir_taehv_w2.2_super` |
+| `--resume_from_checkpoint` | Restore a checkpoint directory or select the latest | `latest` |
+
+**Sample Size Configuration Guide**:
+
+- `--fix_sample_size 720 1280` resizes to cover the canvas and center crops while preserving aspect ratio. Without it, `--video_sample_size` selects a square canvas.
+- A 25-frame clip at stride 3 needs at least 73 source frames. Short inputs shorten the batch to the shortest 4k+1 clip length, recorded as `frames` in the training log.
+- Training uses random temporal starts. Data loading retains the original random replacement/retry behavior on decode failures.
+
+**Loss Configuration Guide**:
+
+- The frozen official encoder supplies latents; only the Super decoder is optimized.
+- LPIPS uses AlexNet; DINOv2 uses intermediate ViT-B/14 patch features. Frozen feature networks still pass reconstruction gradients to the decoder.
+- Grid loss measures residual means by spatial phase with each frame/channel's DC removed. Inspect local crops alongside this metric.
+
+### 3.4 Training Validation
+
+Validation is optional and disabled in the default launcher. To enable it, add
+`--validation_data_meta` (plus the related validation arguments) to the training
+command. Relative video paths resolve under `--train_data_dir`, so use absolute
+video paths when validation data has another root.
+
+| Parameter | Description | Recommended Value |
+|-----------|-------------|-------------------|
+| `--validation_data_meta` | Held-out video metadata array | Your validation index |
+| `--validation_steps` | Evaluate every N optimizer updates | 1000 |
+| `--validation_n_frames` | Validation clip length, 4k+1 | 81 |
+| `--validation_max_samples` | Validation clip cap; zero uses the complete index | 0 |
+
+Validation runs at initialization, periodically, and at the final step. It reports
+PSNR, SSIM, LPIPS and grid, and saves the first sample's input/reconstruction preview
+as `sample/step-N.mp4`. Metrics use floating-point output before video encoding,
+averaging frames within a clip and then clips equally.
+
+With EMA enabled, validation and best-weight selection use EMA; otherwise they
+use raw weights. Selection maximizes PSNR subject to LPIPS/grid ratios relative
+to the first validation, controlled by `--selection_lpips_ratio` and
+`--selection_grid_ratio`. The initialization is also a candidate. Best inference
+weights go to `best/taehv`; resumable checkpoints retain raw training weights.
+
+For a short execution check, use the official demo:
+
+```bash
+accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
+  --pretrained_model_name_or_path="$MODEL_NAME" --tae_path="$TAE_PATH" \
+  --train_data_dir="datasets/X-Fun-Videos-Demo" \
+  --train_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
+  --validation_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
+  --fix_sample_size 64 96 --video_sample_n_frames=5 --validation_n_frames=5 \
+  --video_sample_strides 1 --num_train_epochs=1 --lr_warmup_steps=0 \
+  --checkpointing_steps=1 --validation_steps=1 --validation_max_samples=1 \
+  --dataloader_num_workers=0 --use_ema --ema_decay=0.999 \
+  --output_dir="output_dir_taehv_w2.2_super_smoke"
+```
+
+Sharing the demo across training/validation checks execution only. Use independent
+video sources for actual experiments, and use native-resolution data to assess
+720p quality.
+
+### 3.5 Training Tips
+
+- **Warm-start vs. from scratch**: Fine-tune from `taew2_2_super.pth` or a previous Super directory; use `--from_scratch` for a new decoder. Initialization, dataset and budget all affect the result.
+- **Loss balance**: Start with the four-loss settings in the launcher and inspect both reconstruction metrics and native-resolution crops.
+- **Memory**: Lower the training frame count, canvas size or feature batch size, or enable `--feature_checkpointing` if needed. Perceptual networks are frozen but their prediction branches still need activation memory.
+- **EMA**: The launcher enables EMA. `best/taehv` holds selected inference weights; `checkpoint-N/taehv` holds raw weights for training recovery.
+- **Distributed training**: Use Accelerate's single-device or ordinary DDP configuration. Global batch equals device count times per-device batch times accumulation.
+
+### 3.6 Further Tuning
+
+- Try longer temporal context and check both the initial frames and later frames, accounting for the extra compute.
+- Explore pixel MSE or official-decoder output supervision, balancing pixel accuracy, perceptual quality and teacher compute cost.
+- Adjust LPIPS/grid weights with controlled comparisons, checking fine textures, text and local grid artifacts alongside average metrics.
+
+These are possible extensions; the current training entry uses the four-loss recipe above.
+
+---
+
+## 4. Inference Testing
+
+### 4.1 Checkpoint Layout
+
+Each resumable checkpoint is written as `output_dir/checkpoint-{step}/`:
+
+```text
+output_dir_taehv_w2.2_super/
+  args-step-N.yaml          # resolved invocation record
+  train.jsonl              # total loss, four terms, LR and actual frame count
+  validation.jsonl         # metrics and selection decisions, when validation is enabled
+  best/taehv/              # selected inference weights, when validation is enabled
+  checkpoint-1000/
+    taehv/                 # raw Super weights and config.json
+    optimizer.bin
+    custom_checkpoint_*.pkl # scheduler and optional EMA state
+    random_states_*.pkl
+    trainer_state.json
+  sample/                  # validation videos
+  logs/                    # TensorBoard logs
+```
+
+The `taehv` directories use Diffusers `save_pretrained` format and load through
+`AutoencoderTinyWan.from_pretrained`. `args-step-N.yaml` records the invocation;
+training arguments are configured directly in the shell launcher.
+
+### 4.2 Use the Trained TAE in Predict Scripts
+
+Point `tae_path` in a compatible Wan2.2 5B TAE predict script to the selected
+`best/taehv` directory:
+
+| Script | Family |
+|--------|--------|
+| `examples/wan2.2/predict_ti2v_tae.py` | taew2_2 |
+| `examples/wan2.2_fun/predict_t2v_2.2vae_tae.py` | taew2_2 |
+| `examples/wan2.2_fun/predict_i2v_2.2vae_tae.py` | taew2_2 |
+
+```python
+# e.g. in examples/wan2.2_fun/predict_t2v_2.2vae_tae.py
+tae_path = "output_dir_taehv_w2.2_super/best/taehv"
+```
+
+The trained decoder consumes official Wan2.2 5B latents. A scratch-trained
+checkpoint has no trained tiny encoder; retain the official Wan encoder for
+image/video conditioning when using such a checkpoint.
@@ -0,0 +1,416 @@
+# TAE Super（Tiny AutoEncoder）训练指南
+
+本文档提供 Wan2.2 Super decoder（`AutoencoderTinyWan`）训练/微调的完整流程，
+包括环境配置、数据准备、训练，以及将训练好的 checkpoint 用于推理。
+
+> **说明**：训练使用冻结的官方 Wan2.2 encoder，只更新 Super decoder 约 19M 参数。
+> Tiny encoder 冻结，不参与训练前向。
+>
+> ```text
+> x (video, [-1, 1])
+>   -> official Wan2.2 encoder (frozen) : z = vae.encode(x).latent_dist.mode()
+>   -> Super decoder                  : x_hat = tae(z).sample
+> loss = 0.1 L1 + 0.1 LPIPS + 0.02 DINO + 1.0 grid
+> ```
+>
+> 损失在重建与输入 RGB 映射到 [0,1] 后计算，decoder 直接学习解码官方 Wan latent。
+
+脚本支持 Wan2.2 5B Super 模型：
+
+| 家族 | Latent | 官方完整 VAE | `--config_path` | 适用模型 |
+|------|--------|--------------|-----------------|----------|
+| taew2_2_super | 48ch, patch_size=2 | `AutoencoderKLWan3_8` | `config/wan2.2/wan_civitai_5b.yaml` | Wan2.2 TI2V-5B / Fun-2.2VAE |
+
+**与 [train_taehv.py](train_taehv.py) 的差异**：
+
+| 对比项 | `train_taehv.py` | `train_taehv_super.py` |
+|--------|-----------------|-----------------------|
+| 重建路径 | 小 encoder → 小 decoder，官方 encoder 提供 latent 监督 | 冻结官方 encoder → Super decoder |
+| 更新参数 | 默认训练小 encoder 和 decoder，由 `--trainable_modules` 控制 | 仅 Super decoder；官方和小 encoder 均冻结 |
+| 损失 | 像素 L1 + latent MSE | 像素 L1 + LPIPS + DINO + grid |
+
+---
+
+## 目录
+
+- [一、环境配置](#一环境配置)
+- [二、数据准备](#二数据准备)
+  - [2.1 快速测试数据集](#21-快速测试数据集)
+  - [2.2 数据集结构](#22-数据集结构)
+  - [2.3 metadata.json 格式](#23-metadatajson-格式)
+  - [2.4 相对路径与绝对路径使用方案](#24-相对路径与绝对路径使用方案)
+- [三、TAE Super 训练](#三tae-super-训练)
+  - [3.1 下载预训练模型](#31-下载预训练模型)
+  - [3.2 快速开始](#32-快速开始)
+  - [3.3 训练常用参数解析](#33-训练常用参数解析)
+  - [3.4 训练验证](#34-训练验证)
+  - [3.5 训练技巧](#35-训练技巧)
+  - [3.6 后续调优方向](#36-后续调优方向)
+- [四、推理测试](#四推理测试)
+  - [4.1 Checkpoint 目录结构](#41-checkpoint-目录结构)
+  - [4.2 在 Predict 脚本中使用训练好的 TAE](#42-在-predict-脚本中使用训练好的-tae)
+- [五、更多资源](#五更多资源)
+
+---
+
+## 一、环境配置
+
+以下命令在 VideoX-Fun 仓库根目录执行，使用 pip 将依赖安装到当前 Python 环境。
+
+**方式 1：使用 requirements.txt**
+
+```bash
+pip install -r requirements.txt
+pip install lpips torchvision av modelscope opencv-python
+```
+
+**方式 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 lpips torchvision av modelscope opencv-python
+```
+
+LPIPS 的 AlexNet 主干和 DINOv2 在缺少缓存时自动下载，后续复用 PyTorch 缓存；
+LPIPS 校准权重随安装包提供。离线加载 DINOv2 时，可使用 `--dino_repo` 和 `--dino_weights`。
+
+> 脚本通过 Accelerate 支持单设备与普通 DDP。官方 Wan VAE 和冻结的 LPIPS/DINOv2
+> 特征网络均保留在 GPU 上。
+
+---
+
+## 二、数据准备
+
+### 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
+  }
+]
+```
+
+**绝对路径格式**：
+```json
+[
+  {
+    "file_path": "/mnt/data/videos/sunset.mp4",
+    "text": "A beautiful sunset over the ocean",
+    "type": "video",
+    "width": 1024,
+    "height": 1024
+  }
+]
+```
+
+**关键字段说明**：
+
+- `file_path`：视频路径（相对或绝对路径）
+- `text`：视频描述（TAE 损失不使用，仅为 meta 格式兼容而保留）
+- `type`：数据类型，固定为 `"video"`
+- `width` / `height`：视频宽高（推荐保留）。Super 训练入口使用固定的缩放/裁剪画布。
+
+### 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.json"
+```
+
+> 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。
+
+---
+
+## 三、TAE Super 训练
+
+### 3.1 下载预训练模型
+
+训练脚本需要随模型目录提供的官方 Wan2.2 VAE 权重。微调时同时下载发布版 Super 权重：
+
+```bash
+mkdir -p models/Diffusion_Transformer
+
+modelscope download --model Wan-AI/Wan2.2-TI2V-5B --local_dir models/Diffusion_Transformer/Wan2.2-TI2V-5B
+
+wget -O models/Diffusion_Transformer/taew2_2_super.pth \
+  https://github.com/madebyollin/taehv/raw/main/taew2_2_super.pth
+```
+
+LPIPS 与 DINOv2 权重由损失实现单独加载。
+
+### 3.2 快速开始
+
+**Wan2.2 TI2V-5B / Fun-2.2VAE（taew2_2_super）训练示例**：
+
+启动器采用 720x1280 画布、训练 25 帧和 EMA，不设置 `--max_train_steps`：
+训练按 `--num_train_epochs`（默认 100）个 epoch 运行，直到手动终止进程。
+使用 `accelerate config` 配置 GPU/进程数；四卡、每卡 batch=1、累积=1 时有效 batch 为 4。
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth"
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
+  --config_path="config/wan2.2/wan_civitai_5b.yaml" \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --tae_path="$TAE_PATH" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --fix_sample_size 720 1280 \
+  --video_sample_strides 1 2 3 \
+  --video_sample_n_frames=25 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --checkpointing_steps=1000 \
+  --learning_rate=5e-06 \
+  --lr_scheduler="cosine" \
+  --lr_warmup_steps=200 \
+  --lr_min_ratio=0.2 \
+  --seed=42 \
+  --output_dir="output_dir_taehv_w2.2_super" \
+  --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
+  --adam_weight_decay=1e-2 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=1.0 \
+  --l1_loss_weight=0.1 \
+  --lpips_loss_weight=0.1 \
+  --dino_loss_weight=0.02 \
+  --grid_loss_weight=1.0 \
+  --grid_period=16 \
+  --feature_batch_size=13 \
+  --use_ema \
+  --ema_decay=0.999 \
+  --selection_lpips_ratio=1.02 \
+  --selection_grid_ratio=1.05
+```
+
+相同参数已写入 `scripts/taehv/train_taehv_super.sh`，也可修改其中的训练参数后，
+运行 `bash scripts/taehv/train_taehv_super.sh`。
+
+**从头训练**：
+
+```bash
+accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --from_scratch \
+  --fix_sample_size 720 1280 \
+  --video_sample_n_frames=25 \
+  --video_sample_strides 1 2 3 \
+  --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
+  --use_ema --ema_decay=0.999 \
+  --dataloader_num_workers=8 --checkpointing_steps=1000 \
+  --output_dir="output_dir_taehv_w2.2_super_scratch"
+```
+
+使用 `--from_scratch` 时不要传 `--tae_path`。只训练 decoder，保存的 tiny encoder
+仍是随机初始化权重。
+
+**中断后恢复训练**：
+
+```bash
+accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --fix_sample_size 720 1280 --video_sample_n_frames=25 \
+  --video_sample_strides 1 2 3 \
+  --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
+  --use_ema --ema_decay=0.999 \
+  --dataloader_num_workers=8 --checkpointing_steps=1000 \
+  --output_dir="output_dir_taehv_w2.2_super" \
+  --resume_from_checkpoint=latest
+```
+
+恢复时保持创建 checkpoint 时的配置。未设置 `--max_train_steps` 时，schedule
+预算由 `--num_train_epochs` 决定，恢复时应保留原值。Checkpoint 恢复
+optimizer、scheduler、EMA、RNG、step 和 batch 游标；改变 workers 或采样设置可能改变
+恢复后的样本顺序。新训练应使用新的输出目录。
+
+### 3.3 训练常用参数解析
+
+下表列出启动器的配置值，不代表所有参数的 Python CLI 默认值。
+
+| 参数 | 说明 | 示例值 |
+|------|------|--------|
+| `--config_path` | 官方 Wan VAE 组件配置 | `config/wan2.2/wan_civitai_5b.yaml` |
+| `--pretrained_model_name_or_path` | 包含官方 VAE 权重的模型目录 | `models/Diffusion_Transformer/Wan2.2-TI2V-5B` |
+| `--tae_path` | Super 权重文件或 Diffusers 目录；不传则随机初始化 | `models/Diffusion_Transformer/taew2_2_super.pth` |
+| `--from_scratch` | 随机初始化 Super decoder；不能与加载权重或恢复组合 | - |
+| `--train_data_dir` / `--train_data_meta` | 视频根目录与 JSON 元数据数组 | `datasets/X-Fun-Videos-Demo` |
+| `--fix_sample_size` | 固定 `[高, 宽]`，覆盖正方形 `--video_sample_size` | `720 1280` |
+| `--video_sample_strides` | 每个训练片段随机选择一个采样 stride | `1 2 3` |
+| `--video_sample_n_frames` | 训练片段帧数，必须为 4k+1 | 25 |
+| `--train_batch_size` / `--gradient_accumulation_steps` | 每卡 batch / 梯度累积 | 1 / 1 |
+| `--vae_mini_batch` | 官方 VAE 编码 mini-batch | 1 |
+| `--dataloader_num_workers` | DataLoader workers | 8 |
+| `--num_train_epochs` | 未设置 `--max_train_steps` 时的 epoch 预算 | 100 |
+| `--checkpointing_steps` / `--checkpoints_total_limit` | 保存间隔 / checkpoint 数量上限 | 1000 / `None` |
+| `--learning_rate` / `--lr_scheduler` | 峰值学习率 / 调度器 | `5e-6` / `cosine` |
+| `--lr_warmup_steps` / `--lr_min_ratio` | 预热步数 / cosine 终点与峰值之比 | 200 / 0.2 |
+| `--l1_loss_weight` / `--lpips_loss_weight` | 像素 L1 / LPIPS 损失权重 | 0.1 / 0.1 |
+| `--dino_loss_weight` / `--grid_loss_weight` | DINO / 周期格纹损失权重 | 0.02 / 1.0 |
+| `--feature_batch_size` / `--feature_checkpointing` | 感知网络帧 batch / 激活 checkpointing | 13 / 关闭 |
+| `--use_ema` / `--ema_decay` | 维护用于验证和选模的 decoder EMA | 开启 / 0.999 |
+| `--seed` / `--mixed_precision` | 随机种子 / 混合精度 | 42 / `bf16` |
+| `--output_dir` | 输出目录 | `output_dir_taehv_w2.2_super` |
+| `--resume_from_checkpoint` | 恢复 checkpoint 目录或选择最新断点 | `latest` |
+
+**Sample Size 配置指南**：
+
+- `--fix_sample_size 720 1280` 等比例缩放覆盖画布后中心裁剪；未指定时，`--video_sample_size` 控制正方形画布。
+- 25 帧、stride=3 至少需要 73 个源视频帧。短视频会使 batch 按最短视频截为 4k+1 帧，实际长度记录在训练日志的 `frames` 字段。
+- 训练随机选择时间起点。解码失败时沿用原有随机换视频重试逻辑。
+
+**损失配置指南**：
+
+- 冻结官方 encoder 提供 latent，只优化 Super decoder。
+- LPIPS 使用 AlexNet，DINOv2 使用 ViT-B/14 中间 patch 特征；特征网络冻结，但仍将重建分支的梯度传回 decoder。
+- Grid loss 按空间相位统计残差均值，并去掉每帧每通道的 DC；应结合局部原生裁剪判断伪影。
+
+### 3.4 训练验证
+
+验证为可选功能，默认启动器中未开启。需要时在训练命令中通过 `--validation_data_meta`
+（及相关验证参数）启用。相对视频路径基于 `--train_data_dir` 解析；验证数据位于其他
+根目录时，可使用绝对视频路径。
+
+| 参数 | 说明 | 推荐值 |
+|------|------|--------|
+| `--validation_data_meta` | 独立验证集元数据数组 | 自己的验证索引 |
+| `--validation_steps` | 每 N 次优化器更新执行验证 | 1000 |
+| `--validation_n_frames` | 验证片段帧数，4k+1 | 81 |
+| `--validation_max_samples` | 验证片段数量上限，0 表示全部 | 0 |
+
+初始化、周期节点和最后一步均执行验证，报告 PSNR、SSIM、LPIPS 和 grid，并将首个
+样本的输入/重建预览保存为 `sample/step-N.mp4`。指标使用视频编码前的浮点输出，
+先对片段内帧平均，再对片段等权平均。
+
+开启 EMA 时验证和最佳权重选择使用 EMA，否则使用 raw。选模在相对首次验证的
+LPIPS/grid 比例约束下最大化 PSNR，由 `--selection_lpips_ratio`、`--selection_grid_ratio`
+控制。初始权重也是候选。最佳推理权重写入 `best/taehv`，可恢复 checkpoint 保留 raw
+训练权重。
+
+可使用官方 Demo 进行短流程检查：
+
+```bash
+accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
+  --pretrained_model_name_or_path="$MODEL_NAME" --tae_path="$TAE_PATH" \
+  --train_data_dir="datasets/X-Fun-Videos-Demo" \
+  --train_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
+  --validation_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
+  --fix_sample_size 64 96 --video_sample_n_frames=5 --validation_n_frames=5 \
+  --video_sample_strides 1 --num_train_epochs=1 --lr_warmup_steps=0 \
+  --checkpointing_steps=1 --validation_steps=1 --validation_max_samples=1 \
+  --dataloader_num_workers=0 --use_ema --ema_decay=0.999 \
+  --output_dir="output_dir_taehv_w2.2_super_smoke"
+```
+
+短测中训练和验证共用 Demo 只用于检查执行；正式实验应按视频来源隔离，评估 720p
+质量时应使用原生分辨率数据。
+
+### 3.5 训练技巧
+
+- **热启动与从头训练**：从 `taew2_2_super.pth` 或已有 Super 目录微调；使用 `--from_scratch` 训练新 decoder。初始化、数据和预算都会影响结果。
+- **损失平衡**：先使用启动器中的四项损失配置，同时检查重建指标和原生分辨率裁剪。
+- **显存**：按需降低训练帧数、画布尺寸或特征 batch，或开启 `--feature_checkpointing`。感知网络虽然冻结，预测分支仍需要激活显存。
+- **EMA**：启动器默认开启 EMA。`best/taehv` 为选出的推理权重，`checkpoint-N/taehv` 为恢复训练使用的 raw 权重。
+- **分布式训练**：使用 Accelerate 单设备或普通 DDP 配置。有效 batch 为设备数、每卡 batch 和梯度累积的乘积。
+
+### 3.6 后续调优方向
+
+- 尝试更长的时间上下文，同时检查视频起始段和后续帧，并计入额外计算开销。
+- 探索像素 MSE 或官方 decoder 输出监督，平衡像素精度、感知质量和 teacher 计算成本。
+- 通过受控对照调整 LPIPS/grid 权重，结合细纹理、文字和局部格纹检查平均指标。
+
+以上为可探索的扩展方向，当前训练入口使用前述四项损失配置。
+
+---
+
+## 四、推理测试
+
+### 4.1 Checkpoint 目录结构
+
+可恢复 checkpoint 写入 `output_dir/checkpoint-{step}/`：
+
+```text
+output_dir_taehv_w2.2_super/
+  args-step-N.yaml          # 实际运行参数记录
+  train.jsonl              # 总损失、四项损失、LR 和实际帧数
+  validation.jsonl         # 提供验证索引时保存指标与选模结果
+  best/taehv/              # 提供验证索引时保存最佳推理权重
+  checkpoint-1000/
+    taehv/                 # raw Super 权重与 config.json
+    optimizer.bin
+    custom_checkpoint_*.pkl # scheduler 和可选 EMA 状态
+    random_states_*.pkl
+    trainer_state.json
+  sample/                  # 验证视频
+  logs/                    # TensorBoard 日志
+```
+
+`taehv` 目录使用 Diffusers `save_pretrained` 格式，可由
+`AutoencoderTinyWan.from_pretrained` 加载。`args-step-N.yaml` 记录本次调用参数，
+训练参数直接在 shell 启动器中配置。
+
+### 4.2 在 Predict 脚本中使用训练好的 TAE
+
+将兼容的 Wan2.2 5B TAE predict 脚本中的 `tae_path` 指向选出的 `best/taehv` 目录：
+
+| 脚本 | 家族 |
+|------|------|
+| `examples/wan2.2/predict_ti2v_tae.py` | taew2_2 |
+| `examples/wan2.2_fun/predict_t2v_2.2vae_tae.py` | taew2_2 |
+| `examples/wan2.2_fun/predict_i2v_2.2vae_tae.py` | taew2_2 |
+
+```python
+# 例如 examples/wan2.2_fun/predict_t2v_2.2vae_tae.py 中
+tae_path = "output_dir_taehv_w2.2_super/best/taehv"
+```
+
+训练好的 decoder 接收官方 Wan2.2 5B latent。从头训练的 checkpoint 没有训练好的
+小 encoder；此类 checkpoint 用于图像/视频条件生成时，应保留官方 Wan encoder。
@@ -128,7 +128,7 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun

 ```bash
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
-export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 ```

 **绝对路径**：
@@ -173,42 +173,50 @@ wget https://github.com/madebyollin/taehv/raw/main/taew2_2.safetensors
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
-export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # 可选：用官方发布的 TAE 权重热启动，而不是从头训练。
 export TAE_PATH="models/Diffusion_Transformer/taew2_2.safetensors"
 NCCL_DEBUG=INFO

 accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
  --config_path="config/wan2.2/wan_civitai_5b.yaml" \
-  --pretrained_model_name_or_path=$MODEL_NAME \
-  --tae_path=$TAE_PATH \
-  --train_data_dir=$DATASET_NAME \
-  --train_data_meta=$DATASET_META_NAME \
-  --video_sample_size=512 \
-  --video_sample_stride=1 \
-  --video_sample_n_frames=33 \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --tae_path="$TAE_PATH" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --fix_sample_size 720 1280 \
+  --video_sample_strides 1 2 3 \
+  --video_sample_n_frames=25 \
+  --validation_n_frames=49 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
-  --dataloader_num_workers=4 \
-  --num_train_epochs=100 \
-  --checkpointing_steps=500 \
+  --dataloader_num_workers=8 \
+  --max_train_steps=10000 \
+  --checkpointing_steps=1000 \
+  --validation_steps=1000 \
+  --validation_paths "asset/inpaint_video.mp4" \
  --learning_rate=1e-04 \
  --lr_scheduler="constant_with_warmup" \
-  --lr_warmup_steps=100 \
+  --lr_warmup_steps=200 \
  --latent_loss_weight=1.0 \
+  --use_ema \
+  --ema_decay=0.999 \
  --seed=42 \
  --output_dir="output_dir_taehv_w2.2" \
  --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
  --adam_weight_decay=1e-2 \
  --adam_epsilon=1e-08 \
  --vae_mini_batch=1 \
  --max_grad_norm=1.0 \
-  --random_hw_adapt \
  --enable_bucket \
  --trainable_modules "." \
  --resume_from_checkpoint=latest
 ```

+相同参数已写入 `scripts/taehv/train_taehv.sh`，也可修改其中的训练参数后，运行 `bash scripts/taehv/train_taehv.sh`。
+
 **Wan2.1 / Wan2.2 14B（taew2_1）训练示例**：

 与上面相同，仅做如下修改：
@@ -237,15 +245,17 @@ export TAE_PATH="taew2_1.safetensors"
 | `--use_taehv_sequential` | 以 O(1) 显存的串行模式运行 TAE，而不是并行模式 | - |
 | `--latent_loss_weight` | latent MSE（TAE latent 对齐完整 VAE latent）相对 pixel L1 的权重 | 1.0 |
 | `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` |
-| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata.json` |
+| `--train_data_meta` | 训练数据元文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` |
 | `--train_batch_size` | 每卡批次大小 | 1 |
 | `--video_sample_size` | 训练分辨率 | 512 |
 | `--video_sample_stride` | 视频采样步幅 | 1 |
+| `--video_sample_strides` | 空格分隔的候选步幅，每个片段随机抽取一个（覆盖 `--video_sample_stride`） | `None` |
 | `--video_sample_n_frames` | 采样帧数，**必须为 4k+1**（33、49、81……） | 33 |
 | `--vae_mini_batch` | teacher VAE 编码时的迷你批次大小 | 1 |
 | `--gradient_accumulation_steps` | 梯度累积步数 | 1 |
 | `--dataloader_num_workers` | DataLoader 子进程数 | 4 |
 | `--num_train_epochs` | 训练 epoch 数 | 100 |
+| `--max_train_steps` | 总训练步数；设置后覆盖 `--num_train_epochs` | 10000 |
 | `--checkpointing_steps` | 每 N 步保存 checkpoint | 500 |
 | `--checkpoints_total_limit` | 最多保留的 checkpoint 数量 | `None` |
 | `--learning_rate` | 初始学习率 | 1e-4 |
@@ -253,22 +263,26 @@ export TAE_PATH="taew2_1.safetensors"
 | `--lr_warmup_steps` | 学习率预热步数 | 100 |
 | `--use_8bit_adam` / `--use_came` | 可选优化器 | - |
 | `--use_ema` | 维护 TAE 的 EMA 副本（用于验证与最终保存） | - |
-| `--seed` | 随机种子 | 42 |
+| `--ema_decay` | EMA 衰减系数，同时关闭 EMA warmup；不传则用 `EMAModel` 默认值（0.9999） | `None` |
+| `--seed` | 随机种子；默认 42，直接跑命令行也可复现 | 42 |
 | `--output_dir` | 输出目录 | `output_dir_taehv_w2.2` |
 | `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
 | `--max_grad_norm` | 梯度裁剪阈值 | 1.0 |
 | `--enable_bucket` | 启用分桶训练，不裁剪视频，按分辨率分组训练 | - |
-| `--random_hw_adapt` | 自动缩放视频到一定范围内的随机尺寸 | - |
+| `--fix_sample_size` | 固定画布 `[height, width]`，覆盖方形 `--video_sample_size`（等比缩放后中心裁剪） | `None` |
 | `--low_vram` | 将 teacher VAE 放在 CPU，仅在编码时搬上 GPU | - |
 | `--trainable_modules` | 可训练模块（`"."` 表示所有模块） | `"."` |
 | `--trainable_modules_low_learning_rate` | 以 lr/2 训练的模块 | `[]` |
 | `--resume_from_checkpoint` | 恢复训练路径，使用 `"latest"` 自动选择最新 checkpoint | `latest` |
 | `--validation_steps` / `--validation_epochs` | 每 N 步 / 每 N 个 epoch 执行一次验证 | 2000 / 5 |
-| `--validation_paths` | 用于重建对比验证的视频路径 | `"asset/1.mp4"` |
+| `--validation_paths` | 用于重建对比验证的视频路径 | `"asset/inpaint_video.mp4"` |
+| `--validation_n_frames` | 验证片段帧数（4k+1），不传则复用 `--video_sample_n_frames` | `None` |

 **Sample Size 配置指南**：
- `video_sample_size` 表示训练分辨率；当启用 `random_hw_adapt` 时，表示分辨率的最小值。
+- `video_sample_size` 设置方形训练画布；配合 `--enable_bucket` 可按原始宽高比训练。
+- `--fix_sample_size 720 1280` 固定非方形画布（等比缩放 + 中心裁剪）；配合 `--enable_bucket` 时按批生效，`video_sample_size` 会提升到 `max(fix_sample_size)`。
 - `video_sample_n_frames` 必须满足 `4k+1`（如 33、49、81），因为完整 VAE 与 TAE 都是因果 4 倍时间压缩。
+- `--video_sample_strides 1 2 3` 会为每个片段随机抽取一个时间步幅（多尺度时间增广）；步幅越大请求的帧数越少，因此单个片段覆盖的时间跨度基本不变。

 ### 3.4 训练验证

@@ -278,16 +292,19 @@ export TAE_PATH="taew2_1.safetensors"
 |------|------|--------|
 | `--validation_steps` | 每 N 步执行一次验证 | 2000 |
 | `--validation_epochs` | 每 N 个 epoch 执行一次验证 | 5 |
-| `--validation_paths` | 验证视频路径 | `"asset/1.mp4"` |
+| `--validation_paths` | 验证视频路径 | `"asset/inpaint_video.mp4"` |
+| `--validation_n_frames` | 验证片段帧数，**必须为 4k+1**；不传则复用 `--video_sample_n_frames` | 81 |

 ```bash
-  --validation_paths "asset/1.mp4" \
+  --validation_paths "asset/inpaint_video.mp4" \
  --validation_steps=2000 \
-  --validation_epochs=5
+  --validation_epochs=5 \
+  --validation_n_frames=81
 ```

 **注意事项**：
 - 验证会为每个样本在 `output_dir/sample/` 保存三个视频：`*_input.mp4`（缩放后的输入）、`*_taehv.mp4`（TAE 重建）、`*_fullvae.mp4`（完整 VAE 重建，供参考对比）。
+- 验证片段会先截取 `min(视频长度, --validation_n_frames)` 帧、再向下对齐到 `4k+1`，因此把该值设大不会在短视频素材上失败。
 - 启用 `--use_ema` 时，验证使用 EMA 权重。

 ### 3.5 训练技巧
@@ -297,7 +314,7 @@ export TAE_PATH="taew2_1.safetensors"
 - **Decoder-only 蒸馏**：加上 `--freeze_tae_encoder` 只提升解码质量。
 - **显存**：teacher 完整 VAE 是主要的显存消耗者。使用 `--low_vram` 可以将其放在 CPU、仅在编码时搬上 GPU；也可以降低 `--video_sample_n_frames` / `--video_sample_size`。
 - **串行 TAE**：`--use_taehv_sequential` 以速度换显存，激活显存与视频长度无关（O(1)）。
- **EMA**：推荐开启 `--use_ema` 作为最终交付物；最终保存的 `taehv` 目录里是 EMA 权重。
+- **EMA**：推荐开启 `--use_ema` 作为最终交付物；最终保存的 `taehv` 目录里是 EMA 权重。短训练希望均值更快跟上时，可传 `--ema_decay`（如 `train_taehv.sh` 与 `train_taehv_super.sh` 用的 `0.999`）替代默认的 0.9999。

 ### 3.6 多机分布式训练

@@ -17,6 +17,10 @@ The teacher VAE family is selected by --config_path:
    config/wan2.1/wan_civitai.yaml                -> AutoencoderKLWan    (16ch, taew2_1)
    config/wan2.2/wan_civitai_5b.yaml             -> AutoencoderKLWan3_8 (48ch, taew2_2)
 (any wan2.2 *_2.2vae.yaml also works, they share the same vae_kwargs)
+
+The training loop mirrors scripts/wan2.1_fun/train.py (derived from
+https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py):
+same shared argument surface, dataset/sampler, EMA and checkpoint hooks.
 """
 #!/usr/bin/env python
 # coding=utf-8
@@ -34,8 +38,8 @@ The teacher VAE family is selected by --config_path:
 # See the License for the specific language governing permissions and

 import argparse
+import copy
 import gc
-import inspect
 import logging
 import math
 import os
@@ -61,7 +65,6 @@ from einops import rearrange
 from omegaconf import OmegaConf
 from packaging import version
 from PIL import Image
-from torch.utils.data import RandomSampler
 from torch.utils.tensorboard import SummaryWriter
 from torchvision import transforms
 from transformers.utils import ContextManagers
@@ -91,13 +94,6 @@ check_min_version("0.18.0.dev0")
 logger = get_logger(__name__, log_level="INFO")


-def filter_kwargs(cls, kwargs):
-    """Filter kwargs to only include parameters accepted by the class __init__."""
-    sig = inspect.signature(cls.__init__)
-    valid_keys = set(sig.parameters.keys()) - {"self"}
-    return {k: v for k, v in kwargs.items() if k in valid_keys}
-
-
 def linear_decay(initial_value, final_value, total_steps, current_step):
    if current_step >= total_steps:
        return final_value
@@ -109,8 +105,13 @@ def linear_decay(initial_value, final_value, total_steps, current_step):

 def log_validation(vae, tae, args, accelerator, weight_dtype, global_step):
    """Validation: encode video -> decode with TAE / full VAE, save comparison videos."""
-    try:
-        with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype):
+    # The guard sits *inside* the contexts on purpose: `Accelerator.autocast` enters and
+    # exits AMP without a try/finally, so an exception escaping the block strands autocast
+    # enabled and leaves its cast-weight cache holding tensors built under no_grad. The
+    # next training step would then produce a graph-free loss and `accelerator.backward`
+    # would die with "element 0 of tensors does not require grad".
+    with torch.no_grad(), accelerator.autocast():
+        try:
            logger.info("Running validation...")

            if args.validation_paths is None or len(args.validation_paths) == 0:
@@ -120,6 +121,10 @@ def log_validation(vae, tae, args, accelerator, weight_dtype, global_step):
            from decord import VideoReader

            unwrapped_tae = accelerator.unwrap_model(tae)
+            # Under --low_vram the teacher VAE waits on CPU between encoding steps.
+            teacher_on_gpu = bool(args.low_vram)
+            if teacher_on_gpu:
+                vae.to(accelerator.device, dtype=weight_dtype)

            for i, video_path in enumerate(args.validation_paths):
                if not os.path.exists(video_path):
@@ -128,7 +133,7 @@ def log_validation(vae, tae, args, accelerator, weight_dtype, global_step):

                # Load video frames
                vr = VideoReader(video_path)
-                num_frames = min(len(vr), args.video_sample_n_frames)
+                num_frames = min(len(vr), args.validation_n_frames)
                # Align to temporal compression ratio (4k+1)
                temporal_ratio = vae.config.temporal_compression_ratio
                num_frames = (num_frames - 1) // temporal_ratio * temporal_ratio + 1
@@ -144,8 +149,8 @@ def log_validation(vae, tae, args, accelerator, weight_dtype, global_step):

                # Resize to target size aligned to the spatial compression ratio
                spatial_ratio = int(vae.config.spatial_compression_ratio)
-                h_target = int(args.video_sample_size / spatial_ratio) * spatial_ratio
-                w_target = h_target  # square for simplicity in validation
+                h_target = int(args.video_sample_height / spatial_ratio) * spatial_ratio
+                w_target = int(args.video_sample_width / spatial_ratio) * spatial_ratio
                pixel_values = F.interpolate(
                    pixel_values, size=(h_target, w_target), mode='bilinear', align_corners=False
                )
@@ -161,29 +166,33 @@ def log_validation(vae, tae, args, accelerator, weight_dtype, global_step):
                z_full = vae.encode(pixel_values).latent_dist.mode()
                decoded_full = vae.decode(z_full, return_dict=False)[0]

-                # Save videos
+                # Save videos; save_videos_grid goes through numpy, which has no bf16.
                os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
                save_videos_grid(
-                    pixel_values,
+                    pixel_values.float(),
                    os.path.join(args.output_dir, f"sample/step{global_step}_val{i}_input.mp4"),
                    rescale=True, fps=24
                )
                save_videos_grid(
-                    decoded_tae,
+                    decoded_tae.float(),
                    os.path.join(args.output_dir, f"sample/step{global_step}_val{i}_taehv.mp4"),
                    rescale=True, fps=24
                )
                save_videos_grid(
-                    decoded_full,
+                    decoded_full.float(),
                    os.path.join(args.output_dir, f"sample/step{global_step}_val{i}_fullvae.mp4"),
                    rescale=True, fps=24
                )
                logger.info(f"Saved validation video {i} at step {global_step}")

-    except Exception as e:
-        gc.collect()
-        torch.cuda.empty_cache()
-        print(f"Eval error on rank {accelerator.process_index} with info {e}")
+            if teacher_on_gpu:
+                vae.to("cpu")
+                torch.cuda.empty_cache()
+
+        except Exception as e:
+            gc.collect()
+            torch.cuda.empty_cache()
+            print(f"Eval error on rank {accelerator.process_index} with info {e}")


 def parse_args():
@@ -218,14 +227,14 @@ def parse_args():
    parser.add_argument(
        "--train_data_dir",
        type=str,
-        default=None,
+        default="datasets/X-Fun-Videos-Demo",
        help=("A folder containing the training data."),
    )
    parser.add_argument(
        "--train_data_meta",
        type=str,
-        default=None,
-        help=("A json/csv containing the training data meta."),
+        default="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json",
+        help=("A JSON array containing the training video metadata."),
    )
    parser.add_argument(
        "--max_train_samples",
@@ -255,7 +264,9 @@ def parse_args():
        default=None,
        help="The directory where the downloaded models and datasets will be stored.",
    )
-    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    # The default must be a concrete int: the batch-sampler generators and `worker_init_fn`
+    # below call `manual_seed(args.seed)` / `args.seed + process_index` without a None guard.
+    parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible training.")
    parser.add_argument(
        "--train_batch_size", type=int, default=1, help="Batch size (per device) for the training dataloader."
    )
@@ -318,6 +329,12 @@ def parse_args():
        ),
    )
    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
+    parser.add_argument(
+        "--ema_decay",
+        type=float,
+        default=None,
+        help="EMA decay for the tracked parameters. Defaults to the EMAModel default (0.9999).",
+    )
    parser.add_argument(
        "--non_ema_revision",
        type=str,
@@ -389,6 +406,12 @@ def parse_args():
        default=2000,
        help="Run validation every X steps.",
    )
+    parser.add_argument(
+        "--validation_n_frames",
+        type=int,
+        default=None,
+        help="Validation clip length, 4k+1. Defaults to `--video_sample_n_frames`.",
+    )
    parser.add_argument(
        "--tracker_project_name",
        type=str,
@@ -398,36 +421,40 @@ def parse_args():
    parser.add_argument(
        "--enable_bucket", action="store_true", help="Whether enable bucket sample in datasets."
    )
-    parser.add_argument(
-        "--random_ratio_crop", action="store_true", help="Whether enable random ratio crop sample in datasets."
-    )
-    parser.add_argument(
-        "--random_hw_adapt", action="store_true", help="Whether enable random adapt height and width in datasets."
-    )
    parser.add_argument(
        "--video_sample_size",
        type=int,
        default=512,
        help="Training video sample size.",
    )
+    parser.add_argument(
+        "--fix_sample_size",
+        nargs=2,
+        type=int,
+        default=None,
+        help="Fix Sample size [height, width] to override `--video_sample_size` with a fixed "
+             "non-square shape. Under `--enable_bucket` the scalar `--video_sample_size` is bumped to "
+             "at least max(fix_sample_size), so the bucket ratios still cover the canvas.",
+    )
    parser.add_argument(
        "--video_sample_stride",
        type=int,
        default=1,
        help="Sample stride of the video.",
    )
+    parser.add_argument(
+        "--video_sample_strides",
+        nargs="+",
+        type=int,
+        default=None,
+        help="Optional per-video random sampling strides; one is drawn for every clip.",
+    )
    parser.add_argument(
        "--video_sample_n_frames",
        type=int,
        default=33,
        help="Num frame of video. Must be 4k+1 to match the causal temporal compression.",
    )
-    parser.add_argument(
-        "--video_repeat",
-        type=int,
-        default=0,
-        help="Num of repeat video.",
-    )
    parser.add_argument(
        "--tae_path",
        type=str,
@@ -512,10 +539,45 @@ def parse_args():
    if args.non_ema_revision is None:
        args.non_ema_revision = args.revision

+    # Derive the effective canvas once, so every downstream use agrees.
+    if args.fix_sample_size is not None:
+        args.video_sample_height, args.video_sample_width = args.fix_sample_size
+        # The bucket sampler groups by the scalar sample size; keep it >= the fixed canvas.
+        if args.enable_bucket:
+            args.video_sample_size = max(max(args.fix_sample_size), args.video_sample_size)
+    else:
+        args.video_sample_height = args.video_sample_width = args.video_sample_size
+    if any(size < 32 or size % 16 for size in (args.video_sample_height, args.video_sample_width)):
+        parser.error("Video dimensions must be >=32 and divisible by 16")
+
+    if args.validation_n_frames is None:
+        args.validation_n_frames = args.video_sample_n_frames
    # Both the full VAE and the TAE are causal 4x temporal compressors, so the
    # clip length must be 4k+1 for the decode output to match the input.
-    assert (args.video_sample_n_frames - 1) % 4 == 0, \
-        f"video_sample_n_frames must be 4k+1, got {args.video_sample_n_frames}"
+    if any(frames < 1 or (frames - 1) % 4 for frames in (args.video_sample_n_frames, args.validation_n_frames)):
+        parser.error("Video frame counts must be 4k+1")
+    if args.max_train_steps is not None and args.max_train_steps < 1:
+        parser.error("max_train_steps must be positive")
+    if min(args.num_train_epochs, args.train_batch_size, args.vae_mini_batch,
+           args.gradient_accumulation_steps, args.checkpointing_steps, args.video_sample_stride) < 1:
+        parser.error("Epochs, batch sizes, checkpoint interval and stride must be positive")
+    if min(args.dataloader_num_workers, args.validation_steps, args.validation_epochs, args.lr_warmup_steps) < 0:
+        parser.error("Workers, validation intervals and warmup cannot be negative")
+    if args.video_sample_strides is not None and (not args.video_sample_strides or min(args.video_sample_strides) < 1):
+        parser.error("video_sample_strides must be positive")
+    if args.checkpoints_total_limit is not None and args.checkpoints_total_limit < 1:
+        parser.error("checkpoints_total_limit must be positive")
+    if args.ema_decay is not None and not 0 <= args.ema_decay < 1:
+        parser.error("ema_decay must be in [0, 1)")
+    if not 0 <= args.adam_beta1 < 1 or not 0 <= args.adam_beta2 < 1:
+        parser.error("Invalid Adam betas")
+    if args.adam_weight_decay < 0:
+        parser.error("Adam weight decay cannot be negative")
+    if any(not math.isfinite(value) or value <= 0 for value in
+           (args.learning_rate, args.adam_epsilon, args.max_grad_norm)):
+        parser.error("Learning rate, epsilon and gradient norm must be finite and positive")
+    if args.latent_loss_weight < 0:
+        parser.error("latent_loss_weight cannot be negative")

    return args

@@ -652,9 +714,9 @@ def main():
            arch_variant=args.tae_arch_variant,
        )
    tae.parallel = not args.use_taehv_sequential
-    # The TAE is tiny (~20MB); keep it fully in weight_dtype so forward works
-    # regardless of whether autocast is active.
-    tae.to(weight_dtype)
+    # Trainable parameters (and the Adam moments built below) stay in FP32 like the
+    # other training scripts; mixed precision comes from `accelerator.autocast()`.
+    tae.float()

    # Set trainable parameters
    tae.requires_grad_(False)
@@ -687,10 +749,17 @@ def main():
            # every local shard of the copy pairs 1:1 with the live shard.
            ema_tae = FSDPEMA(ema_module, source=tae, accelerator=accelerator, fsdp_plugin=fsdp_plugin)
        else:
-            ema_module = ema_module.to(weight_dtype)
+            ema_module = ema_module.float()
            if args.tae_path is not None:
                ema_module.load_state_dict(tae.state_dict())
-            ema_tae = EMAModel(ema_module.parameters(), model_cls=AutoencoderTinyWan, model_config=ema_module.config)
+            # An explicit decay also disables the warmup ramp, mirroring train_taehv_super.py.
+            ema_kwargs = {} if args.ema_decay is None else {
+                "decay": args.ema_decay, "min_decay": args.ema_decay, "use_ema_warmup": False,
+            }
+            ema_tae = EMAModel(
+                ema_module.parameters(), model_cls=AutoencoderTinyWan, model_config=ema_module.config,
+                **ema_kwargs,
+            )

    # ==================== Save/Load Hooks ====================
    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
@@ -822,7 +891,8 @@ def main():

    train_dataset = VideoDataset(
        args.train_data_meta, args.train_data_dir,
-        sample_size=args.video_sample_size, sample_stride=args.video_sample_stride,
+        sample_size=(args.video_sample_height, args.video_sample_width),
+        sample_stride=args.video_sample_strides or args.video_sample_stride,
        sample_n_frames=args.video_sample_n_frames,
        enable_bucket=args.enable_bucket, enable_inpaint=False,
    )
@@ -844,6 +914,9 @@ def main():
            aspect_ratios=aspect_ratio_sample_size,
        )

+        # A fixed canvas overrides the closest bucket size, so every batch is uniform.
+        fix_size = None if args.fix_sample_size is None else [args.video_sample_height, args.video_sample_width]
+
        def collate_fn(examples):
            new_examples = {}
            new_examples["pixel_values"] = []
@@ -851,9 +924,13 @@ def main():
            pixel_value = examples[0]["pixel_values"]
            f, h, w, c = np.shape(pixel_value)

-            closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size)
-            spatial_ratio = int(vae.config.spatial_compression_ratio)
-            closest_size = [int(x / spatial_ratio) * spatial_ratio for x in closest_size]
+            if fix_size is None:
+                closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size)
+                spatial_ratio = int(vae.config.spatial_compression_ratio)
+                closest_size = [int(x / spatial_ratio) * spatial_ratio for x in closest_size]
+                target_size = closest_size
+            else:
+                target_size = fix_size

            min_example_length = min([example["pixel_values"].shape[0] for example in examples])
            batch_video_length = int(min(args.video_sample_n_frames + sample_n_frames_bucket_interval, min_example_length))
@@ -865,14 +942,19 @@ def main():
                pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous()
                pixel_values = pixel_values / 255.0

-                if closest_size[0] / h > closest_size[1] / w:
-                    resize_size = closest_size[0], int(w * closest_size[0] / h)
+                if fix_size is None:
+                    if closest_size[0] / h > closest_size[1] / w:
+                        resize_size = closest_size[0], int(w * closest_size[0] / h)
+                    else:
+                        resize_size = int(h * closest_size[1] / w), closest_size[1]
                else:
-                    resize_size = int(h * closest_size[1] / w), closest_size[1]
+                    # Aspect-preserving cover resize, same as scripts/taehv/train_taehv_super.py.
+                    scale = max(fix_size[0] / h, fix_size[1] / w)
+                    resize_size = math.ceil(h * scale), math.ceil(w * scale)

                transform = transforms.Compose([
                    transforms.Resize(resize_size, interpolation=transforms.InterpolationMode.BILINEAR),
-                    transforms.CenterCrop(closest_size),
+                    transforms.CenterCrop(target_size),
                    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
                ])
                pixel_values = transform(pixel_values)[:batch_video_length]
@@ -900,6 +982,11 @@ def main():
            worker_init_fn=worker_init_fn(args.seed + accelerator.process_index)
        )

+    # Snapshot the sampler's empty bucket layout so every epoch can restore it. The shape is
+    # sampler-specific -- ImageVideoSampler holds flat lists while AspectRatioBatchImageVideoSampler
+    # holds one list per aspect ratio -- so the initial value is copied rather than rebuilt.
+    initial_bucket = copy.deepcopy(batch_sampler.bucket)
+
    # ==================== LR Scheduler ====================
    overrode_max_train_steps = False
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
@@ -989,8 +1076,18 @@ def main():
    )

    for epoch in range(first_epoch, args.num_train_epochs):
+        if global_step >= args.max_train_steps:
+            break
        train_loss = 0.0
+        train_dataloader.set_epoch(epoch)
        batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
+        # Indices that never filled a batch stay behind in `bucket` and would leak into the next
+        # epoch, silently changing the batch composition, so the layout is restored here.
+        # `_pos_start` is deliberately NOT reset: RandomSampler.__iter__ already zeroes it at the
+        # end of a full pass, and this script's resume relies on load_model_hook writing the saved
+        # position back into it. train_taehv_super.py can reset it unconditionally only because it
+        # pairs that with `resume_step` + accelerator.skip_first_batches, which this script lacks.
+        batch_sampler.bucket = copy.deepcopy(initial_bucket)
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(tae):
                pixel_values = batch["pixel_values"].to(weight_dtype)  # [B, F, C, H, W]
@@ -1015,10 +1112,11 @@ def main():
                    vae.to('cpu')
                    torch.cuda.empty_cache()

-                # 2. Student forward: TAE encode -> decode
+                # 2. Student forward: TAE encode -> decode (FP32 weights, autocast activations)
                # Use the unwrapped model since encode/decode are not exposed by DDP.
-                z_tae = unwrap_model(tae).encode(pixel_values).latent_dist.mode()
-                decoded = unwrap_model(tae).decode(z_tae).sample
+                with accelerator.autocast():
+                    z_tae = unwrap_model(tae).encode(pixel_values).latent_dist.mode()
+                    decoded = unwrap_model(tae).decode(z_tae).sample
                # Trim to the input temporal span in case the TAE padded extra frames.
                decoded = decoded[:, :, :pixel_values.shape[2]]

@@ -1129,6 +1227,12 @@ def main():
                log_validation(vae, tae, args, accelerator, weight_dtype, global_step)
                if args.use_ema:
                    ema_tae.restore(tae.parameters())
+        # Epoch-level break, mirroring train_taehv_super.py: the `break` inside the step loop only
+        # ends that epoch, so without this the remaining epochs would each run one extra step past
+        # max_train_steps and write a spurious checkpoint. Placed after the epoch validation so the
+        # final epoch still gets validated before exiting.
+        if global_step >= args.max_train_steps:
+            break

    # 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
@@ -1,6 +1,6 @@
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
-export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata.json"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export TAE_PATH="models/Diffusion_Transformer/taew2_2.safetensors"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
 # export NCCL_IB_DISABLE=1
@@ -9,30 +9,36 @@ NCCL_DEBUG=INFO

 accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
  --config_path="config/wan2.2/wan_civitai_5b.yaml" \
-  --pretrained_model_name_or_path=$MODEL_NAME \
-  --tae_path=$TAE_PATH \
-  --train_data_dir=$DATASET_NAME \
-  --train_data_meta=$DATASET_META_NAME \
-  --video_sample_size=512 \
-  --video_sample_stride=1 \
-  --video_sample_n_frames=33 \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --tae_path="$TAE_PATH" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --fix_sample_size 720 1280 \
+  --video_sample_strides 1 2 3 \
+  --video_sample_n_frames=25 \
+  --validation_n_frames=49 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
-  --dataloader_num_workers=4 \
-  --num_train_epochs=100 \
-  --checkpointing_steps=500 \
+  --dataloader_num_workers=8 \
+  --max_train_steps=10000 \
+  --checkpointing_steps=1000 \
+  --validation_steps=1000 \
+  --validation_paths "asset/inpaint_video.mp4" \
  --learning_rate=1e-04 \
  --lr_scheduler="constant_with_warmup" \
-  --lr_warmup_steps=100 \
+  --lr_warmup_steps=200 \
  --latent_loss_weight=1.0 \
+  --use_ema \
+  --ema_decay=0.999 \
  --seed=42 \
  --output_dir="output_dir_taehv_w2.2" \
  --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
  --adam_weight_decay=1e-2 \
  --adam_epsilon=1e-08 \
  --vae_mini_batch=1 \
  --max_grad_norm=1.0 \
-  --random_hw_adapt \
  --enable_bucket \
  --trainable_modules "." \
  --resume_from_checkpoint=latest
@@ -0,0 +1,946 @@
+"""Super decoder training, copied and adapted from train_taehv.py.
+
+Frozen official Wan2.2 encoder -> trainable Super decoder, using the four
+TinyVAE losses and the stage 4a B training recipe.
+"""
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2025 The VideoX-Fun Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+# ruff: noqa: E402 -- local imports follow the repository sys.path bootstrap.
+import argparse
+import json
+import logging
+import math
+import os
+import shutil
+import sys
+
+import diffusers
+import torch
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import EMAModel
+from diffusers.utils import check_min_version
+from omegaconf import OmegaConf
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+from torchvision import transforms
+
+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.data import ImageVideoSampler, RandomSampler, VideoDataset
+from videox_fun.models import AutoencoderKLWan3_8, AutoencoderTinyWan
+from videox_fun.utils.tqdm_bar import PauseAwareTqdm
+from videox_fun.utils.utils import save_videos_grid
+
+check_min_version("0.18.0.dev0")
+logger = get_logger(__name__, log_level="INFO")
+
+
+# TinyVAE reconstruction losses: RGB [0,1], BCTHW, FP32 reductions.
+def grid_loss_per_frame(pred, target, period=16):
+    """Reference: ZYT-World: A Real-Time Controllable World Model for
+    Closed-Loop Autonomous-Driving Simulation (https://arxiv.org/abs/2609.21712).
+
+    Penalize NCHW residual phase means (16x16 by default) after per-frame/channel
+    DC removal to suppress periodic grid artifacts.
+    """
+    n, c, h, w = pred.shape
+    h, w = h // period * period, w // period * period
+    if min(h, w) == 0:
+        raise ValueError("Images must be at least one grid period in size")
+    residual = pred.float()[..., :h, :w] - target.float()[..., :h, :w]
+    phase = residual.reshape(n, c, h // period, period, w // period, period).mean((2, 4))
+    return (phase - phase.mean((-2, -1), keepdim=True)).abs().mean((1, 2, 3))
+
+class TinyVAEReconstructionLoss(nn.Module):
+    def __init__(
+        self,
+        weights,
+        dino_repo=None,
+        dino_weights=None,
+        feature_batch_size=13,
+        feature_checkpointing=False,
+        grid_period=16,
+        lpips_model=None,
+        dino_model=None,
+    ):
+        super().__init__()
+        self.weights = {
+            k: float(weights[k]) for k in ("l1", "lpips", "dino", "grid")
+        }
+        if any(not math.isfinite(v) or v < 0 for v in self.weights.values()):
+            raise ValueError("Loss weights must be finite and nonnegative")
+        if not any(self.weights.values()) or min(feature_batch_size, grid_period) < 1:
+            raise ValueError("Need a positive objective, feature batch size and grid period")
+        if lpips_model is None:
+            import lpips
+
+            lpips_model = lpips.LPIPS(net="alex", version="0.1", verbose=False)
+        self.lpips = lpips_model.eval().requires_grad_(False)
+        if self.weights["dino"] and dino_model is None:
+            dino_model = torch.hub.load(
+                dino_repo or "facebookresearch/dinov2:7764ea0f912e53c92e82eb78a2a1631e92725fc8",
+                "dinov2_vitb14",
+                source="local" if dino_repo else "github",
+                pretrained=not bool(dino_weights),
+                trust_repo=True,
+                skip_validation=True,
+            )
+            if dino_weights:
+                dino_model.load_state_dict(
+                    torch.load(dino_weights, map_location="cpu", weights_only=True), strict=True
+                )
+        self.dino = dino_model.eval().requires_grad_(False) if self.weights["dino"] else None
+        self.feature_batch_size = feature_batch_size
+        self.feature_checkpointing = feature_checkpointing
+        self.grid_period = grid_period
+        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406])[None, :, None, None])
+        self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225])[None, :, None, None])
+        self.eval()
+
+    def train(self, mode=True):
+        return super().train(False)
+
+    def dino_features(self, x):
+        x = F.interpolate(x.float(), (224, 560), mode="bilinear", align_corners=False)
+        return self.dino.get_intermediate_layers(
+            (x - self.mean) / self.std,
+            n=(2, 5, 8, 11),
+            reshape=False,
+            return_class_token=False,
+            norm=True,
+        )
+
+    def feature_terms(self, pred, target):
+        lp = (
+            self.lpips(pred * 2 - 1, target * 2 - 1).float().flatten()
+            if self.weights["lpips"]
+            else pred.new_zeros(len(pred))
+        )
+        semantic = pred.new_zeros(len(pred))
+        if self.dino is not None:
+            with torch.no_grad():
+                reference = self.dino_features(target)
+            features = self.dino_features(pred)
+            semantic = torch.stack(
+                [
+                    (fp.float() - ft.float()).square().sum((-2, -1))
+                    / ft.float().square().sum((-2, -1)).clamp_min(1e-8)
+                    for fp, ft in zip(features, reference, strict=True)
+                ]
+            ).mean(0)
+        return lp, semantic
+
+    def forward(self, pred, target):
+        if pred.ndim != 5 or pred.shape != target.shape or pred.shape[1] != 3:
+            raise ValueError("Expected equal BCTHW RGB tensors in [0,1]")
+        p = pred.permute(0, 2, 1, 3, 4).flatten(0, 1).float()
+        x = target.permute(0, 2, 1, 3, 4).flatten(0, 1).float()
+        terms = {
+            "l1": (p - x).abs().mean(),
+            "grid": grid_loss_per_frame(p, x, self.grid_period).mean(),
+        }
+        chunks = []
+        for start in range(0, len(p), self.feature_batch_size):
+            inputs = (
+                p[start : start + self.feature_batch_size],
+                x[start : start + self.feature_batch_size],
+            )
+            if self.feature_checkpointing and torch.is_grad_enabled() and p.requires_grad:
+                chunks.append(checkpoint(self.feature_terms, *inputs, use_reentrant=False))
+            else:
+                chunks.append(self.feature_terms(*inputs))
+        terms.update(
+            lpips=torch.cat([v[0] for v in chunks]).mean(),
+            dino=torch.cat([v[1] for v in chunks]).mean(),
+        )
+        return sum(self.weights[k] * v for k, v in terms.items()), terms
+
+    @torch.no_grad()
+    def metrics(self, pred, target):
+        """Unweighted frame means; RGB Gaussian 11x11 valid SSIM, no DINO in eval."""
+        p = pred.permute(0, 2, 1, 3, 4).flatten(0, 1).float()
+        x = target.permute(0, 2, 1, 3, 4).flatten(0, 1).float()
+        rows = []
+        for start in range(0, len(p), self.feature_batch_size):
+            a, b = (
+                p[start : start + self.feature_batch_size],
+                x[start : start + self.feature_batch_size],
+            )
+            lp = self.lpips(a * 2 - 1, b * 2 - 1).float().flatten()
+            with torch.autocast(p.device.type, enabled=False):
+                q = torch.arange(11, device=p.device, dtype=torch.float32) - 5
+                kernel = torch.exp(-q.square() / (2 * 1.5**2))
+                kernel /= kernel.sum()
+                window = (kernel[:, None] * kernel[None, :])[None, None].expand(3, 1, 11, 11)
+                mu_a, mu_b = (F.conv2d(v, window, groups=3) for v in (a, b))
+                va = (F.conv2d(a.square(), window, groups=3) - mu_a.square()).clamp_min(0)
+                vb = (F.conv2d(b.square(), window, groups=3) - mu_b.square()).clamp_min(0)
+                cov = F.conv2d(a * b, window, groups=3) - mu_a * mu_b
+                ss = ((2 * mu_a * mu_b + 0.01**2) * (2 * cov + 0.03**2)) / (
+                    (mu_a.square() + mu_b.square() + 0.01**2) * (va + vb + 0.03**2)
+                )
+                psnr = -10 * (a - b).square().mean((1, 2, 3)).clamp_min(1e-12).log10()
+                grid = grid_loss_per_frame(a, b, self.grid_period)
+            rows.append(torch.stack((psnr, ss.mean((1, 2, 3)), lp, grid), dim=1))
+        return dict(zip(("psnr", "ssim", "lpips", "grid"), torch.cat(rows).mean(0).tolist()))
+
+@torch.no_grad()
+def log_validation(vae, tae, args, accelerator, weight_dtype, global_step,
+                   criterion, validation_dataloader, selection_state, ema_tae=None):
+    """Validate the Super decoder using official latents, then select inference weights."""
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        logger.info("Running validation...")
+        unwrapped_tae = accelerator.unwrap_model(tae)
+        trainable_params = list(unwrapped_tae.model.decoder.parameters())
+        was_training = unwrapped_tae.training
+        if ema_tae is not None:
+            ema_tae.store(trainable_params)
+            ema_tae.copy_to(trainable_params)
+        unwrapped_tae.eval()
+        rows = []
+        try:
+            for i, batch in enumerate(validation_dataloader):
+                if args.validation_max_samples and i >= args.validation_max_samples:
+                    break
+                pixel_values = batch["pixel_values"].to(accelerator.device)
+                with accelerator.autocast():
+                    gt_latents = vae.encode(pixel_values.to(weight_dtype)).latent_dist.mode()
+                    decoded = unwrapped_tae(gt_latents).sample[:, :, :pixel_values.shape[2]]
+                    decoded = decoded.float().add(1).mul(0.5)
+                    target = pixel_values.float().add(1).mul(0.5)
+                    rows.append(criterion.metrics(decoded, target))
+                if i == 0:
+                    save_videos_grid(
+                        torch.cat([target, decoded]),
+                        os.path.join(args.output_dir, f"sample/step-{global_step}.mp4"), n_rows=2,
+                    )
+            if not rows:
+                raise ValueError("Validation set is empty")
+            metrics = {key: sum(row[key] for row in rows) / len(rows) for key in rows[0]}
+            if selection_state["reference"] is None:
+                selection_state["reference"] = metrics
+            reference = selection_state["reference"]
+            eligible = (metrics["lpips"] <= reference["lpips"] * args.selection_lpips_ratio and
+                        metrics["grid"] <= reference["grid"] * args.selection_grid_ratio)
+            best = eligible and (selection_state["best_psnr"] is None or metrics["psnr"] > selection_state["best_psnr"])
+            if best:
+                selection_state["best_psnr"] = metrics["psnr"]
+                unwrapped_tae.save_pretrained(os.path.join(args.output_dir, "best", "taehv"))
+            row = {"step": global_step, "weights": "ema" if ema_tae is not None else "raw",
+                   "eligible": eligible, "best": best, **metrics}
+            with open(os.path.join(args.output_dir, "validation.jsonl"), "a") as file:
+                file.write(json.dumps(row) + "\n")
+            accelerator.log({f"validation/{key}": value for key, value in metrics.items()}, step=global_step)
+            logger.info(row)
+        finally:
+            if ema_tae is not None:
+                ema_tae.restore(trainable_params)
+            unwrapped_tae.train(was_training)
+    accelerator.wait_for_everyone()
+
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Training script for the Wan2.2 Super decoder.")
+    parser.add_argument(
+        "--config_path",
+        type=str,
+        default="config/wan2.2/wan_civitai_5b.yaml",
+        help="Path to the Wan2.2 model config yaml; vae_kwargs configures the frozen official VAE.",
+    )
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model directory (contains the full VAE weights).",
+    )
+    parser.add_argument(
+        "--train_data_dir",
+        type=str,
+        default="datasets/X-Fun-Videos-Demo",
+        help=("A folder containing the training data."),
+    )
+    parser.add_argument(
+        "--train_data_meta",
+        type=str,
+        default="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json",
+        help=("A JSON array containing the training video metadata."),
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="output_dir_taehv_super",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--train_batch_size", type=int, default=1, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument(
+        "--vae_mini_batch", type=int, default=1, help="Mini batch size for VAE encoding."
+    )
+    parser.add_argument("--num_train_epochs", type=int, default=100)
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-6,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="cosine",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=200, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--lr_min_ratio", type=float, default=0.2,
+        help="Final LR / peak LR for the cosine scheduler (stage 4a: 1e-6 / 5e-6).",
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training."
+        ),
+    )
+    parser.add_argument("--use_ema", action=argparse.BooleanOptionalAction, default=False, help="Whether to use EMA model.")
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help="Number of subprocesses to use for data loading.",
+    )
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.99, help="The beta2 parameter for the Adam optimizer.")
+    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(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="TensorBoard log directory.",
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="bf16",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose between fp16 and bf16.",
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        help='The integration to report the results and logs to.',
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help="Save a checkpoint of the training state every X updates.",
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=("Max number of checkpoints to store."),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="Whether training should be resumed from a previous checkpoint.",
+    )
+    parser.add_argument(
+        "--validation_steps",
+        type=int,
+        default=2000,
+        help="Run validation every X steps.",
+    )
+    parser.add_argument(
+        "--tracker_project_name",
+        type=str,
+        default="taehv-train",
+        help="The `project_name` argument passed to Accelerator.init_trackers.",
+    )
+    parser.add_argument(
+        "--video_sample_size",
+        type=int,
+        default=512,
+        help="Training video sample size.",
+    )
+    parser.add_argument(
+        "--video_sample_stride",
+        type=int,
+        default=1,
+        help="Sample stride of the video.",
+    )
+    parser.add_argument(
+        "--video_sample_n_frames",
+        type=int,
+        default=25,
+        help="Num frame of video. Must be 4k+1 to match the causal temporal compression.",
+    )
+    parser.add_argument(
+        "--tae_path",
+        type=str,
+        default=None,
+        help="Optional Super weights to warm-start from (taew2_2_super.pth/.safetensors, a directory, "
+             "or a previous checkpoint-xxx/taehv dir). If not given, train from scratch.",
+    )
+    parser.add_argument(
+        "--use_taehv_sequential", action="store_true",
+        help="Run the TAE in sequential (O(1)-memory) mode instead of parallel.",
+    )
+    parser.add_argument(
+        "--from_scratch",
+        action="store_true",
+        help="Initialize a new Super decoder without student weights.",
+    )
+    parser.add_argument(
+        "--fix_sample_size",
+        nargs=2,
+        type=int,
+        default=None,
+        help="Fix Sample size [height, width] to override `--video_sample_size` with a fixed "
+             "non-square shape.",
+    )
+    parser.add_argument(
+        "--video_sample_strides",
+        nargs="+",
+        type=int,
+        default=None,
+        help="Optional per-video random sampling strides; one is drawn for every clip.",
+    )
+    parser.add_argument(
+        "--validation_data_meta",
+        type=str,
+        default=None,
+        help="Held-out validation JSON index.",
+    )
+    parser.add_argument(
+        "--validation_n_frames",
+        type=int,
+        default=81,
+        help="Validation clip length, 4k+1.",
+    )
+    parser.add_argument(
+        "--validation_max_samples",
+        type=int,
+        default=0,
+        help="Validation sample cap; zero evaluates the complete index.",
+    )
+    parser.add_argument(
+        "--dino_repo",
+        type=str,
+        default=None,
+        help="Optional local pinned DINOv2 source directory.",
+    )
+    parser.add_argument(
+        "--dino_weights",
+        type=str,
+        default=None,
+        help="Optional local ViT-B/14 weights.",
+    )
+    parser.add_argument(
+        "--feature_batch_size",
+        type=int,
+        default=13,
+        help="Frame batch size for perceptual networks.",
+    )
+    parser.add_argument(
+        "--feature_checkpointing",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Checkpoint perceptual feature computation.",
+    )
+    parser.add_argument(
+        "--grid_period",
+        type=int,
+        default=16,
+        help="Spatial residual period for grid loss.",
+    )
+    parser.add_argument(
+        "--l1_loss_weight",
+        type=float,
+        default=0.1,
+        help="RGB L1 reconstruction loss weight.",
+    )
+    parser.add_argument(
+        "--lpips_loss_weight",
+        type=float,
+        default=0.1,
+        help="LPIPS reconstruction loss weight.",
+    )
+    parser.add_argument(
+        "--dino_loss_weight",
+        type=float,
+        default=0.02,
+        help="DINOv2 patch feature loss weight.",
+    )
+    parser.add_argument(
+        "--grid_loss_weight",
+        type=float,
+        default=1.0,
+        help="Periodic spatial residual loss weight.",
+    )
+    parser.add_argument(
+        "--ema_decay",
+        type=float,
+        default=0.999,
+        help="EMA decay for decoder parameters.",
+    )
+    parser.add_argument(
+        "--selection_lpips_ratio",
+        type=float,
+        default=1.02,
+        help="Maximum LPIPS relative to the first validation.",
+    )
+    parser.add_argument(
+        "--selection_grid_ratio",
+        type=float,
+        default=1.05,
+        help="Maximum grid loss relative to the first validation.",
+    )
+
+    args = parser.parse_args()
+
+    if args.fix_sample_size is not None:
+        args.video_sample_height, args.video_sample_width = args.fix_sample_size
+    else:
+        args.video_sample_height = args.video_sample_width = args.video_sample_size
+    if any(size < 32 or size % 16 for size in (args.video_sample_height, args.video_sample_width)):
+        parser.error("Video dimensions must be >=32 and divisible by 16")
+    if any(frames < 1 or (frames - 1) % 4 for frames in (args.video_sample_n_frames, args.validation_n_frames)):
+        parser.error("Video frame counts must be 4k+1")
+    if args.from_scratch and (args.tae_path or args.resume_from_checkpoint):
+        parser.error("--from_scratch cannot be combined with weight initialization or resume")
+    if args.max_train_steps is not None and args.max_train_steps < 1:
+        parser.error("max_train_steps must be positive")
+    if min(args.num_train_epochs, args.train_batch_size, args.vae_mini_batch,
+           args.gradient_accumulation_steps, args.checkpointing_steps, args.video_sample_stride,
+           args.feature_batch_size, args.grid_period) < 1:
+        parser.error("Epochs, batch sizes, checkpoint interval, stride and grid period must be positive")
+    if min(args.dataloader_num_workers, args.validation_steps, args.validation_max_samples, args.lr_warmup_steps) < 0:
+        parser.error("Workers, validation intervals/sample cap and warmup cannot be negative")
+    if args.video_sample_strides is not None and (not args.video_sample_strides or min(args.video_sample_strides) < 1):
+        parser.error("video_sample_strides must be positive")
+    if args.checkpoints_total_limit is not None and args.checkpoints_total_limit < 1:
+        parser.error("checkpoints_total_limit must be positive")
+    if not 0 <= args.ema_decay < 1 or not 0 <= args.adam_beta1 < 1 or not 0 <= args.adam_beta2 < 1:
+        parser.error("Invalid EMA decay or Adam betas")
+    if args.adam_weight_decay < 0:
+        parser.error("Adam weight decay cannot be negative")
+    if not 0 <= args.lr_min_ratio <= 1:
+        parser.error("lr_min_ratio must be between 0 and 1")
+    if any(not math.isfinite(value) or value <= 0 for value in
+           (args.learning_rate, args.adam_epsilon, args.max_grad_norm,
+            args.selection_lpips_ratio, args.selection_grid_ratio)):
+        parser.error("Learning rate, epsilon, gradient norm and selection ratios must be finite and positive")
+    loss_weights = (args.l1_loss_weight, args.lpips_loss_weight, args.dino_loss_weight,
+                    args.grid_loss_weight)
+    if not any(loss_weights) or any(not math.isfinite(value) or value < 0 for value in loss_weights):
+        parser.error("Loss weights must be finite and nonnegative, with at least one positive term")
+    return args
+
+
+
+
+def main():
+    args = parse_args()
+    logging_dir = os.path.join(args.output_dir, args.logging_dir)
+    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=None if args.report_to == "none" else args.report_to,
+        project_config=accelerator_project_config,
+    )
+    if accelerator.distributed_type.name not in ("NO", "MULTI_GPU", "MULTI_CPU"):
+        raise ValueError("This decoder trainer supports single device and ordinary DDP")
+
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        diffusers.utils.logging.set_verbosity_error()
+    set_seed(args.seed)
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+    if accelerator.is_main_process:
+        if args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+
+    # For mixed precision training, keep trainable parameters and Adam moments FP32.
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    # ==================== Load Config ====================
+    config = OmegaConf.load(args.config_path)
+    if config['vae_kwargs'].get('vae_type') != 'AutoencoderKLWan3_8':
+        raise ValueError("Super training requires the official 48-channel Wan2.2 VAE")
+
+    # ==================== Load Models ====================
+    vae = AutoencoderKLWan3_8.from_pretrained(
+        os.path.join(args.pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')),
+        additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
+    )
+    vae.eval()
+    vae.requires_grad_(False)
+    vae.to(accelerator.device, dtype=weight_dtype)
+
+    tae_load_kwargs = {"latent_channels": 48, "patch_size": 2, "arch_variant": "super"}
+    if args.tae_path is not None and not args.resume_from_checkpoint:
+        logger.info(f"Warm-starting Super TAE from checkpoint: {args.tae_path}")
+        tae = AutoencoderTinyWan.from_pretrained(
+            args.tae_path, additional_kwargs=tae_load_kwargs,
+        )
+    else:
+        logger.info("Initializing Super TAE (training state is restored below when resuming)")
+        tae = AutoencoderTinyWan(**tae_load_kwargs)
+    if tae.config.arch_variant != "super" or tae.config.latent_channels != 48 or tae.config.patch_size != 2:
+        raise ValueError("Expected a 48-channel, patch-size-2 Super checkpoint")
+    tae.parallel = not args.use_taehv_sequential
+    tae.float()
+    tae.requires_grad_(False)
+    tae.model.decoder.requires_grad_(True)
+    tae.train()
+
+    trainable_params = list(filter(lambda p: p.requires_grad, tae.parameters()))
+    total_params = sum(p.numel() for p in tae.parameters())
+    trainable_count = sum(p.numel() for p in trainable_params)
+    logger.info(f"Total params: {total_params / 1e6:.2f}M, Trainable params: {trainable_count / 1e6:.2f}M")
+    ema_tae = None
+    if args.use_ema:
+        ema_tae = EMAModel(
+            trainable_params, decay=args.ema_decay, min_decay=args.ema_decay, use_ema_warmup=False,
+        )
+
+    # ==================== Save/Load Hooks ====================
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            accelerator.unwrap_model(models[0]).save_pretrained(os.path.join(output_dir, "taehv"))
+            trainer_state = {"step": global_step, "epoch": epoch, "batch": resume_step, **selection_state}
+            with open(os.path.join(output_dir, "trainer_state.json"), "w") as file:
+                json.dump(trainer_state, file, indent=2)
+        # The model is saved in Diffusers format by this hook, as in train_taehv.py.
+        weights.pop()
+
+    def load_model_hook(models, input_dir):
+        while models:
+            model = accelerator.unwrap_model(models.pop())
+            load_model = AutoencoderTinyWan.from_pretrained(input_dir, subfolder="taehv")
+            model.register_to_config(**load_model.config)
+            model.load_state_dict(load_model.state_dict(), strict=True)
+            del load_model
+
+    accelerator.register_save_state_pre_hook(save_model_hook)
+    accelerator.register_load_state_pre_hook(load_model_hook)
+
+    # ==================== Optimizer ====================
+    optimizer = torch.optim.AdamW(
+        trainable_params, lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay, eps=args.adam_epsilon,
+    )
+
+    # ==================== Dataset ====================
+    train_dataset = VideoDataset(
+        args.train_data_meta, args.train_data_dir,
+        sample_size=args.video_sample_size,
+        sample_stride=args.video_sample_strides or args.video_sample_stride,
+        sample_n_frames=args.video_sample_n_frames,
+        enable_bucket=True, enable_inpaint=False, text_drop_ratio=0,
+    )
+    if not len(train_dataset):
+        raise ValueError("Training set is empty")
+
+    def collate_fn(examples):
+        """Fixed canvas, aspect-preserving resize/crop; align the shortest clip to 4k+1."""
+        frames = (min(len(e["pixel_values"]) for e in examples) - 1) // 4 * 4 + 1
+        height, width = args.video_sample_height, args.video_sample_width
+        videos = []
+        for example in examples:
+            x = torch.from_numpy(example["pixel_values"][:frames]).permute(0, 3, 1, 2).float()
+            h, w = x.shape[-2:]
+            scale = max(height / h, width / w)
+            resize_size = (math.ceil(h * scale), math.ceil(w * scale))
+            x = transforms.functional.resize(x, resize_size, antialias=True)
+            x = transforms.functional.center_crop(x, (height, width))
+            videos.append(x.div_(127.5).sub_(1))
+        return {"pixel_values": torch.stack(videos).permute(0, 2, 1, 3, 4).contiguous()}
+
+    batch_sampler_generator = torch.Generator().manual_seed(args.seed)
+    batch_sampler = ImageVideoSampler(
+        RandomSampler(train_dataset, generator=batch_sampler_generator),
+        train_dataset, args.train_batch_size, drop_last=True,
+    )
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_sampler=batch_sampler,
+        collate_fn=collate_fn,
+        persistent_workers=True if args.dataloader_num_workers != 0 else False,
+        num_workers=args.dataloader_num_workers,
+    )
+    validation_dataloader = None
+    if args.validation_data_meta is not None:
+        validation_dataset = VideoDataset(
+            args.validation_data_meta, args.train_data_dir,
+            sample_n_frames=args.validation_n_frames, sample_stride=1,
+            enable_bucket=True, enable_inpaint=False, text_drop_ratio=0,
+            random_start=False,
+        )
+        validation_dataloader = torch.utils.data.DataLoader(
+            validation_dataset, batch_size=1, collate_fn=collate_fn,
+        )
+
+    # ==================== LR Scheduler ====================
+    tae, optimizer, train_dataloader = accelerator.prepare(tae, optimizer, train_dataloader)
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if num_update_steps_per_epoch == 0:
+        raise ValueError("Training data does not contain a complete batch")
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler, optimizer=optimizer.optimizer,
+        num_warmup_steps=args.lr_warmup_steps,
+        num_training_steps=args.max_train_steps,
+    )
+    if args.lr_scheduler == "cosine":
+        # Keep Diffusers warmup/cosine and add stage 4a's nonzero final LR.
+        cosine_schedule = lr_scheduler.lr_lambdas[0]
+
+        def lr_lambda(step):
+            factor = cosine_schedule(step)
+            if step < args.lr_warmup_steps:
+                return factor
+            return args.lr_min_ratio + (1 - args.lr_min_ratio) * factor
+
+        lr_scheduler.lr_lambdas = [lr_lambda] * len(lr_scheduler.lr_lambdas)
+    # Register explicitly so it advances once per optimizer update, at any world size.
+    accelerator.register_for_checkpointing(lr_scheduler)
+    if args.use_ema:
+        ema_tae.to(accelerator.device)
+        accelerator.register_for_checkpointing(ema_tae)
+
+    # ==================== TinyVAE Losses ====================
+    loss_weights = {
+        "l1": args.l1_loss_weight,
+        "lpips": args.lpips_loss_weight,
+        "dino": args.dino_loss_weight,
+        "grid": args.grid_loss_weight,
+    }
+    criterion = TinyVAEReconstructionLoss(
+        loss_weights,
+        dino_repo=args.dino_repo, dino_weights=args.dino_weights,
+        feature_batch_size=args.feature_batch_size,
+        feature_checkpointing=args.feature_checkpointing, grid_period=args.grid_period,
+    ).to(accelerator.device)
+
+    if accelerator.is_main_process:
+        tracker_config = dict(vars(args))
+        keys_to_pop = [key for key, value in tracker_config.items() if isinstance(value, list)]
+        for key in keys_to_pop:
+            tracker_config.pop(key)
+        accelerator.init_trackers(args.tracker_project_name, tracker_config)
+
+    # ==================== Training ====================
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+    logger.info("***** Running Super decoder training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    logger.info(f"  Loss weights = {loss_weights}")
+    global_step = 0
+    first_epoch = 0
+    resume_step = 0
+    selection_state = {"reference": None, "best_psnr": None}
+
+    # Resume from checkpoint
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint != "latest":
+            path = args.resume_from_checkpoint
+        else:
+            dirs = [name for name in os.listdir(args.output_dir)
+                    if name.startswith("checkpoint-") and name.split("-")[-1].isdigit()]
+            dirs = sorted(dirs, key=lambda name: int(name.split("-")[-1]))
+            if not dirs:
+                raise FileNotFoundError(f"No checkpoints in {args.output_dir}")
+            path = os.path.join(args.output_dir, dirs[-1])
+        with open(os.path.join(path, "trainer_state.json")) as file:
+            trainer_state = json.load(file)
+        global_step = trainer_state["step"]
+        first_epoch = trainer_state["epoch"]
+        resume_step = trainer_state["batch"]
+        selection_state["reference"] = trainer_state["reference"]
+        selection_state["best_psnr"] = trainer_state["best_psnr"]
+        accelerator.print(f"Resuming from checkpoint {path}")
+        accelerator.load_state(path)
+        if args.use_ema:
+            ema_tae.to(accelerator.device)
+    initial_global_step = global_step
+    epoch = first_epoch
+    if accelerator.is_main_process:
+        OmegaConf.save(OmegaConf.create(vars(args)), os.path.join(args.output_dir, f"args-step-{global_step}.yaml"))
+    if validation_dataloader is not None and selection_state["reference"] is None:
+        log_validation(vae, tae, args, accelerator, weight_dtype, global_step,
+                       criterion, validation_dataloader, selection_state, ema_tae)
+
+    progress_bar = PauseAwareTqdm(
+        range(0, args.max_train_steps), initial=initial_global_step, desc="Steps",
+        disable=not accelerator.is_local_main_process,
+    )
+    optimizer.zero_grad()
+    for epoch in range(first_epoch, args.num_train_epochs):
+        if global_step >= args.max_train_steps:
+            break
+        train_loss = torch.zeros(1 + len(loss_weights), device=accelerator.device)
+        micro_steps = 0
+        train_dataloader.set_epoch(epoch)
+        batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
+        batch_sampler.sampler._pos_start = 0
+        batch_sampler.bucket = {"image": [], "video": []}
+        epoch_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step) if resume_step else train_dataloader
+        for step, batch in enumerate(epoch_dataloader, start=resume_step):
+            with accelerator.accumulate(tae):
+                pixel_values = batch["pixel_values"].to(accelerator.device)  # [B, C, F, H, W], FP32 target
+                bsz = pixel_values.shape[0]
+
+                # 1. Official latents, frozen.
+                with torch.no_grad(), accelerator.autocast():
+                    gt_latents_list = []
+                    for i in range(0, bsz, args.vae_mini_batch):
+                        encoded = vae.encode(pixel_values[i:i + args.vae_mini_batch].to(weight_dtype)).latent_dist.mode()
+                        gt_latents_list.append(encoded)
+                    gt_latents = torch.cat(gt_latents_list, dim=0).detach()
+
+                # 2. Student forward: official latent -> Super decoder, through DDP.
+                with accelerator.autocast():
+                    decoded = tae(gt_latents).sample[:, :, :pixel_values.shape[2]]
+                    decoded = decoded.float().add(1).mul(0.5)
+                    target = pixel_values.float().add(1).mul(0.5)
+                    loss, loss_terms = criterion(decoded, target)
+                if not torch.isfinite(loss):
+                    raise FloatingPointError(f"Nonfinite loss at step {global_step}")
+
+                # 3. Gather losses for logging and backpropagate.
+                train_loss += torch.stack([loss.detach(), *[loss_terms[key].detach() for key in loss_weights]])
+                micro_steps += 1
+                accelerator.backward(loss)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
+                optimizer.step()
+                optimizer.zero_grad()
+                if accelerator.sync_gradients and not accelerator.optimizer_step_was_skipped:
+                    lr_scheduler.step()
+                    if args.use_ema:
+                        ema_tae.step(trainable_params)
+
+            resume_step = step + 1
+            if accelerator.sync_gradients and not accelerator.optimizer_step_was_skipped:
+                progress_bar.update(1)
+                global_step += 1
+                averaged_loss = accelerator.reduce(train_loss / micro_steps, reduction="mean").tolist()
+                logs = dict(zip(("loss", *loss_weights), averaged_loss))
+                logs.update(step=global_step, frames=pixel_values.shape[2], lr=lr_scheduler.get_last_lr()[0])
+                accelerator.log(logs, step=global_step)
+                if accelerator.is_main_process:
+                    with open(os.path.join(args.output_dir, "train.jsonl"), "a") as file:
+                        file.write(json.dumps(logs) + "\n")
+                train_loss.zero_()
+                micro_steps = 0
+                progress_bar.set_postfix(step_loss=logs["loss"], lr=logs["lr"])
+
+                # Validation
+                if validation_dataloader is not None and args.validation_steps and global_step % args.validation_steps == 0:
+                    with progress_bar.paused():
+                        log_validation(vae, tae, args, accelerator, weight_dtype, global_step,
+                                       criterion, validation_dataloader, selection_state, ema_tae)
+
+                # Checkpointing
+                if global_step % args.checkpointing_steps == 0:
+                    if accelerator.is_main_process and args.checkpoints_total_limit is not None:
+                        checkpoints = [name for name in os.listdir(args.output_dir)
+                                       if name.startswith("checkpoint-") and name.split("-")[-1].isdigit()]
+                        checkpoints = sorted(checkpoints, key=lambda name: int(name.split("-")[-1]))
+                        num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
+                        for removing_checkpoint in checkpoints[:max(num_to_remove, 0)]:
+                            shutil.rmtree(os.path.join(args.output_dir, removing_checkpoint))
+                    accelerator.wait_for_everyone()
+                    save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
+                    with progress_bar.paused():
+                        accelerator.save_state(save_path)
+                    logger.info(f"Saved state to {save_path}")
+                if global_step >= args.max_train_steps:
+                    break
+        if global_step >= args.max_train_steps:
+            break
+        resume_step = 0
+    progress_bar.close()
+
+    # Final validation and save. Restore raw parameters after EMA validation.
+    if validation_dataloader is not None and (not args.validation_steps or global_step % args.validation_steps):
+        log_validation(vae, tae, args, accelerator, weight_dtype, global_step,
+                       criterion, validation_dataloader, selection_state, ema_tae)
+    accelerator.wait_for_everyone()
+    save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
+    accelerator.save_state(save_path)
+    logger.info(f"Saved state to {save_path}")
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth"
+# 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/taehv/train_taehv_super.py \
+  --config_path="config/wan2.2/wan_civitai_5b.yaml" \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --tae_path="$TAE_PATH" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --fix_sample_size 720 1280 \
+  --video_sample_strides 1 2 3 \
+  --video_sample_n_frames=25 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --checkpointing_steps=1000 \
+  --learning_rate=5e-06 \
+  --lr_scheduler="cosine" \
+  --lr_warmup_steps=200 \
+  --lr_min_ratio=0.2 \
+  --seed=42 \
+  --output_dir="output_dir_taehv_w2.2_super" \
+  --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
+  --adam_weight_decay=1e-2 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=1.0 \
+  --l1_loss_weight=0.1 \
+  --lpips_loss_weight=0.1 \
+  --dino_loss_weight=0.02 \
+  --grid_loss_weight=1.0 \
+  --grid_period=16 \
+  --feature_batch_size=13 \
+  --use_ema \
+  --ema_decay=0.999 \
+  --selection_lpips_ratio=1.02 \
+  --selection_grid_ratio=1.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 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.

 ---

@@ -172,13 +173,14 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --num_train_epochs=100 \
  --checkpointing_steps=200 \
  --learning_rate=2.0e-06 \
-  --learning_rate_critic=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_causal_forcing_dmd" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=0.0 \
  --adam_beta1=0.0 \
  --adam_beta2=0.999 \
@@ -190,8 +192,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --enable_bucket \
  --num_frame_per_block=3 \
  --use_kv_cache_training \
-  --denoising_step_indices_list 1000 667 334 1 \
-  --real_guidance_scale=6.0 \
+  --use_ema \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --randomize_step_indices \
  --fake_guidance_scale=0.0 \
  --gen_update_interval=5 \
@@ -212,7 +216,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 |
@@ -220,7 +224,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--num_train_epochs` | Number of training epochs | 100 |
 | `--checkpointing_steps` | Save checkpoint every N steps | 200 |
 | `--learning_rate` | Generator learning rate | 2e-06 |
-| `--learning_rate_critic` | Critic learning rate | 2e-06 |
+| `--learning_rate_critic` | Critic learning rate | 4.0e-07 |
 | `--lr_scheduler` | LR scheduler type | `constant_with_warmup` |
 | `--lr_warmup_steps` | LR warmup steps | 100 |
 | `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_dmd` |
@@ -233,6 +237,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--max_grad_norm` | Gradient clipping threshold | 10.0 |
 | `--trainable_modules` | Trainable modules (`"."` = all) | `"."` |
 | `--low_vram` | Enable low VRAM mode (offload VAE/text encoder) | - |
+| `--use_ema` | Maintain an EMA copy of the generator weights | - |

 **Video Sampling Parameters**:

@@ -253,8 +258,8 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh

 | Parameter | Description | Example Value |
 |-----------|-------------|---------------|
-| `--denoising_step_indices_list` | Denoising step indices (DMD core param). Example runs 4-step; parser default `[1000, 500]` = 2-step | `1000 667 334 1` |
-| `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 6.0 |
+| `--denoising_step_indices_list` | Denoising step indices (DMD core param). Example runs 4-step; parser default `[1000, 500]` = 2-step | `1000 750 500 250` |
+| `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 4.0 |
 | `--randomize_step_indices` | Randomize the denoising step indices during training | - |
 | `--fake_guidance_scale` | CFG scale for the fake-score (generator). 0.0 = no CFG | 0.0 |
 | `--gen_update_interval` | Generator update interval (generator updates every N critic steps) | 5 |
@@ -264,11 +269,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 remain off, and the DMD launcher additionally enables `--score_timestep_clamp`. AR diffusion and CCD do not gain GAN losses.
+
+The discriminator shares the **fake-score backbone**, with one learnable register token per selected feature layer. Its head is optimized by the critic optimizer, not the generator optimizer. Generator loss is `DMD + weighted GAN`; critic loss is `denoising + weighted GAN + R1 + R2`.
+
+| 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 (enabled in the DMD launcher) |
+
+For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` for the generator and `softplus(-D_real) + softplus(D_fake)` for the discriminator. Relativistic losses instead use `softplus(-(D_fake - D_real))` and `softplus(-(D_real - D_fake))`.
+
+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 +342,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 训练还会读取视频文件。

 ---

@@ -172,13 +173,14 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --num_train_epochs=100 \
  --checkpointing_steps=200 \
  --learning_rate=2.0e-06 \
-  --learning_rate_critic=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_causal_forcing_dmd" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=0.0 \
  --adam_beta1=0.0 \
  --adam_beta2=0.999 \
@@ -190,8 +192,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --enable_bucket \
  --num_frame_per_block=3 \
  --use_kv_cache_training \
-  --denoising_step_indices_list 1000 667 334 1 \
-  --real_guidance_scale=6.0 \
+  --use_ema \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --randomize_step_indices \
  --fake_guidance_scale=0.0 \
  --gen_update_interval=5 \
@@ -212,7 +216,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 |
@@ -220,7 +224,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--num_train_epochs` | 训练 epoch 数 | 100 |
 | `--checkpointing_steps` | 每 N 步保存 checkpoint | 200 |
 | `--learning_rate` | 生成器学习率 | 2e-06 |
-| `--learning_rate_critic` | 判别器学习率 | 2e-06 |
+| `--learning_rate_critic` | 判别器学习率 | 4.0e-07 |
 | `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
 | `--lr_warmup_steps` | 学习率预热步数 | 100 |
 | `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_dmd` |
@@ -233,6 +237,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--max_grad_norm` | 梯度裁剪阈值 | 10.0 |
 | `--trainable_modules` | 可训练模块（`"."` = 全部） | `"."` |
 | `--low_vram` | 低显存模式（卸载 VAE/文本编码器） | - |
+| `--use_ema` | 维护生成器权重的 EMA 副本 | - |

 **视频采样参数**：

@@ -253,8 +258,8 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh

 | 参数 | 说明 | 示例值 |
 |------|------|--------|
-| `--denoising_step_indices_list` | 去噪步骤索引（DMD 核心参数）。示例为 4 步；解析器默认 `[1000, 500]` = 2 步 | `1000 667 334 1` |
-| `--real_guidance_scale` | real-score（14B teacher）的 CFG scale | 6.0 |
+| `--denoising_step_indices_list` | 去噪步骤索引（DMD 核心参数）。示例为 4 步；解析器默认 `[1000, 500]` = 2 步 | `1000 750 500 250` |
+| `--real_guidance_scale` | real-score（14B teacher）的 CFG scale | 4.0 |
 | `--randomize_step_indices` | 训练时随机化去噪步骤索引 | - |
 | `--fake_guidance_scale` | fake-score（生成器）的 CFG scale。0.0 = 无 CFG | 0.0 |
 | `--gen_update_interval` | 生成器更新间隔（每 N 步判别器更新后更新 1 次生成器） | 5 |
@@ -264,23 +269,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 仍关闭，DMD 的 launcher 额外启用 `--score_timestep_clamp`。AR diffusion 和 CCD 不添加 GAN 损失。
+
+判别器共享 **fake-score backbone**，每个选定特征层对应一个可学习 register token。判别器头由 critic optimizer 更新，不加入 generator optimizer。生成器损失为 `DMD + 加权 GAN`，critic 损失为 `denoising + 加权 GAN + R1 + R2`。
+
+| 参数 | 说明 | 默认值 |
+|------|------|--------|
+| `--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` | 关闭（DMD launcher 已启用） |
+
+设 logits 为 `D_fake` 和 `D_real`，普通 GAN 的生成器损失为 `softplus(-D_fake)`，判别器损失为 `softplus(-D_real) + softplus(D_fake)`。相对论模式分别使用 `softplus(-(D_fake - D_real))` 和 `softplus(-(D_real - D_fake))`。
+
+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 +341,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(
@@ -37,6 +37,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -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(
@@ -41,6 +41,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --output_dir="output_dir_wan2.1_causal_forcing_ccd" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=0.0 \
  --adam_beta1=0.0 \
  --adam_beta2=0.999 \
@@ -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()
@@ -38,13 +38,14 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --num_train_epochs=100 \
  --checkpointing_steps=200 \
  --learning_rate=2.0e-06 \
-  --learning_rate_critic=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_causal_forcing_dmd" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=0.0 \
  --adam_beta1=0.0 \
  --adam_beta2=0.999 \
@@ -56,8 +57,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --enable_bucket \
  --num_frame_per_block=3 \
  --use_kv_cache_training \
-  --denoising_step_indices_list 1000 667 334 1 \
-  --real_guidance_scale=6.0 \
+  --use_ema \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --randomize_step_indices \
  --fake_guidance_scale=0.0 \
  --gen_update_interval=5 \
@@ -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.
@@ -36,6 +34,7 @@ This document provides the complete workflow for **Flex-Forcing Stage 2 — Dist
  - [3.5 Training with FSDP](#35-training-with-fsdp)
  - [3.6 Other Backends](#36-other-backends)
  - [3.7 Multi-Node Distributed Training](#37-multi-node-distributed-training)
+  - [3.8 Optional One-Forcing GAN, EMA, and Precision Controls](#38-optional-one-forcing-gan-ema-and-precision-controls)
 - [4. Inference Testing](#4-inference-testing)
  - [4.1 Inference Parameters](#41-inference-parameters)
  - [4.2 Text-to-Video (T2V) Inference](#42-text-to-video-t2v-inference)
@@ -109,6 +108,8 @@ DMD uses a **TextDataset** (`--train_mode="normal"`) — the generator creates i

 `train_distill.sh` defaults to `datasets/internal_datasets/metadata.json`. For a real run, swap in the **VidProM extended** prompts the paper uses.

+Enabling any GAN/R1/R2 weight (see **3.8**) switches the loader to **ImageVideoDataset**: captioned real videos are then required, cropped to the critic's fake input frame count.
+
 ---

 ## 3. Distillation Training
@@ -176,15 +177,16 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -206,11 +208,13 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```

-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/`.

@@ -222,6 +226,7 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
 |-----------|-------------|---------------|
 | `--config_path` | Component subpaths / `transformer_additional_kwargs` | `config/wan2.1/wan_civitai.yaml` |
 | `--pretrained_model_name_or_path` | Base model directory (1.3B) | `$MODEL_NAME` |
+| `--real_score_pretrained_model_name_or_path` | 14B real-score teacher directory; defaults to `--pretrained_model_name_or_path` | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
 | `--train_data_meta` | Prompt JSON (only the `"text"` field is read) | `$DATASET_META_NAME` |
 | `--ode_transformer_path` | Stage 1 ODE checkpoint for generator/critic init | `$STAGE1_CKPT` |
 | `--fix_sample_size` | Fixed `H W` of the training clip | `480 832` |
@@ -232,9 +237,9 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
 | `--dataloader_num_workers` | DataLoader workers | 8 |
 | `--num_train_epochs` | Number of training epochs (overridden by `--max_train_steps`) | 100 |
 | `--max_train_steps` | Hard step cap; the paper uses 600 | 600 |
-| `--checkpointing_steps` | Save checkpoint every N steps | 50 |
-| `--learning_rate` | Generator learning rate | `2e-06` |
-| `--learning_rate_critic` | Critic learning rate | `2e-06` |
+| `--checkpointing_steps` | Save checkpoint every N steps | 200 |
+| `--learning_rate` | Generator learning rate | `2.0e-06` |
+| `--learning_rate_critic` | Critic learning rate | `4.0e-07` |
 | `--lr_scheduler` | LR scheduler type | `constant_with_warmup` |
 | `--lr_warmup_steps` | LR warmup steps | 100 |
 | `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 4.5 |
@@ -243,7 +248,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 +261,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.
@@ -334,15 +329,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -364,6 +360,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -412,15 +410,16 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -442,10 +441,13 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```

+`--use_ema` is dropped here: `train_distill.py` raises `NotImplementedError` for EMA under DeepSpeed Zero-3.
+
 #### 3.6.2 Training without DeepSpeed and FSDP

 **This approach is not recommended because there is no memory-saving backend, which can easily cause out-of-memory errors**. We only provide the training shell for reference.
@@ -476,15 +478,16 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -506,6 +509,8 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -550,15 +555,16 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -580,6 +586,8 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -614,6 +622,44 @@ NCCL_DEBUG=INFO

 - **Data Synchronization**: All machines must be able to access the same data paths (NFS/shared storage)

+### 3.8 Optional One-Forcing GAN, EMA, and Precision Controls
+
+These options port the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms already used by the causal-forcing DMD trainer. All GAN/R1/R2 weights default to **0** and stay off in the launcher. `--use_ema`, `--fsdp_reduce_fp32` and `--score_timestep_clamp` default to **off** in argparse, but `train_distill.sh` enables all three; the DeepSpeed Zero-3 command in **3.6.1** drops `--use_ema`, which Zero-3 rejects.
+
+The discriminator shares the **fake-score backbone**, with one learnable register token per selected feature layer. Its head is optimized by the critic optimizer, not the generator optimizer. Generator loss is `DMD + weighted GAN`; critic loss is `denoising + weighted GAN + R1 + R2`.
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--use_ema` | EMA over the generator; checkpoints save a `transformer_ema` subfolder and validation samples from the EMA weights | off (enabled in `train_distill.sh`) |
+| `--gan_g_weight` | Non-saturating generator GAN loss weight | `0.0` |
+| `--gan_d_weight` | Discriminator GAN loss weight | `0.0` |
+| `--gan_feature_layers` | Fake-score block indices used by the head | `21 29` |
+| `--gan_block_ffn_dim` | Head cross-attention FFN dimension | `2048` |
+| `--gan_block_num_heads` | Head count; must divide the fake-score hidden dimension | `12` |
+| `--relativistic_discriminator` | Use paired real/fake logit differences | off |
+| `--r1_weight` | APT-style finite-difference penalty on noisy real latents | `0.0` |
+| `--r1_sigma` | R1 perturbation standard deviation | `0.01` |
+| `--r2_weight` | Normalized finite-difference penalty on noisy generated latents | `0.0` |
+| `--r2_sigma` | R2 perturbation standard deviation | `0.01` |
+| `--fsdp_reduce_fp32` | FSDP FP32 gradient reduction; parameter gathering stays low precision | off (enabled in `train_distill.sh`) |
+| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off (enabled in `train_distill.sh`) |
+
+For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` for the generator and `softplus(-D_real) + softplus(D_fake)` for the discriminator. Relativistic losses instead use `softplus(-(D_fake - D_real))` and `softplus(-(D_real - D_fake))`.
+
+R1 follows the source's APT form: `r1_weight * mean((D_real_perturbed - D_real)^2)`, **without dividing by sigma**. R2 follows `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`. These are finite-difference penalties, not input-gradient double-backward regularizers. Real/fake samples share timestep and noise; perturbations use independent noise. The non-relativistic critic runs fake and real as two sequential 1x-batch forwards with a backward after each (the memory-optimized split), the relativistic critic runs one packed 2x-batch forward, and R1/R2 each re-forward their own 2x batch after the main backward(s).
+
+To enable GAN, pass positive generator/discriminator weights and, optionally, R1/R2 weights to `train_distill.py`. Requirements:
+- Captioned real videos, not a prompt-only JSON. The real latents are cropped to the critic's fake input frame count, so the clip must cover at least that many latent frames.
+- `--gradient_accumulation_steps=1`; DeepSpeed is not supported by the GAN branch.
+- Valid feature-layer indices for the selected fake-score model. The default `21 29` layers and `12` heads fit Wan2.1-T2V-1.3B; other backbones may need different settings.
+- Enabling only `gan_g_weight` does not supply a discriminator classification-learning objective. Choose the discriminator weight deliberately; nonzero values are not a guaranteed quality improvement.
+
+Trainable fake-score and head master weights stay **FP32**, while forward uses the configured autocast/FSDP precision. Logits, softplus losses, and finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient communication precision, not the forward dtype or model architecture. The clamp changes score-training timesteps, not the generator rollout or validation schedule.
+
+Resume GAN checkpoints with the same GAN/head configuration. The head is saved with fake-score, including FP32 head tensors, and is recreated before loading. It is not part of the exported generator, so inference needs no discriminator or new GAN flags.
+
+Under `--use_ema`, the shadow stays in FP32 and mirrors the live model's mixed-precision policy; FSDP EMA requires FSDP1 and DeepSpeed Zero-3 is rejected.
+
 ---

 ## 4. Inference Testing
@@ -757,15 +803,16 @@ 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.

 ---

 ## 5. Additional Resources

 - **Flex-Forcing paper**: https://arxiv.org/abs/2607.03509
+- **One-Forcing (optional GAN mechanism)**: https://github.com/Aurora-edu/One-Forcing
 - **Stage 1**: [README_TRAIN_ODE.md](./README_TRAIN_ODE.md)
- **Shared chunking helpers**: `videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`, `sample_flexible_chunks`, `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)。
@@ -36,6 +34,7 @@
  - [3.5 使用 FSDP 训练](#35-使用-fsdp-训练)
  - [3.6 其他后端](#36-其他后端)
  - [3.7 多机分布式训练](#37-多机分布式训练)
+  - [3.8 可选的 One-Forcing GAN、EMA 与精度开关](#38-可选的-one-forcing-ganema-与精度开关)
 - [四、推理测试](#四推理测试)
  - [4.1 推理参数解析](#41-推理参数解析)
  - [4.2 文生视频（T2V）推理](#42-文生视频t2v推理)
@@ -109,6 +108,8 @@ DMD 使用 **TextDataset**（`--train_mode="normal"`）——生成器通过 rol

 `train_distill.sh` 默认指向 `datasets/internal_datasets/metadata.json`。真实训练请换成论文使用的 **VidProM extended** prompt 集。

+开启任何 GAN/R1/R2 权重（见 **3.8**）会把数据加载切到 **ImageVideoDataset**：此时需要带标注的真实视频，并会裁剪到 critic fake 输入的帧数。
+
 ---

 ## 三、蒸馏训练
@@ -176,15 +177,16 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -206,11 +208,13 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```

-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/`。

@@ -222,6 +226,7 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
 |-----------|-------------|---------------|
 | `--config_path` | 组件子路径 / `transformer_additional_kwargs` | `config/wan2.1/wan_civitai.yaml` |
 | `--pretrained_model_name_or_path` | 底座模型目录（1.3B） | `$MODEL_NAME` |
+| `--real_score_pretrained_model_name_or_path` | 14B 真实分数教师目录；默认回落 `--pretrained_model_name_or_path` | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
 | `--train_data_meta` | prompt JSON（只读 `"text"` 字段） | `$DATASET_META_NAME` |
 | `--ode_transformer_path` | 用于生成器 / critic 初始化的阶段 1 ODE checkpoint | `$STAGE1_CKPT` |
 | `--fix_sample_size` | 固定训练片段的 `H W` | `480 832` |
@@ -232,9 +237,9 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
 | `--dataloader_num_workers` | DataLoader 工作进程数 | 8 |
 | `--num_train_epochs` | 训练 epoch 数（被 `--max_train_steps` 覆盖） | 100 |
 | `--max_train_steps` | 硬性步数上限；论文用 600 | 600 |
-| `--checkpointing_steps` | 每 N 步保存一次 checkpoint | 50 |
-| `--learning_rate` | 生成器学习率 | `2e-06` |
-| `--learning_rate_critic` | critic 学习率 | `2e-06` |
+| `--checkpointing_steps` | 每 N 步保存一次 checkpoint | 200 |
+| `--learning_rate` | 生成器学习率 | `2.0e-06` |
+| `--learning_rate_critic` | critic 学习率 | `4.0e-07` |
 | `--lr_scheduler` | 学习率调度器类型 | `constant_with_warmup` |
 | `--lr_warmup_steps` | 学习率预热步数 | 100 |
 | `--real_guidance_scale` | 真实分数（14B 教师）的 CFG scale | 4.5 |
@@ -243,7 +248,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 +261,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 训练验证
@@ -345,15 +340,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -375,6 +371,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -417,15 +415,16 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -447,10 +446,13 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```

+此处去掉 `--use_ema`：DeepSpeed Zero-3 下 `train_distill.py` 会对 EMA 抛出 `NotImplementedError`。
+
 DeepSpeed-Zero-3 训练完成之后，得到 `checkpoint-xxx` 状态的保存权重，我们需要将权重文件合并为单个，可以使用提供的`scripts/zero_to_bf16.py`脚本。

 ```shell
@@ -485,15 +487,16 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -515,6 +518,8 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -560,15 +565,16 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -590,6 +596,8 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -616,6 +624,44 @@ export RANK=1
 > - Use `NCCL_DEBUG=INFO` to troubleshoot issues in multi-machine training.
 > - If there is no RDMA environment, set `export NCCL_IB_DISABLE=1` and `export NCCL_P2P_DISABLE=1` to avoid communication issues.

+### 3.8 可选的 One-Forcing GAN、EMA 与精度开关
+
+这些开关把 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 的机制（因果 forcing 的 DMD trainer 已在用）移植过来。所有 GAN/R1/R2 权重默认是 **0**，launcher 中也保持关闭；`--use_ema`、`--fsdp_reduce_fp32` 与 `--score_timestep_clamp` 在 argparse 中默认**关闭**，但 `train_distill.sh` 把三者都打开了；**3.6.1** 的 DeepSpeed Zero-3 命令则去掉了 Zero-3 不支持的 `--use_ema`。
+
+判别器与 **fake-score 骨干共享**，每个被选中的特征层有一个可学习的 register token。判别器头由 critic 优化器优化，而不是生成器优化器。生成器损失为 `DMD + 加权 GAN`；critic 损失为 `去噪 + 加权 GAN + R1 + R2`。
+
+| 参数 | 说明 | 默认 |
+|-----------|-------------|---------|
+| `--use_ema` | 对生成器做 EMA；checkpoint 额外保存 `transformer_ema` 子目录，验证样片采样自 EMA 权重 | 关（launcher 已启用） |
+| `--gan_g_weight` | 非饱和生成器 GAN 损失权重 | `0.0` |
+| `--gan_d_weight` | 判别器 GAN 损失权重 | `0.0` |
+| `--gan_feature_layers` | 判别器头使用的 fake-score block 索引 | `21 29` |
+| `--gan_block_ffn_dim` | 判别器头 cross-attention FFN 维度 | `2048` |
+| `--gan_block_num_heads` | 判别器头 head 数；必须整除 fake-score 隐维度 | `12` |
+| `--relativistic_discriminator` | 使用成对的 real/fake logit 差值 | 关 |
+| `--r1_weight` | 对含噪真实 latent 的 APT 式有限差分惩罚 | `0.0` |
+| `--r1_sigma` | R1 扰动标准差 | `0.01` |
+| `--r2_weight` | 对含噪生成 latent 的归一化有限差分惩罚 | `0.0` |
+| `--r2_sigma` | R2 扰动标准差 | `0.01` |
+| `--fsdp_reduce_fp32` | FSDP 梯度归约用 FP32；参数收集保持低精度 | 关（launcher 已启用） |
+| `--score_timestep_clamp` | 把 DMD、生成器 GAN 与 critic 的时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关（launcher 已启用） |
+
+对 logit `D_fake` 与 `D_real`，普通 GAN 损失是生成器 `softplus(-D_fake)`、判别器 `softplus(-D_real) + softplus(D_fake)`；相对论版本改用 `softplus(-(D_fake - D_real))` 与 `softplus(-(D_real - D_fake))`。
+
+R1 沿用来源的 APT 形式：`r1_weight * mean((D_real_perturbed - D_real)^2)`，**不除以 sigma**。R2 是 `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`。两者都是有限差分惩罚，不是输入梯度二重反向的正则。real/fake 样本共享时间步与噪声；扰动使用独立噪声。非相对论 critic 把 fake 与 real 拆成两次 1x batch 前向、各自一次 backward（省显存的拆分路径），相对论 critic 跑一次打包的 2x batch 前向，R1/R2 则在主 backward 之后各自再前向自己的 2x batch。
+
+启用 GAN 需要给 `train_distill.py` 传入正的生成器/判别器权重，可选再加 R1/R2 权重。要求：
+- 带标注的真实视频，不能是只有 prompt 的 JSON。真实 latent 会裁剪到 critic fake 输入的帧数，因此片段至少要覆盖那么多 latent 帧。
+- `--gradient_accumulation_steps=1`；GAN 分支不支持 DeepSpeed。
+- 特征层索引必须对所选 fake-score 模型有效。默认的 `21 29` 层与 `12` head 适配 Wan2.1-T2V-1.3B；换骨干可能需要不同配置。
+- 只开 `gan_g_weight` 不会给判别器提供分类学习目标。请有意识地选择判别器权重；非零值不保证质量提升。
+
+可训练的 fake-score 与判别器头 master 权重保持 **FP32**，前向使用配置的 autocast/FSDP 精度。logit、softplus 损失与有限差分归约都在 FP32 下计算。`--fsdp_reduce_fp32` 只改梯度通信精度，不改前向 dtype 或模型结构。时间步 clamp 只改分数训练的时间步，不改生成器 rollout 或验证调度。
+
+恢复 GAN checkpoint 时请使用相同的 GAN/判别器头配置。判别器头随 fake-score 一起保存（含 FP32 的 head 张量），加载前重建。它不属于导出的生成器，因此推理不需要判别器或新的 GAN 旗标。
+
+开启 `--use_ema` 时，shadow 保持 FP32 并与 live 模型的混合精度策略对齐；FSDP EMA 仅支持 FSDP1，DeepSpeed Zero-3 会被拒绝。
+
 ---

 ## 四、推理测试
@@ -746,15 +792,12 @@ 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)
+- **One-Forcing（可选 GAN 机制）**：https://github.com/Aurora-edu/One-Forcing
 - **官方 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
 ```

@@ -17,6 +17,7 @@

 import argparse
 import contextlib
+import dataclasses
 import gc
 import json
 import logging
@@ -76,18 +77,16 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512,
 from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
                               WanTransformer3DModel,
                               WanTransformer3DModel_FlexForcing,
-                               WanTransformer3DModel_SelfForcing)
-from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
-                                 WanFlexForcingPipeline,
-                                 WanSelfForcingPipeline)
+                               WanTransformer3DModel_SelfForcing,
+                               add_wan_gan_discriminator)
+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.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,
@@ -175,6 +174,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi
    return pred_for_score, score_num_frames, True


+def repeat_batch(value, repeats):
+    """Repeat a batched tensor or per-sample list along the batch dimension."""
+    if value is None or repeats == 1:
+        return value
+    if torch.is_tensor(value):
+        return torch.cat([value] * repeats)
+    return list(value) * repeats
+
+
 def filter_kwargs(cls, kwargs):
    import inspect
    sig = inspect.signature(cls.__init__)
@@ -299,6 +307,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 +320,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 +337,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 +381,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 +403,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 +542,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
@@ -806,6 +755,15 @@ def parse_args():
            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
        ),
    )
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.")
+    parser.add_argument(
+        "--fsdp_reduce_fp32", action="store_true",
+        help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
+    )
+    parser.add_argument(
+        "--score_timestep_clamp", action="store_true",
+        help="Clamp DMD, GAN and critic timesteps to [0.02, 0.98] * num_train_timesteps.",
+    )
    parser.add_argument(
        "--non_ema_revision",
        type=str,
@@ -1028,6 +986,16 @@ def parse_args():
        default=None,
        help=("If you want to load the ode-trained weight into generator transformer3d, input its path."),
    )
+    parser.add_argument(
+        "--real_score_pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        help=(
+            "Path to a non-causal pretrained model used as the DMD real_score teacher. "
+            "For CF Stage 3 this should point at Wan2.1-T2V-14B. If unset, falls back to "
+            "`--pretrained_model_name_or_path` (1.3B) — useful only for sanity testing."
+        ),
+    )
    parser.add_argument(
        "--vae_path",
        type=str,
@@ -1137,21 +1105,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 +1148,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",
@@ -1227,6 +1201,47 @@ def parse_args():
        help="Probability of applying teacher forcing per step (1.0 = always)"
    )

+    parser.add_argument(
+        "--gan_g_weight", type=float, default=0.0,
+        help="One-Forcing: non-saturating generator GAN loss weight (0 disables). Positive GAN/R1/R2 weights require real videos.",
+    )
+    parser.add_argument(
+        "--gan_d_weight", type=float, default=0.0,
+        help="One-Forcing: discriminator GAN loss weight added to critic denoising loss (0 disables).",
+    )
+    parser.add_argument(
+        "--gan_feature_layers", nargs="+", type=int, default=[21, 29],
+        help="fake_score block indices feeding the discriminator head, one register token per layer.",
+    )
+    parser.add_argument(
+        "--gan_block_ffn_dim", type=int, default=2048,
+        help="FFN dimension of the discriminator cross-attention blocks.",
+    )
+    parser.add_argument(
+        "--gan_block_num_heads", type=int, default=12,
+        help="Number of heads of the discriminator cross-attention blocks.",
+    )
+    parser.add_argument(
+        "--relativistic_discriminator", action="store_true",
+        help="Use relativistic GAN losses based on real/fake logit differences.",
+    )
+    parser.add_argument(
+        "--r1_weight", type=float, default=0.0,
+        help="Finite-difference R1 penalty weight on noisy real latents (APT form, not divided by sigma).",
+    )
+    parser.add_argument(
+        "--r1_sigma", type=float, default=0.01,
+        help="Perturbation standard deviation for the R1 penalty.",
+    )
+    parser.add_argument(
+        "--r2_weight", type=float, default=0.0,
+        help="Finite-difference R2 penalty weight on noisy generated latents (0.5 * squared logit difference / sigma squared)."
+    )
+    parser.add_argument(
+        "--r2_sigma", type=float, default=0.01,
+        help="Perturbation standard deviation for the R2 penalty.",
+    )
+
    args = parser.parse_args()
    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
    if env_local_rank != -1 and env_local_rank != args.local_rank:
@@ -1241,6 +1256,14 @@ def parse_args():

 def main():
    args = parse_args()
+    use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight))
+    if any(not math.isfinite(w) or w < 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight)):
+        raise ValueError("GAN and regularization weights must be finite and non-negative.")
+    if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1):
+        raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.")
+    if use_gan and (args.r1_weight > 0 and (not math.isfinite(args.r1_sigma) or args.r1_sigma <= 0)
+                    or args.r2_weight > 0 and (not math.isfinite(args.r2_sigma) or args.r2_sigma <= 0)):
+        raise ValueError("Enabled R1/R2 penalties require finite, positive perturbation standard deviations.")

    if args.report_to == "wandb" and args.hub_token is not None:
        raise ValueError(
@@ -1304,6 +1327,9 @@ def main():
        else:
            fsdp_stage = 0
        print(f"Using FSDP stage: {fsdp_stage}")
+        if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
+            fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
+            print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")

        args.use_fsdp = True
        if fsdp_stage == 3:
@@ -1431,8 +1457,11 @@ def main():
        transformer_additional_kwargs=generator_transformer_kwargs,
        low_cpu_mem_usage=True,
    ).to(weight_dtype)
+    # CF Stage 3 uses a 14B non-causal teacher as real_score; pass an explicit
+    # `--real_score_pretrained_model_name_or_path` to point at Wan2.1-T2V-14B.
+    real_score_base = args.real_score_pretrained_model_name_or_path or args.pretrained_model_name_or_path
    real_score_transformer3d = WanTransformer3DModel.from_pretrained(
-        os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+        os.path.join(real_score_base, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
        low_cpu_mem_usage=True,
    ).to(weight_dtype)
@@ -1533,6 +1562,41 @@ def main():
                param.requires_grad = True
                break

+    if use_gan:
+        if (not args.gan_feature_layers or any(i < 0 or i >= len(fake_score_transformer3d.blocks) for i in args.gan_feature_layers)
+                or args.gan_block_ffn_dim <= 0 or args.gan_block_num_heads <= 0
+                or fake_score_transformer3d.dim % args.gan_block_num_heads != 0):
+            raise ValueError("GAN feature layers must exist and the positive head count must divide the fake_score hidden dimension.")
+        # Keep trainable discriminator master weights in fp32; forward uses autocast.
+        add_wan_gan_discriminator(
+            fake_score_transformer3d,
+            args.gan_feature_layers,
+            args.gan_block_ffn_dim,
+            args.gan_block_num_heads,
+            torch.float32,
+            seed=args.seed or 0,
+        )
+
+    # Create EMA for the generator_transformer3d.
+    if args.use_ema:
+        if zero_stage == 3:
+            raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.")
+
+        if fsdp_plugin is not None:
+            if getattr(fsdp_plugin, "fsdp_version", 1) != 1:
+                raise NotImplementedError("FSDPEMA 仅支持 FSDP1。")
+            ema_module = generator_transformer_cls.from_config(generator_transformer3d.config)
+            ema_generator_transformer3d = FSDPEMA(
+                ema_module, source=generator_transformer3d,
+                accelerator=accelerator, fsdp_plugin=fsdp_plugin,
+            )
+        else:
+            ema_generator_transformer3d = EMAModel(
+                generator_transformer3d.parameters(),
+                model_cls=generator_transformer_cls,
+                model_config=generator_transformer3d.config,
+            )
+
    # `accelerate` 0.16.0 will have better support for customized saving
    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
        # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
@@ -1543,7 +1607,10 @@ def main():
                    from safetensors.torch import save_file

                    safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
-                    accelerate_state_dict = {k: v.to(dtype=weight_dtype) for k, v in accelerate_state_dict.items()}
+                    accelerate_state_dict = {
+                        k: v.to(dtype=torch.float32 if k.startswith("gan_head.") else weight_dtype)
+                        for k, v in accelerate_state_dict.items()
+                    }
                    save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})

                    with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
@@ -1560,7 +1627,14 @@ def main():
            # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
            def save_model_hook(models, weights, output_dir):
                if accelerator.is_main_process:
-                    models[0].save_pretrained(os.path.join(output_dir, "transformer"))
+                    model = accelerator.unwrap_model(models[0])
+                    save_dir = os.path.join(output_dir, "transformer")
+                    if hasattr(model, "gan_head"):
+                        from safetensors.torch import save_file
+                        model.save_config(save_dir)
+                        save_file(model.state_dict(), os.path.join(save_dir, "diffusion_pytorch_model.safetensors"), metadata={"format": "pt"})
+                    else:
+                        model.save_pretrained(save_dir)
                    if not args.use_deepspeed:
                        weights.pop()

@@ -1586,14 +1660,19 @@ def main():
                        else WanTransformer3DModel
                    )

-                    # load diffusers style into model
-                    load_model = model_cls.from_pretrained(
-                        input_dir, subfolder="transformer"
-                    )
-                    unwrapped.register_to_config(**load_model.config)
+                    # GAN heads are recreated from the same CLI configuration before resume.
+                    if hasattr(unwrapped, "gan_head"):
+                        from safetensors.torch import load_file
+                        unwrapped.load_state_dict(load_file(os.path.join(input_dir, "transformer", "diffusion_pytorch_model.safetensors")), strict=True)
+                    else:
+                        # load diffusers style into model
+                        load_model = model_cls.from_pretrained(
+                            input_dir, subfolder="transformer"
+                        )
+                        unwrapped.register_to_config(**load_model.config)

-                    unwrapped.load_state_dict(load_model.state_dict())
-                    del load_model
+                        unwrapped.load_state_dict(load_model.state_dict())
+                        del load_model

                pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
                if os.path.exists(pkl_path):
@@ -1607,6 +1686,32 @@ def main():
        accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
        accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)

+        if args.use_ema:
+            def save_ema_hook(models, weights, output_dir):
+                if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process:
+                    ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
+
+            def load_ema_hook(models, input_dir):
+                ema_path = os.path.join(input_dir, "transformer_ema")
+                if isinstance(ema_generator_transformer3d, FSDPEMA):
+                    ema_generator_transformer3d.load_pretrained(ema_path)
+                else:
+                    _, ema_kwargs = generator_transformer_cls.load_config(ema_path, return_unused_kwargs=True)
+                    load_model = generator_transformer_cls.from_pretrained(
+                        input_dir, subfolder="transformer_ema",
+                        transformer_additional_kwargs=generator_transformer_kwargs,
+                        low_cpu_mem_usage=True,
+                        torch_dtype=torch.float32,
+                    )
+                    load_model = EMAModel(load_model.parameters(), model_cls=generator_transformer_cls, model_config=load_model.config)
+                    load_model.load_state_dict(ema_kwargs)
+                    ema_generator_transformer3d.load_state_dict(load_model.state_dict())
+                    ema_generator_transformer3d.to(accelerator.device)
+                    del load_model
+
+            accelerator.register_save_state_pre_hook(save_ema_hook)
+            accelerator.register_load_state_pre_hook(load_ema_hook)
+
    if args.gradient_checkpointing:
        generator_transformer3d.enable_gradient_checkpointing()
        fake_score_transformer3d.enable_gradient_checkpointing()
@@ -1703,6 +1808,12 @@ def main():
                    print(f"Set {name} to lr : {args.learning_rate_critic / 2}")
                break

+    if use_gan:
+        fake_trainable_params_optim[0]['params'].extend(
+            param for name, param in fake_score_transformer3d.named_parameters()
+            if name.startswith("gan_head.") and name not in in_already
+        )
+
    if args.use_came:
        optimizer = optimizer_cls(
            trainable_params_optim,
@@ -1744,7 +1855,7 @@ def main():
        args.random_hw_adapt = False

    # Get the dataset
-    if args.train_mode != "normal" or args.use_teacher_forcing:
+    if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
        train_dataset = ImageVideoDataset(
            args.train_data_meta, args.train_data_dir,
            video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames, 
@@ -1772,7 +1883,7 @@ def main():

        return length_to_frame_num

-    if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing:
+    if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan:
        aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
        batch_sampler_generator = torch.Generator().manual_seed(args.seed)
        batch_sampler = AspectRatioBatchImageVideoSampler(
@@ -2115,6 +2226,9 @@ def main():
    if args.train_mode != "normal":
        clip_image_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)

+    if args.use_ema:
+        ema_generator_transformer3d.to(accelerator.device)
+
    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if overrode_max_train_steps:
@@ -2205,6 +2319,8 @@ def main():
        vae_stream_2 = None

    idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
+    min_score_timestep = int(0.02 * noise_scheduler.config.num_train_timesteps)
+    max_score_timestep = int(0.98 * noise_scheduler.config.num_train_timesteps)

    def randomize_denoising_step_indices(
        denoising_step_indices_list,
@@ -2265,6 +2381,8 @@ def main():
        # must be averaged by contribution count, not by gradient_accumulation_steps.
        train_gen_log_count = 0
        train_denoising_loss = 0.0
+        train_gan_g_loss = 0.0
+        train_gan_d_loss = 0.0
        batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
        for step, batch in enumerate(train_dataloader):
            # Data batch sanity check
@@ -2397,7 +2515,7 @@ def main():
                                pixel_values_bs = pixel_values_bs.sample()
                                new_pixel_values.append(pixel_values_bs)
                            return torch.cat(new_pixel_values, dim = 0)
-                        if args.use_teacher_forcing:
+                        if args.use_teacher_forcing or use_gan:
                            clean_latents = _batch_encode_vae(pixel_values)
                        else:
                            clean_latents = None
@@ -2429,7 +2547,7 @@ def main():
                                clip_context.append(_clip_context)
                            clip_context = torch.cat(clip_context)

-                    if args.use_teacher_forcing:
+                    if args.use_teacher_forcing or use_gan:
                        target_shape = clean_latents.size()
                    else:
                        target_shape = mask_latents.size()
@@ -3177,6 +3295,8 @@ def main():
                    # Common code for both KV cache and block mask training
                    indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
                    generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
+                    if args.score_timestep_clamp:
+                        generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep)
                    generator_denoised_input = add_noise(
                        generator_pred,
                        torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
@@ -3282,17 +3402,53 @@ def main():
                    train_dmd_loss += avg_dmd_loss.item()
                    train_gen_log_count += 1

+                    # --- One-Forcing Generator GAN Loss ---
+                    generator_loss = dmd_loss
+                    if args.gan_g_weight > 0:
+                        indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
+                        gan_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
+                        if args.score_timestep_clamp:
+                            gan_timestep = gan_timestep.clamp(min_score_timestep, max_score_timestep)
+                        gan_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng)
+                        gan_inputs = [add_noise(generator_pred, gan_noise, gan_timestep)]
+                        if args.relativistic_discriminator:
+                            real_latents = clean_latents[:, :, :generator_pred.size(2)].to(generator_pred.dtype)
+                            gan_inputs.append(add_noise(real_latents, gan_noise, gan_timestep))
+                        num_gan_inputs = len(gan_inputs)
+                        with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                            _, gan_logits = fake_score_transformer3d(
+                                x=torch.cat(gan_inputs).to(dtype=weight_dtype),
+                                context=repeat_batch(prompt_embeds, num_gan_inputs),
+                                t=gan_timestep.repeat(num_gan_inputs),
+                                seq_len=seq_len,
+                                y=repeat_batch(inpaint_latents, num_gan_inputs) if args.train_mode != "normal" else None,
+                                clip_fea=repeat_batch(clip_context, num_gan_inputs) if args.train_mode != "normal" else None,
+                                gan_classify=True,
+                            )
+                        gan_logits = gan_logits.float().chunk(num_gan_inputs)
+                        if args.relativistic_discriminator:
+                            gan_g_loss = F.softplus(-(gan_logits[0] - gan_logits[1])).mean()
+                        else:
+                            gan_g_loss = F.softplus(-gan_logits[0]).mean()
+                        gan_g_loss = gan_g_loss * args.gan_g_weight
+                        train_gan_g_loss += accelerator.gather(gan_g_loss.detach().repeat(args.train_batch_size)).mean().item()
+                        generator_loss = generator_loss + gan_g_loss
+
                    if args.low_vram:
                        real_score_transformer3d = real_score_transformer3d.to("cpu")
-                        fake_score_transformer3d = fake_score_transformer3d.to("cpu")
+                        if args.gan_g_weight == 0:
+                            fake_score_transformer3d = fake_score_transformer3d.to("cpu")
                        torch.cuda.empty_cache()

-                    accelerator.backward(dmd_loss)
+                    accelerator.backward(generator_loss)
                    if accelerator.sync_gradients:
                        accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()
+                    if args.gan_g_weight > 0:
+                        # Discard discriminator gradients from the generator GAN backward.
+                        critic_optimizer.zero_grad()
                    
                    if args.low_vram:
                        fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device)
@@ -3683,6 +3839,8 @@ def main():

                indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
                critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
+                if args.score_timestep_clamp:
+                    critic_timestep = critic_timestep.clamp(min_score_timestep, max_score_timestep)
                critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng)

                fake_score_denoised_input = add_noise(
@@ -3691,15 +3849,25 @@ def main():
                    critic_timestep
                )

-                with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
-                    fake_score_denoised_output = fake_score_transformer3d(
-                        x=fake_score_denoised_input,
-                        context=prompt_embeds,
-                        t=critic_timestep,
-                        seq_len=seq_len,
-                        y=inpaint_latents if args.train_mode != "normal" else None,
-                        clip_fea=clip_context if args.train_mode != "normal" else None,
-                    )
+                # Generated and real latents share critic_noise / critic_timestep. Prepare the
+                # real input and the R1/R2 perturbations once; they are consumed by the critic
+                # forward below and by the R1/R2 penalties after the main backward(s).
+                real_noisy = None
+                r1_perturbation = None
+                r2_perturbation = None
+                if use_gan:
+                    real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype)
+                    real_noisy = add_noise(real_latents, critic_noise, critic_timestep)
+                    del real_latents
+                    if args.r1_weight > 0:
+                        r1_perturbation = args.r1_sigma * torch.randn(
+                            real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng
+                        )
+                    if args.r2_weight > 0:
+                        r2_perturbation = args.r2_sigma * torch.randn(
+                            fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype,
+                            device=fake_score_denoised_input.device, generator=torch_rng
+                        )

                def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
                    noise_pred = noise_pred.float()
@@ -3713,11 +3881,136 @@ def main():
                    final_loss = masked_loss.mean()
                    return final_loss

-                denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
-                avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
-                train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
-                
-                accelerator_fake_score_transformer3d.backward(denoising_loss)
+                # --- One-Forcing Discriminator GAN Loss ---
+                if use_gan and not args.relativistic_discriminator:
+                    # Memory-optimized critic: the non-relativistic gan_d_loss =
+                    # softplus(-real) + softplus(fake) is separable, so run fake and real as two
+                    # sequential 1x-batch forwards and backward the fake pass first. Backwarding
+                    # fake frees its gradient-checkpointing block inputs before the real forward,
+                    # dropping the packed fake+real+R1+R2 forward peak (up to 4x batch) down to
+                    # 1x batch here; R1/R2 re-forward their own 2x batch below. Both backward(s)
+                    # accumulate into fake_score.grad, giving the identical update as the packed
+                    # single-forward critic.
+                    # -- Fake pass (1x batch): denoising output + fake discriminator logit --
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        fake_score_denoised_output, fake_gan_logits = fake_score_transformer3d(
+                            x=fake_score_denoised_input,
+                            context=prompt_embeds,
+                            t=critic_timestep,
+                            seq_len=seq_len,
+                            y=inpaint_latents if args.train_mode != "normal" else None,
+                            clip_fea=clip_context if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                    denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
+                    avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
+                    train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
+
+                    fake_logit = fake_gan_logits.float()
+                    fake_d_loss = F.softplus(fake_logit).mean() * args.gan_d_weight
+                    critic_loss = denoising_loss + fake_d_loss
+                    accelerator_fake_score_transformer3d.backward(critic_loss)
+                    # Fake graph freed: the transformer blocks' stored inputs are released here.
+                    fake_d_detached = fake_d_loss.detach()
+                    del critic_loss, fake_gan_logits, fake_logit, fake_d_loss, fake_score_denoised_output
+
+                    # -- Real pass (1x batch): discriminator logit only --
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        _, real_gan_logits = fake_score_transformer3d(
+                            x=real_noisy,
+                            context=prompt_embeds,
+                            t=critic_timestep,
+                            seq_len=seq_len,
+                            y=inpaint_latents if args.train_mode != "normal" else None,
+                            clip_fea=clip_context if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                    real_logit = real_gan_logits.float()
+                    real_d_loss = F.softplus(-real_logit).mean() * args.gan_d_weight
+                    accelerator_fake_score_transformer3d.backward(real_d_loss)
+                    # Real graph freed. Combine the two detached halves for logging.
+                    gan_d_loss = fake_d_detached + real_d_loss.detach()
+                    train_gan_d_loss += accelerator.gather(gan_d_loss.repeat(args.train_batch_size)).mean().item()
+                    del fake_d_detached, real_gan_logits, real_logit, real_d_loss
+                else:
+                    # Combined path: 2x-batch forward for relativistic GAN (both logits coupled by
+                    # softplus(-(real - fake))), or 1x-batch forward when GAN is disabled.
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        if use_gan:
+                            fake_score_denoised_output, gan_logits = fake_score_transformer3d(
+                                x=torch.cat([fake_score_denoised_input, real_noisy]),
+                                context=repeat_batch(prompt_embeds, 2),
+                                t=critic_timestep.repeat(2),
+                                seq_len=seq_len,
+                                y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
+                                clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
+                                gan_classify=True,
+                            )
+                            fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0]
+                            gan_logits = gan_logits.float().chunk(2)
+                        else:
+                            fake_score_denoised_output = fake_score_transformer3d(
+                                x=fake_score_denoised_input,
+                                context=prompt_embeds,
+                                t=critic_timestep,
+                                seq_len=seq_len,
+                                y=inpaint_latents if args.train_mode != "normal" else None,
+                                clip_fea=clip_context if args.train_mode != "normal" else None,
+                            )
+
+                    denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
+                    avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
+                    train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
+
+                    critic_loss = denoising_loss
+                    if use_gan:
+                        fake_logit, real_logit = gan_logits[0], gan_logits[1]
+                        if args.relativistic_discriminator:
+                            gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
+                        else:
+                            gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
+                        gan_d_loss = gan_d_loss * args.gan_d_weight
+                        train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
+                        critic_loss = critic_loss + gan_d_loss
+
+                    accelerator_fake_score_transformer3d.backward(critic_loss)
+                    del critic_loss
+
+                # R1/R2 gradient penalties: each re-forwards its own 2x batch after the main
+                # backward(s) and finite-differences the perturbed logit against the clean one,
+                # matching the packed-formula value while keeping the peak at 2x batch.
+                if use_gan and args.r1_weight > 0:
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        _, r1_logits = fake_score_transformer3d(
+                            x=torch.cat([real_noisy + r1_perturbation, real_noisy]),
+                            context=repeat_batch(prompt_embeds, 2),
+                            t=critic_timestep.repeat(2),
+                            seq_len=seq_len,
+                            y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
+                            clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                        r1_logits = r1_logits.float().chunk(2)
+                        r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean()
+                    accelerator_fake_score_transformer3d.backward(r1_loss)
+                    del r1_loss
+
+                if use_gan and args.r2_weight > 0:
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        _, r2_logits = fake_score_transformer3d(
+                            x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]),
+                            context=repeat_batch(prompt_embeds, 2),
+                            t=critic_timestep.repeat(2),
+                            seq_len=seq_len,
+                            y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
+                            clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                        r2_logits = r2_logits.float().chunk(2)
+                        r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean()
+                    accelerator_fake_score_transformer3d.backward(r2_loss)
+                    del r2_loss
+
                if accelerator_fake_score_transformer3d.sync_gradients:
                    accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm)
                critic_optimizer.step()
@@ -3731,12 +4024,19 @@ def main():
            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:

+                if args.use_ema:
+                    ema_generator_transformer3d.step(generator_transformer3d.parameters())
                progress_bar.update(1)
                global_step += 1
-                accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step)
+                logs = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}
+                if use_gan:
+                    logs.update({"train_gan_g_loss": train_gan_g_loss / max(train_gen_log_count, 1), "train_gan_d_loss": train_gan_d_loss})
+                accelerator.log(logs, step=global_step)
                train_dmd_loss = 0.0
                train_gen_log_count = 0
                train_denoising_loss = 0.0
+                train_gan_g_loss = 0.0
+                train_gan_d_loss = 0.0

                if global_step % args.checkpointing_steps == 0:
                    if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
@@ -3779,6 +4079,10 @@ def main():

                if args.validation_prompts is not None and global_step % args.validation_steps == 0:
                    with progress_bar.paused():
+                        if args.use_ema:
+                            # Store the generator parameters temporarily and load the EMA parameters to perform inference.
+                            ema_generator_transformer3d.store(generator_transformer3d.parameters())
+                            ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
                        log_validation(
                            vae,
                            text_encoder,
@@ -3791,6 +4095,9 @@ def main():
                            weight_dtype,
                            global_step,
                        )
+                        if args.use_ema:
+                            # Switch back to the original generator parameters.
+                            ema_generator_transformer3d.restore(generator_transformer3d.parameters())

            logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
            progress_bar.set_postfix(**logs)
@@ -3800,6 +4107,10 @@ def main():

        if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
            with progress_bar.paused():
+                if args.use_ema:
+                    # Store the generator parameters temporarily and load the EMA parameters to perform inference.
+                    ema_generator_transformer3d.store(generator_transformer3d.parameters())
+                    ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
                log_validation(
                    vae,
                    text_encoder,
@@ -3812,6 +4123,9 @@ def main():
                    weight_dtype,
                    global_step,
                )
+                if args.use_ema:
+                    # Switch back to the original generator parameters.
+                    ema_generator_transformer3d.restore(generator_transformer3d.parameters())

    # Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever
    # something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto
@@ -28,15 +28,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_flex_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -58,5 +59,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --randomize_step_indices \
  --flow_euler_rollout \
  --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
  --num_frame_per_block=3 \
  --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
@@ -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
@@ -0,0 +1,356 @@
+# Wan2.1 One-Forcing: Stable One-Step Autoregressive Video Generation
+
+This document provides the complete workflow for **One-Forcing** on Wan2.1-T2V-1.3B, a recipe ported from [Aurora-edu/One-Forcing](https://github.com/Aurora-edu/One-Forcing).
+
+> **What is One-Forcing?**
+>
+> One-Forcing (*Towards Stable One-Step Autoregressive Video Generation*) is **Causal-Forcing Stage 3 — Distribution Matching Distillation (DMD)** *plus* a **shared noised-latent adversarial critic**. The warm-start generator is produced by the first two Causal-Forcing stages, and One-Forcing then polishes it into a **stable one-step, frame-wise autoregressive** video model:
+>
+> 1. **Stage 1 — AR Diffusion** (`scripts/wan2.1_causal_forcing/train_ar_diffusion.py`): Train the model with **teacher forcing** on clean video latents, learning to denoise chunk-by-chunk (or frame-by-frame) in a causal autoregressive manner. This produces a strong AR backbone. See [README_TRAIN_AR_DIFFUSION.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md).
+> 2. **Stage 2 — Causal Consistency Distillation (CCD)** (`scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py`): Distill the multi-step AR model into a **one-step-per-block** consistency model using an EMA teacher with CFG guidance. Its output checkpoint is the **warm-start generator** for One-Forcing. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md).
+> 3. **Stage 3 — One-Forcing** (`train_one_forcing.py`, this document), which adds two mechanisms on top of DMD:
+>    - **DMD core**: the generator matches the real distribution through a 14B real-score teacher and a 1.3B fake-score critic (`train_causal_dmd.py`).
+>    - **Adversarial critic**: a discriminator head rides on the **fake-score backbone** (one register token per selected feature layer). Both the DMD denoising loss and the adversarial loss operate on the **same noised generated latent**, and the critic uses **one shared noise/timestep** for fake and real samples.
+>    - **One step**: `denoising_step_list = [1000]` — each frame is generated in a single denoising step (the "One" in One-Forcing).
+>
+> Instead of training Stage 1/2 yourself, you can also warm-start from the **official ODE-initialized checkpoint** released by the One-Forcing authors (see Section 3).
+>
+> **The trainer** `scripts/wan2.1_one_forcing/train_one_forcing.py` implements the One-Forcing generator/critic GAN losses, the relativistic option, and finite-difference R1/R2 penalties. For the underlying DMD stage, see [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md).
+
+---
+
+## Table of Contents
+- [1. Prerequisites](#1-prerequisites)
+- [2. Environment Setup](#2-environment-setup)
+- [3. Download Pretrained Models](#3-download-pretrained-models)
+- [4. Prepare Training Data](#4-prepare-training-data)
+  - [4.1 Quick Demo Dataset](#41-quick-demo-dataset)
+  - [4.2 Dataset Structure](#42-dataset-structure)
+  - [4.3 metadata.json Format](#43-metadatajson-format)
+- [5. Training](#5-training)
+  - [5.1 Quick Start](#51-quick-start)
+  - [5.2 Recipe Parameters](#52-recipe-parameters)
+  - [5.3 GAN Prerequisites](#53-gan-prerequisites)
+  - [5.4 Optional Regularizers (Relativistic / R1 / R2)](#54-optional-regularizers-relativistic--r1--r2)
+- [6. Inference](#6-inference)
+  - [6.1 Quick Start](#61-quick-start)
+  - [6.2 1-Step vs 4-Step Selector](#62-1-step-vs-4-step-selector)
+  - [6.3 First-Frame Schedule](#63-first-frame-schedule)
+- [7. Checkpoint Round-Trip](#7-checkpoint-round-trip)
+- [8. Additional Resources](#8-additional-resources)
+
+---
+
+## 1. Prerequisites
+
+One-Forcing is the final polish on top of a distilled generator, so it needs:
+
+1. A **warm-start generator checkpoint** — a distilled/ODE generator that the DMD + adversarial critic then polishes (e.g. a **Causal-Forcing Stage-2 CCD** checkpoint, or the official ODE-initialized checkpoint in Section 3).
+2. A **Wan2.1-T2V-14B** model as the DMD real-score teacher.
+3. **Captioned real videos** — the adversarial critic needs clean real latents, so a prompt-only JSON is **not** sufficient (unlike plain DMD).
+
+---
+
+## 2. Environment Setup
+
+**Method 1: Using requirements.txt**
+
+```bash
+pip install -r requirements.txt
+```
+
+**Method 2: Manual Installation**
+
+```bash
+pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
+pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
+pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
+pip install yunchang xfuser modelscope openpyxl
+pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
+pip install opencv-python-headless
+pip install deepspeed==0.17.0 numpy==1.26.4
+```
+
+**Method 3: Using Docker**
+
+```bash
+# pull image
+docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+
+# enter image
+docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+```
+
+---
+
+## 3. Download Pretrained Models
+
+One-Forcing requires **two** pretrained models:
+
+- **Wan2.1-T2V-1.3B**: base model for the generator/critic.
+- **Wan2.1-T2V-14B**: the non-causal real-score teacher used by DMD to compute the real distribution score.
+
+```bash
+# Create model directory
+mkdir -p models/Diffusion_Transformer
+
+# Download Wan2.1 T2V 1.3B (student base model)
+modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
+
+# Download Wan2.1 T2V 14B (DMD real-score teacher)
+modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B
+```
+
+The official One-Forcing checkpoints are hosted on HuggingFace (`JiaqiFeng/OneForcing`):
+
+```bash
+# Warm-start generator for training (the official ODE-initialized checkpoint),
+# loaded via --ode_transformer_path.
+hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
+
+# Trained One-Forcing generator for inference (see Section 6).
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+If you trained Causal-Forcing Stage 1/2 yourself, you can instead point `--ode_transformer_path` at your own Stage-2 CCD checkpoint.
+
+---
+
+## 4. Prepare Training Data
+
+One-Forcing enables the GAN branch by default (`--gan_g_weight=0.03 --gan_d_weight=0.03`), so the data loader always uses an **ImageVideoDataset**: the adversarial critic needs **captioned real videos** to build its clean latents. A prompt-only JSON is **not** sufficient (unlike plain DMD).
+
+### 4.1 Quick Demo Dataset
+
+```bash
+# Download official demo dataset
+modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
+```
+
+### 4.2 Dataset Structure
+
+Same as Causal-Forcing Stage 1/2. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md#42-dataset-structure) for details.
+
+### 4.3 metadata.json Format
+
+Each entry needs `file_path`, `type="video"`, and `text`; the launcher uses the `metadata_add_width_height.json` variant, which additionally carries `width`/`height` so bucketing does not need to decode the media files:
+
+```json
+[
+  {
+    "file_path": "train/00000000.mp4",
+    "text": "A young woman gently turns her head to the right, tilting it slightly as she offers a soft greeting...",
+    "type": "video",
+    "width": 832,
+    "height": 480
+  }
+]
+```
+
+> **Note**: Keep enough frames per video for `--video_sample_stride=1` and `--video_sample_n_frames=81`: the GAN branch aborts when the source yields fewer latent frames than `score_num_frames=21` (see Section 5.3).
+
+---
+
+## 5. Training
+
+### 5.1 Quick Start
+
+The ready-to-use launcher is [train_one_forcing.sh](./train_one_forcing.sh). It calls the One-Forcing trainer (`train_one_forcing.py`) with the GAN recipe enabled:
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
+
+accelerate launch --mixed_precision="bf16" --use_fsdp \
+    --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+    scripts/wan2.1_one_forcing/train_one_forcing.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --ode_transformer_path=$STAGE2_CKPT \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --fix_sample_size 480 832 \
+  --video_sample_stride=1 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --max_train_steps=200 \
+  --checkpointing_steps=100 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
+  --lr_scheduler="constant" \
+  --seed=42 \
+  --output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.01 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=10.0 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --num_frame_per_block=1 \
+  --use_kv_cache_training \
+  --use_ema \
+  --ema_decay=0.99 \
+  --denoising_step_indices_list 1000 \
+  --real_guidance_scale=6.0 \
+  --score_timestep_clamp \
+  --fake_guidance_scale=0.0 \
+  --gen_update_interval=5 \
+  --gan_g_weight=0.03 \
+  --gan_d_weight=0.03 \
+  --gan_feature_layers 21 29 \
+  --gan_block_ffn_dim=2048 \
+  --gan_block_num_heads=12 \
+  --resume_from_checkpoint="latest" \
+  --trainable_modules "."
+```
+
+Or simply:
+
+```bash
+bash scripts/wan2.1_one_forcing/train_one_forcing.sh
+```
+
+### 5.2 Recipe Parameters
+
+The One-Forcing recipe sets these launcher flags:
+
+| Parameter | Launcher flag | Value |
+|------------------------|---------------|-------|
+| `denoising_step_list: [1000]` | `--denoising_step_indices_list` | `1000` (one step) |
+| `num_frame_per_block: 1` | `--num_frame_per_block` | `1` (frame-wise) |
+| `num_training_frames: 21` / `score_num_frames: 21` | `--video_sample_n_frames=81` @ 480×832 | 21 latent frames |
+| `guidance_scale: 5.0` | `--real_guidance_scale` | `6.0` (see CFG note below) |
+| `lr: 2e-6` / `lr_critic: 4e-7` (gen:critic=5:1) | `--learning_rate` / `--learning_rate_critic` | `2.0e-06` / `4.0e-07` |
+| `beta1: 0.0`, `beta2: 0.999`, `weight_decay: 0.01` | `--adam_beta1`, `--adam_beta2`, `--adam_weight_decay` | `0.0`, `0.999`, `0.01` |
+| `ema_weight: 0.99` | `--use_ema`, `--ema_decay` | on, `0.99` |
+| `max_steps: 200` | `--max_train_steps` | `200` |
+| `dfake_gen_update_ratio: 5` | `--gen_update_interval` | `5` |
+| `gan_feature_layers: [21, 29]` | `--gan_feature_layers` | `21 29` |
+| `gan_block_ffn_dim: 2048` | `--gan_block_ffn_dim` | `2048` |
+| `gan_block_num_heads: 12` | `--gan_block_num_heads` | `12` |
+| `gan_g_weight: 0.03`, `gan_d_weight: 0.03` | `--gan_g_weight`, `--gan_d_weight` | `0.03`, `0.03` |
+| `relativistic_discriminator: false` | `--relativistic_discriminator` | omitted (off) |
+| `r1_weight: 0.0`, `r2_weight: 0.0` | `--r1_weight`, `--r2_weight` | omitted (off) |
+| `context_noise: 0` | `--context_noise` | `0` (default) |
+
+> **CFG convention**: this repo computes guidance as `cond + scale * (cond - uncond)`, and this recipe uses `--real_guidance_scale=6.0`.
+
+For the general training/video-sampling parameters shared with the DMD stage, see [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md#52-key-parameters) Section 5.2.
+
+### 5.3 GAN Prerequisites
+
+The GAN branch is validated at startup in `train_one_forcing.py`. It requires:
+
+- **Captioned real videos** (`type="video"` + `file_path` + `text`), not a prompt-only JSON. Enabling any positive GAN/R1/R2 weight switches the loader from `TextDataset` to `ImageVideoDataset`.
+- **`--gradient_accumulation_steps=1`** and **no DeepSpeed** — the trainer raises `ValueError` otherwise.
+- **`--video_sample_stride=1`**: with `stride=2`, an 81-frame source yields ~10 latent frames, which is below `score_num_frames=21` and aborts.
+- **Valid feature-layer indices** for the fake-score backbone. The defaults `21 29` (layers) and `12` (heads) fit Wan2.1-T2V-1.3B; the head count must divide the fake-score hidden dimension.
+
+Trainable fake-score and head master weights stay **FP32**; the forward pass uses the configured autocast/FSDP precision, and logits/softplus/finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient-communication precision only.
+
+### 5.4 Optional Regularizers (Relativistic / R1 / R2)
+
+These regularizers are **off** by default. To experiment, append to the launcher:
+
+```bash
+# Relativistic discriminator: use paired real/fake logit differences
+--relativistic_discriminator
+
+# Finite-difference gradient penalties (APT form, not double-backward)
+--r1_weight=1.0 --r1_sigma=0.1     # on noisy real latents:   r1_weight * mean((D_real' - D_real)^2)
+--r2_weight=0.001 --r2_sigma=0.01  # on noisy gen latents: 0.5 * r2_weight * mean(((D_fake' - D_fake)/sigma)^2)
+```
+
+For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` (generator) and `softplus(-D_real) + softplus(D_fake)` (discriminator); relativistic losses use the logit differences instead. Real/fake samples share timestep and noise; perturbations use independent noise.
+
+---
+
+## 6. Inference
+
+### 6.1 Quick Start
+
+The inference entry is [examples/wan2.1_one_forcing/predict_t2v.py](../../examples/wan2.1_one_forcing/predict_t2v.py). It loads the **1.3B student** (the 14B model is only the DMD teacher at training time) and runs frame-wise one-step generation. By default it uses the **official One-Forcing checkpoint**:
+
+```bash
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+```python
+# examples/wan2.1_one_forcing/predict_t2v.py
+model_name          = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+transformer_path    = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
+use_ema             = False   # only applies to trainer-output checkpoint dirs
+stochastic_sampling = True    # DMD-distilled student
+guidance_scale      = 1.0     # CFG is baked into the distilled weights
+num_inference_steps = 1       # One-Forcing one-step generation
+num_frame_per_block = 1       # frame-wise
+sample_size         = [480, 832]
+video_length        = 81
+```
+
+Run it:
+
+```bash
+python examples/wan2.1_one_forcing/predict_t2v.py
+```
+
+`transformer_path` may be the official checkpoint file, a plain `.pt`/`.safetensors` file, or a trainer-output **checkpoint dir** (its live/EMA weights are auto-resolved). The GAN discriminator head is a training-only submodule of fake_score, so inference needs **no discriminator wiring**.
+
+### 6.2 1-Step vs 4-Step Selector
+
+With `stochastic_sampling=True`, `num_inference_steps` selects the warped timestep schedule:
+
+| `num_inference_steps` | Denoise timesteps | Use case |
+|-----------------------|-------------------|----------|
+| `1` (One-Forcing default) | `[1000]` | Fastest; frame-wise one-step generation |
+| `4` | `[1000, 750, 500, 250]` | More stable, slower |
+
+### 6.3 First-Frame Schedule
+
+The shared `WanSelfForcingPipeline` applies **one denoising schedule to all blocks**. A separate first-frame schedule (e.g. `[1000, 750, 500, 250]` for the first frame and `[1000]` for later frames) is therefore not applied; the pipeline is shared by multiple examples and is intentionally left unmodified.
+
+---
+
+## 7. Checkpoint Round-Trip
+
+The FSDP trainer writes each `checkpoint-{step}/` as follows, so training output feeds inference with no key surgery:
+
+| Component | Saved to | Contains `gan_head`? |
+|-----------|----------|----------------------|
+| **Generator** (live) | `checkpoint-{step}/diffusion_pytorch_model.safetensors` (root, via FSDP hook) | No |
+| **Generator** (EMA) | `checkpoint-{step}/transformer_ema/` | No |
+| **Fake-score + discriminator head** | `checkpoint-{step}/fake_score/` | Yes (FP32 head tensors) |
+
+- The **generator never contains `gan_head`** keys (the head lives on fake-score), so `predict_t2v.py` loads it cleanly with `strict=False`.
+- `_resolve_transformer_path` in the inference script resolves, in order: `transformer_ema/`, `ema_transformer/` (compat alias), `transformer/`, then the root safetensors.
+- **Resume** GAN training with the same GAN/head configuration: the head is recreated from CLI flags before loading and restored from `fake_score/`.
+
+---
+
+## 8. Additional Resources
+
+- **One-Forcing (source recipe)**: https://github.com/Aurora-edu/One-Forcing
+- **Causal-Forcing (DMD stage)**: https://github.com/thu-ml/Causal-Forcing
+- **Stage 1 (AR Diffusion) training guide**: [README_TRAIN_AR_DIFFUSION.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md)
+- **Stage 2 (CCD) training guide**: [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md)
+- **Underlying DMD trainer guide**: [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md)
+- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
@@ -0,0 +1,355 @@
+# Wan2.1 One-Forcing：稳定的单步自回归视频生成
+
+本文档提供在 Wan2.1-T2V-1.3B 上进行 **One-Forcing** 的完整工作流，该配方移植自 [Aurora-edu/One-Forcing](https://github.com/Aurora-edu/One-Forcing)。
+
+> **什么是 One-Forcing？**
+>
+> One-Forcing（*Towards Stable One-Step Autoregressive Video Generation*）= **Causal-Forcing Stage 3 — 分布匹配蒸馏 (DMD)** *加上*一个**共享噪声 latent 的对抗判别器（critic）**。热启动生成器由 Causal-Forcing 的前两个阶段产出，One-Forcing 再将其打磨为**稳定的单步、逐帧自回归**视频模型：
+>
+> 1. **Stage 1 — AR Diffusion**（`scripts/wan2.1_causal_forcing/train_ar_diffusion.py`）：在干净视频 latent 上以 **teacher forcing** 训练模型，学习按块（或按帧）因果自回归地去噪，产出强 AR 基座。参见 [README_TRAIN_AR_DIFFUSION_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md)。
+> 2. **Stage 2 — 因果一致性蒸馏 (CCD)**（`scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py`）：使用 EMA teacher + CFG 将多步 AR 模型蒸馏为**每块一步**的一致性模型，其产出 checkpoint 即 One-Forcing 的**热启动生成器**。参见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md)。
+> 3. **Stage 3 — One-Forcing**（`train_one_forcing.py`，即本文档），在 DMD 之上增加两个机制：
+>    - **DMD 核心**：生成器通过 14B real-score teacher 与 1.3B fake-score critic 匹配真实分布（`train_causal_dmd.py`）。
+>    - **对抗 critic**：判别器头挂在 **fake-score backbone** 上（每个选定特征层对应一个 register token）。DMD 去噪损失与对抗损失作用在**同一个噪声生成 latent** 上，critic 对真实/生成样本使用**同一份共享噪声与时间步**。
+>    - **单步**：`denoising_step_list = [1000]` —— 每帧在单个去噪步内生成（即 One-Forcing 中的 “One”）。
+>
+> 你也可以不自己训练 Stage 1/2，而是直接使用 One-Forcing 作者发布的**官方 ODE 初始化 checkpoint** 作为热启动权重（见第三节）。
+>
+> **训练器** `scripts/wan2.1_one_forcing/train_one_forcing.py` 实现了 One-Forcing 的生成器/critic GAN 损失、相对论（relativistic）选项以及有限差分 R1/R2 正则。底层 DMD 阶段请参见 [README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md)。
+
+---
+
+## 目录
+- [一、前置条件](#一前置条件)
+- [二、环境配置](#二环境配置)
+- [三、下载预训练模型](#三下载预训练模型)
+- [四、数据准备](#四数据准备)
+  - [4.1 快速测试数据集](#41-快速测试数据集)
+  - [4.2 数据集结构](#42-数据集结构)
+  - [4.3 metadata.json 格式](#43-metadatajson-格式)
+- [五、训练](#五训练)
+  - [5.1 快速开始](#51-快速开始)
+  - [5.2 配方参数](#52-配方参数)
+  - [5.3 GAN 前置要求](#53-gan-前置要求)
+  - [5.4 可选正则（相对论 / R1 / R2）](#54-可选正则相对论--r1--r2)
+- [六、推理](#六推理)
+  - [6.1 快速开始](#61-快速开始)
+  - [6.2 单步与 4 步选择器](#62-单步与-4-步选择器)
+  - [6.3 首帧调度](#63-首帧调度)
+- [七、Checkpoint 往返（训练→推理）](#七checkpoint-往返训练推理)
+- [八、更多资源](#八更多资源)
+
+---
+
+## 一、前置条件
+
+One-Forcing 是在已蒸馏生成器之上的最后打磨，因此需要：
+
+1. **热启动生成器 checkpoint** —— 一个已蒸馏/ODE 生成器，由 DMD + 对抗 critic 进一步打磨（例如 **Causal-Forcing Stage-2 CCD** checkpoint，或第三节的官方 ODE 初始化 checkpoint）。
+2. **Wan2.1-T2V-14B** 模型，作为 DMD real-score teacher。
+3. **带描述的真实视频** —— 对抗 critic 需要干净的真实 latent，因此仅有提示词的 JSON **不够用**（这一点与普通 DMD 不同）。
+
+---
+
+## 二、环境配置
+
+**方式 1：使用 requirements.txt**
+
+```bash
+pip install -r requirements.txt
+```
+
+**方式 2：手动安装依赖**
+
+```bash
+pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
+pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
+pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
+pip install yunchang xfuser modelscope openpyxl
+pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
+pip install opencv-python-headless
+pip install deepspeed==0.17.0 numpy==1.26.4
+```
+
+**方式 3：使用 Docker**
+
+```bash
+# Pull the image
+docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+
+# Enter the container
+docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+```
+
+---
+
+## 三、下载预训练模型
+
+One-Forcing 需要**两个**预训练模型：
+
+- **Wan2.1-T2V-1.3B**：生成器/判别器的基础模型。
+- **Wan2.1-T2V-14B**：非因果的 real-score teacher，用于计算真实分布得分。
+
+```bash
+# Create the model directory
+mkdir -p models/Diffusion_Transformer
+
+# Download the Wan2.1 T2V 1.3B student base model
+modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
+
+# Download the Wan2.1 T2V 14B real-score teacher
+modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B
+```
+
+One-Forcing 官方 checkpoint 托管在 HuggingFace（`JiaqiFeng/OneForcing`）：
+
+```bash
+# 训练用热启动生成器（官方 ODE 初始化 checkpoint），通过 --ode_transformer_path 加载。
+hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
+
+# 推理用已训练 One-Forcing 生成器（见第六节）。
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+如果你自己训练了 Causal-Forcing Stage 1/2，也可以将 `--ode_transformer_path` 指向你产出的 Stage-2 CCD checkpoint。
+
+---
+
+## 四、数据准备
+
+One-Forcing 默认开启 GAN 分支（`--gan_g_weight=0.03 --gan_d_weight=0.03`），因此数据加载始终使用 **ImageVideoDataset**：对抗 critic 需要**带描述的真实视频**来构建干净 latent，仅有提示词的 JSON **不够用**（这一点与普通 DMD 不同）。
+
+### 4.1 快速测试数据集
+
+```bash
+# Download the official demo dataset
+modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
+```
+
+### 4.2 数据集结构
+
+与 Causal-Forcing Stage 1/2 相同。详见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md#42-数据集结构)。
+
+### 4.3 metadata.json 格式
+
+每条记录需要 `file_path`、`type="video"` 与 `text`；启动脚本使用 `metadata_add_width_height.json` 变体，额外携带 `width`/`height`，分桶训练无需现场解码媒体文件：
+
+```json
+[
+  {
+    "file_path": "train/00000000.mp4",
+    "text": "A young woman gently turns her head to the right, tilting it slightly as she offers a soft greeting...",
+    "type": "video",
+    "width": 832,
+    "height": 480
+  }
+]
+```
+
+> **说明**：每条视频须有足够的帧数以满足 `--video_sample_stride=1` 与 `--video_sample_n_frames=81`：当源视频产出的 latent 帧数少于 `score_num_frames=21` 时，GAN 分支会中止（见 5.3）。
+
+---
+
+## 五、训练
+
+### 5.1 快速开始
+
+可直接使用的启动脚本为 [train_one_forcing.sh](./train_one_forcing.sh)。它调用 One-Forcing 训练器（`train_one_forcing.py`）并开启 GAN 配方：
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
+
+accelerate launch --mixed_precision="bf16" --use_fsdp \
+    --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+    scripts/wan2.1_one_forcing/train_one_forcing.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --ode_transformer_path=$STAGE2_CKPT \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --fix_sample_size 480 832 \
+  --video_sample_stride=1 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --max_train_steps=200 \
+  --checkpointing_steps=100 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
+  --lr_scheduler="constant" \
+  --seed=42 \
+  --output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.01 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=10.0 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --num_frame_per_block=1 \
+  --use_kv_cache_training \
+  --use_ema \
+  --ema_decay=0.99 \
+  --denoising_step_indices_list 1000 \
+  --real_guidance_scale=6.0 \
+  --score_timestep_clamp \
+  --fake_guidance_scale=0.0 \
+  --gen_update_interval=5 \
+  --gan_g_weight=0.03 \
+  --gan_d_weight=0.03 \
+  --gan_feature_layers 21 29 \
+  --gan_block_ffn_dim=2048 \
+  --gan_block_num_heads=12 \
+  --resume_from_checkpoint="latest" \
+  --trainable_modules "."
+```
+
+或直接运行：
+
+```bash
+bash scripts/wan2.1_one_forcing/train_one_forcing.sh
+```
+
+### 5.2 配方参数
+
+One-Forcing 配方设置以下 launcher 参数：
+
+| 参数 | Launcher 参数 | 取值 |
+|--------------------|---------------|------|
+| `denoising_step_list: [1000]` | `--denoising_step_indices_list` | `1000`（单步） |
+| `num_frame_per_block: 1` | `--num_frame_per_block` | `1`（逐帧） |
+| `num_training_frames: 21` / `score_num_frames: 21` | `--video_sample_n_frames=81` @ 480×832 | 21 个 latent 帧 |
+| `guidance_scale: 5.0` | `--real_guidance_scale` | `6.0`（见下方 CFG 说明） |
+| `lr: 2e-6` / `lr_critic: 4e-7`（gen:critic=5:1） | `--learning_rate` / `--learning_rate_critic` | `2.0e-06` / `4.0e-07` |
+| `beta1: 0.0`、`beta2: 0.999`、`weight_decay: 0.01` | `--adam_beta1`、`--adam_beta2`、`--adam_weight_decay` | `0.0`、`0.999`、`0.01` |
+| `ema_weight: 0.99` | `--use_ema`、`--ema_decay` | 开启、`0.99` |
+| `max_steps: 200` | `--max_train_steps` | `200` |
+| `dfake_gen_update_ratio: 5` | `--gen_update_interval` | `5` |
+| `gan_feature_layers: [21, 29]` | `--gan_feature_layers` | `21 29` |
+| `gan_block_ffn_dim: 2048` | `--gan_block_ffn_dim` | `2048` |
+| `gan_block_num_heads: 12` | `--gan_block_num_heads` | `12` |
+| `gan_g_weight: 0.03`、`gan_d_weight: 0.03` | `--gan_g_weight`、`--gan_d_weight` | `0.03`、`0.03` |
+| `relativistic_discriminator: false` | `--relativistic_discriminator` | 不传（关闭） |
+| `r1_weight: 0.0`、`r2_weight: 0.0` | `--r1_weight`、`--r2_weight` | 不传（关闭） |
+| `context_noise: 0` | `--context_noise` | `0`（默认） |
+
+> **CFG 约定**：本仓库的引导计算为 `cond + scale * (cond - uncond)`，本配方取 `--real_guidance_scale=6.0`。
+
+与 DMD 阶段共用的通用训练/视频采样参数，参见 [README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md#52-关键参数) 5.2 节。
+
+### 5.3 GAN 前置要求
+
+GAN 分支在 `train_one_forcing.py` 启动时会被校验，要求：
+
+- **带描述的真实视频**（`type="video"` + `file_path` + `text`），不能只提供提示词 JSON。任一 GAN/R1/R2 权重为正时，数据加载会从 `TextDataset` 切换为 `ImageVideoDataset`。
+- **`--gradient_accumulation_steps=1`** 且**不使用 DeepSpeed** —— 否则训练器抛出 `ValueError`。
+- **`--video_sample_stride=1`**：当 `stride=2` 时，81 帧源仅得到约 10 个 latent 帧，低于 `score_num_frames=21` 而中止。
+- fake-score backbone 上**有效的特征层索引**。默认 `21 29`（层）与 `12`（头）适用于 Wan2.1-T2V-1.3B；头数须整除 fake-score 隐藏维度。
+
+可训练的 fake-score 与判别器头 master weights 保持 **FP32**；前向按配置使用 autocast/FSDP 混合精度，logits/softplus/有限差分归约在 FP32 下计算。`--fsdp_reduce_fp32` 只改变梯度通信精度。
+
+### 5.4 可选正则（相对论 / R1 / R2）
+
+这些正则默认**关闭**。若要实验，向 launcher 追加：
+
+```bash
+# Relativistic discriminator: use paired real/fake logit differences
+--relativistic_discriminator
+
+# Finite-difference gradient penalties (APT form, not double-backward)
+--r1_weight=1.0 --r1_sigma=0.1     # on noisy real latents:   r1_weight * mean((D_real' - D_real)^2)
+--r2_weight=0.001 --r2_sigma=0.01  # on noisy gen latents: 0.5 * r2_weight * mean(((D_fake' - D_fake)/sigma)^2)
+```
+
+设 logits 为 `D_fake` 与 `D_real`，普通 GAN 损失为 `softplus(-D_fake)`（生成器）与 `softplus(-D_real) + softplus(D_fake)`（判别器）；相对论模式改用 logit 差值。真实/生成样本共用时间步与噪声，扰动噪声独立抽样。
+
+---
+
+## 六、推理
+
+### 6.1 快速开始
+
+推理入口为 [examples/wan2.1_one_forcing/predict_t2v.py](../../examples/wan2.1_one_forcing/predict_t2v.py)。它加载 **1.3B 学生模型**（14B 仅在训练时作为 DMD teacher），运行逐帧单步生成。默认使用 **One-Forcing 官方 checkpoint**：
+
+```bash
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+```python
+# examples/wan2.1_one_forcing/predict_t2v.py
+model_name          = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+transformer_path    = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
+use_ema             = False   # only applies to trainer-output checkpoint dirs
+stochastic_sampling = True    # DMD-distilled student
+guidance_scale      = 1.0     # CFG is baked into the distilled weights
+num_inference_steps = 1       # One-Forcing one-step generation
+num_frame_per_block = 1       # frame-wise
+sample_size         = [480, 832]
+video_length        = 81
+```
+
+运行：
+
+```bash
+python examples/wan2.1_one_forcing/predict_t2v.py
+```
+
+`transformer_path` 既可以是官方 checkpoint 文件、纯 `.pt`/`.safetensors` 文件，也可以是训练输出的 **checkpoint 目录**（其 live/EMA 权重会被自动解析）。GAN 判别器头是 fake_score 的训练专用子模块，因此推理**无需任何判别器接线**。
+
+### 6.2 单步与 4 步选择器
+
+当 `stochastic_sampling=True` 时，`num_inference_steps` 决定 warped 时间步调度：
+
+| `num_inference_steps` | 去噪时间步 | 适用场景 |
+|-----------------------|------------|----------|
+| `1`（One-Forcing 默认） | `[1000]` | 最快；逐帧单步生成 |
+| `4` | `[1000, 750, 500, 250]` | 更稳定、更慢 |
+
+### 6.3 首帧调度
+
+共享的 `WanSelfForcingPipeline` 对**所有块使用同一份去噪调度**，因此此处不单独为首帧设置调度（例如首帧 `[1000, 750, 500, 250]`、后续帧 `[1000]`）；该 pipeline 被多个 example 共用，有意保持不变。
+
+---
+
+## 七、Checkpoint 往返（训练→推理）
+
+FSDP 训练器将每个 `checkpoint-{step}/` 按如下布局写出，因此训练产物可直接喂给推理，无需改动 key：
+
+| 组件 | 保存位置 | 是否含 `gan_head`？ |
+|------|----------|---------------------|
+| **生成器**（live） | `checkpoint-{step}/diffusion_pytorch_model.safetensors`（根目录，经 FSDP hook） | 否 |
+| **生成器**（EMA） | `checkpoint-{step}/transformer_ema/` | 否 |
+| **Fake-score + 判别器头** | `checkpoint-{step}/fake_score/` | 是（FP32 判别器头张量） |
+
+- **生成器从不包含 `gan_head`** 键（判别器头挂在 fake-score 上），因此 `predict_t2v.py` 可用 `strict=False` 干净加载。
+- 推理脚本中的 `_resolve_transformer_path` 依次解析：`transformer_ema/`、`ema_transformer/`（兼容别名）、`transformer/`，最后到根目录 safetensors。
+- **恢复** GAN 训练时须使用相同的 GAN/判别器头配置：加载前会根据 CLI 参数重建判别器头，并从 `fake_score/` 还原。
+
+---
+
+## 八、更多资源
+
+- **One-Forcing（源配方）**：https://github.com/Aurora-edu/One-Forcing
+- **Causal-Forcing（DMD 阶段）**：https://github.com/thu-ml/Causal-Forcing
+- **Stage 1（AR Diffusion）训练指南**：[README_TRAIN_AR_DIFFUSION_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md)
+- **Stage 2（CCD）训练指南**：[README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md)
+- **底层 DMD 训练器指南**：[README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md)
+- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
@@ -0,0 +1,85 @@
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
+# One-Forcing needs REAL videos to build the adversarial critic's clean_latents, so the dataset
+# must be a video meta (type=video + file_path), NOT a prompt-only json. The official demo
+# dataset X-Fun-Videos-Demo ships captioned videos and is normalized by --enable_bucket
+# --random_hw_adapt --training_with_video_token_length.
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+# Warm-start generator: a distilled/ODE checkpoint that the DMD + adversarial critic then polish.
+# Default: the official One-Forcing ODE-initialized checkpoint, downloaded via
+#   hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
+# You can also point it at your own Causal-Forcing Stage-2 CCD checkpoint.
+export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+# One-Forcing = Causal-Forcing Stage-3 DMD + a shared noised-latent adversarial critic on fake_score.
+# - 1-step frame-wise generation: --denoising_step_indices_list 1000, --num_frame_per_block=1;
+#   21 latent frames from --video_sample_n_frames=81 @ 480x832.
+# - GAN critic: gan_g_weight/gan_d_weight=0.03, feature layers [21, 29], ffn 2048, heads 12,
+#   non-relativistic, R1/R2 off. The head rides on fake_score and is training-only.
+# - GAN requires real videos, --gradient_accumulation_steps=1, no DeepSpeed, and --video_sample_stride=1
+#   (stride=2 on 81-frame sources yields ~10 latent frames < score_num_frames=21 and aborts).
+# - guidance: local CFG is cond + scale*(cond - uncond), so --real_guidance_scale=6.0.
+# - optimizer/EMA/steps: lr=lr_critic=1e-5, betas (0.0, 0.999), weight_decay=0.01,
+#   --use_ema --ema_decay=0.99, gen_update_interval=5, max_steps=200.
+accelerate launch --mixed_precision="bf16" --use_fsdp \
+    --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+    scripts/wan2.1_one_forcing/train_one_forcing.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --ode_transformer_path=$STAGE2_CKPT \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --fix_sample_size 480 832 \
+  --video_sample_stride=1 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --max_train_steps=200 \
+  --checkpointing_steps=100 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
+  --lr_scheduler="constant" \
+  --seed=42 \
+  --output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.01 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=10.0 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --num_frame_per_block=1 \
+  --use_kv_cache_training \
+  --use_ema \
+  --ema_decay=0.99 \
+  --denoising_step_indices_list 1000 \
+  --real_guidance_scale=6.0 \
+  --score_timestep_clamp \
+  --fake_guidance_scale=0.0 \
+  --gen_update_interval=5 \
+  --gan_g_weight=0.03 \
+  --gan_d_weight=0.03 \
+  --gan_feature_layers 21 29 \
+  --gan_block_ffn_dim=2048 \
+  --gan_block_num_heads=12 \
+  --resume_from_checkpoint="latest" \
+  --trainable_modules "."
@@ -179,6 +179,7 @@ The difference between DeepSpeed-Zero-2 and FSDP lies in whether to shard model

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -189,6 +190,7 @@ NCCL_DEBUG=INFO
 accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -204,21 +206,26 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -247,17 +254,20 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | `--num_train_epochs` | Number of training epochs | 100 |
 | `--checkpointing_steps` | Save checkpoint every N steps | 50 |
 | `--learning_rate` | Initial learning rate (generator) | 2e-06 |
-| `--learning_rate_critic` | Initial learning rate (critic) | 2e-06 |
+| `--learning_rate_critic` | Initial learning rate (critic) | 4.0e-07 |
 | `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` |
 | `--lr_warmup_steps` | Learning rate warmup steps | 100 |
 | `--seed` | Random seed | 42 |
 | `--output_dir` | Output directory | `output_dir_wan2.1_self_forcing_distill` |
 | `--gradient_checkpointing` | Enable gradient checkpointing | - |
 | `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
-| `--adam_weight_decay` | AdamW weight decay | 3e-2 |
+| `--adam_weight_decay` | AdamW weight decay | 0.0 |
+| `--adam_beta1` | AdamW beta1 | 0.0 |
+| `--adam_beta2` | AdamW beta2 | 0.999 |
 | `--adam_epsilon` | AdamW epsilon | 1e-10 |
 | `--vae_mini_batch` | VAE encoding mini-batch size | 1 |
-| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
+| `--max_grad_norm` | Gradient clipping threshold | 10.0 |
+| `--fsdp_reduce_fp32` | Optional: reduce gradients in fp32 under FSDP mixed precision (default reduces in bf16) | - |
 | `--enable_bucket` | Enable bucket training, no cropping, group by resolution | - |
 | `--random_hw_adapt` | Auto-scale images/videos to random sizes in `[min_size, max_size]` range | - |
 | `--training_with_video_token_length` | Train based on token length, supports arbitrary resolutions | - |
@@ -275,12 +285,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` |
-| `--real_guidance_scale` | Real guidance scale for scoring | 6.0 |
+| `--randomize_step_indices` | Randomize the denoising timestep indices during training | - |
+| `--real_guidance_scale` | Real guidance scale for scoring | 4.0 |
+| `--score_timestep_clamp` | Optional: clamp DMD/critic timesteps to `[0.02, 0.98] * num_train_timesteps` | - |
 | `--fake_guidance_scale` | Fake guidance scale for scoring | 0.0 |
 | `--gen_update_interval` | Generator update interval | 5 |
 | `--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 +309,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.
@@ -342,6 +374,7 @@ You can configure validation parameters to periodically generate test videos dur

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -352,6 +385,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -367,21 +401,26 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -404,6 +443,7 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/
 Training shell command is as follows:
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -414,6 +454,7 @@ NCCL_DEBUG=INFO
 accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -429,21 +470,26 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -458,6 +504,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -468,6 +515,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -483,21 +531,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -517,6 +570,7 @@ Assuming 2 machines, each with 8 GPUs:
 **Machine 0 (Master)**:
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # Master machine IP
@@ -532,6 +586,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -547,21 +602,26 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -573,6 +633,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
 **Machine 1 (Worker)**:
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # Same as Master
@@ -180,6 +180,7 @@ hf download gdhe17/Self-Forcing --local-dir models/Diffusion_Transformer/Self-Fo

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -190,6 +191,7 @@ NCCL_DEBUG=INFO
 accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -205,21 +207,26 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -248,17 +255,20 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | `--num_train_epochs` | 训练 epoch 数 | 100 |
 | `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 |
 | `--learning_rate` | 初始学习率（生成器） | 2e-06 |
-| `--learning_rate_critic` | 初始学习率（判别器） | 2e-06 |
+| `--learning_rate_critic` | 初始学习率（判别器） | 4.0e-07 |
 | `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
 | `--lr_warmup_steps` | 学习率预热步数 | 100 |
 | `--seed` | 随机种子 | 42 |
 | `--output_dir` | 输出目录 | `output_dir_wan2.1_self_forcing_distill` |
 | `--gradient_checkpointing` | 激活重计算 | - |
 | `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
-| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
+| `--adam_weight_decay` | AdamW 权重衰减 | 0.0 |
+| `--adam_beta1` | AdamW beta1 | 0.0 |
+| `--adam_beta2` | AdamW beta2 | 0.999 |
 | `--adam_epsilon` | AdamW epsilon 值 | 1e-10 |
 | `--vae_mini_batch` | VAE 编码时的迷你批次大小 | 1 |
-| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
+| `--max_grad_norm` | 梯度裁剪阈值 | 10.0 |
+| `--fsdp_reduce_fp32` | 可选：FSDP 混精下梯度以 fp32 做跨卡 reduce（默认 bf16） | - |
 | `--enable_bucket` | 启用分桶训练，不裁剪图片/视频，按分辨率分组训练 | - |
 | `--random_hw_adapt` | 自动缩放图片/视频到 `[min_size, max_size]` 范围内的随机尺寸 | - |
 | `--training_with_video_token_length` | 根据 token 长度训练，支持任意分辨率 | - |
@@ -276,12 +286,17 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | 参数 | 说明 | 示例值 |
 |-----|------|-------|
 | `--denoising_step_indices_list` | 去噪步骤列表（蒸馏核心参数） | `1000 750 500 250` |
-| `--real_guidance_scale` | 用于评分的真实 guidance scale | 6.0 |
+| `--randomize_step_indices` | 训练时随机化去噪 timestep 索引 | - |
+| `--real_guidance_scale` | 用于评分的真实 guidance scale | 4.0 |
+| `--score_timestep_clamp` | 可选：将 DMD/critic 的 timestep 截断到 `[0.02, 0.98] * num_train_timesteps` | - |
 | `--fake_guidance_scale` | 用于评分的虚拟 guidance scale | 0.0 |
 | `--gen_update_interval` | 生成器更新间隔 | 5 |
 | `--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 +310,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 时，表示视频和图像分辨率的最大值。
@@ -343,6 +375,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -353,6 +386,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -368,21 +402,26 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -405,6 +444,7 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/
 训练 shell 命令如下：
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -415,6 +455,7 @@ NCCL_DEBUG=INFO
 accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -430,21 +471,26 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -459,6 +505,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -469,6 +516,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -484,21 +532,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -518,6 +571,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
 **机器 0（Master）**：
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # Master 机器 IP
@@ -533,6 +587,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -548,21 +603,26 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -574,6 +634,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
 **机器 1（Worker）**：
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # 与 Master 相同
@@ -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
@@ -1,4 +1,5 @@
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/internal_datasets/"
 export DATASET_META_NAME="datasets/internal_datasets/metadata.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -9,6 +10,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -24,21 +26,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -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>
@@ -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>
@@ -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
@@ -19,4 +19,4 @@ from .utils import (VIDEO_READER_TIMEOUT, Camera, VideoReader_contextmanager,
                    prepare_lingbot_dit_cond_dict,
                    prepare_lingbot_dit_cond_dict_from_c2ws, process_pose_file,
                    process_pose_params, ray_condition, resize_frame,
-                    resize_image_with_target_area)
+                    resize_image_with_target_area, sample_clip_start)
@@ -30,7 +30,7 @@ except ImportError:
 from .utils import (VIDEO_READER_TIMEOUT, VideoReader_contextmanager,
                    get_random_mask, get_video_reader_batch, padding_image,
                    process_pose_file, process_pose_params, resize_frame,
-                    resize_image_with_target_area)
+                    resize_image_with_target_area, sample_clip_start)


 class ImageVideoSampler(BatchSampler):
@@ -99,6 +99,7 @@ class ImageVideoDataset(Dataset):
        enable_inpaint=False,
        inpaint_mask_fill_value=0,
        return_file_name=False,
+        random_start=True,
    ):
        # Loading annotations from files
        print(f"loading annotations from {ann_path} ...")
@@ -136,6 +137,8 @@ class ImageVideoDataset(Dataset):

        self.video_length_drop_start = video_length_drop_start
        self.video_length_drop_end = video_length_drop_end
+        # Whether the clip start is drawn randomly (training) or pinned (validation/eval); see `sample_clip_start`.
+        self.random_start = random_start

        # Video params: resize, center crop, normalize to [-1, 1]
        self.video_sample_stride    = video_sample_stride
@@ -188,10 +191,13 @@ class ImageVideoDataset(Dataset):
                        f"Video too short: sampled {min_sample_n_frames} frames < required {min_video_sample_n_frames}."
                    )

-                # Select contiguous clip with random start position
+                # Select contiguous clip; the start is random unless this is an eval dataset
                video_length = int(self.video_length_drop_end * len(video_reader))
                clip_length = min(video_length, (min_sample_n_frames - 1) * self.video_sample_stride + 1)
-                start_idx   = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
+                start_idx = sample_clip_start(
+                    video_length, clip_length,
+                    drop_start=self.video_length_drop_start, random_start=self.random_start,
+                )
                batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)

                try:
@@ -363,7 +369,11 @@ class LingbotImageVideoDataset(ImageVideoDataset):

                video_length = int(self.video_length_drop_end * len(video_reader))
                clip_length = min(video_length, (min_sample_n_frames - 1) * self.video_sample_stride + 1)
-                start_idx = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
+                # The start is random unless this is an eval dataset
+                start_idx = sample_clip_start(
+                    video_length, clip_length,
+                    drop_start=self.video_length_drop_start, random_start=self.random_start,
+                )
                batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)

                try:
@@ -503,6 +513,7 @@ class ImageVideoControlDataset(Dataset):
        enable_subject_info=False,
        padding_subject_info=True,
        return_file_name=False,
+        random_start=True,
    ):
        # Loading annotations from files
        print(f"loading annotations from {ann_path} ...")
@@ -543,6 +554,8 @@ class ImageVideoControlDataset(Dataset):

        self.video_length_drop_start = video_length_drop_start
        self.video_length_drop_end = video_length_drop_end
+        # Whether the clip start is drawn randomly (training) or pinned (validation/eval); see `sample_clip_start`.
+        self.random_start = random_start

        # Video params: resize, center crop, normalize to [-1, 1]
        self.video_sample_stride    = video_sample_stride
@@ -603,10 +616,13 @@ class ImageVideoControlDataset(Dataset):
                        f"Video too short: sampled {min_sample_n_frames} frames < required {min_video_sample_n_frames}."
                    )

-                # Select contiguous clip with random start position
+                # Select contiguous clip; the start is random unless this is an eval dataset
                video_length = int(self.video_length_drop_end * len(video_reader))
                clip_length = min(video_length, (min_sample_n_frames - 1) * self.video_sample_stride + 1)
-                start_idx   = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
+                start_idx = sample_clip_start(
+                    video_length, clip_length,
+                    drop_start=self.video_length_drop_start, random_start=self.random_start,
+                )
                batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)

                try:
@@ -34,7 +34,8 @@ except ImportError:
    from .utils import AVVideoReader as VideoReader

 from .utils import (VIDEO_READER_TIMEOUT, VideoReader_contextmanager,
-                    get_random_mask, get_video_reader_batch, resize_frame)
+                    get_random_mask, get_video_reader_batch, resize_frame,
+                    sample_clip_start)


 def load_audio(path, sr, mono=True, native_sr=False, res_type=None):
@@ -87,6 +88,7 @@ class WebVid10M(Dataset):
        enable_bucket=False, 
        enable_inpaint=False, 
        is_image=False,
+        random_start=True,
    ):
        print(f"loading annotations from {csv_path} ...")
        with open(csv_path, 'r') as csvfile:
@@ -100,6 +102,7 @@ class WebVid10M(Dataset):
        self.enable_bucket   = enable_bucket
        self.enable_inpaint  = enable_inpaint
        self.is_image        = is_image
+        self.random_start    = random_start
        
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size)
        self.pixel_transforms = transforms.Compose([
@@ -118,10 +121,10 @@ class WebVid10M(Dataset):
        
        if not self.is_image:
            clip_length = min(video_length, (self.sample_n_frames - 1) * self.sample_stride + 1)
-            start_idx   = random.randint(0, video_length - clip_length)
+            start_idx   = sample_clip_start(video_length, clip_length, random_start=self.random_start)
            batch_index = np.linspace(start_idx, start_idx + clip_length - 1, self.sample_n_frames, dtype=int)
        else:
-            batch_index = [random.randint(0, video_length - 1)]
+            batch_index = [sample_clip_start(video_length, 1, random_start=self.random_start)]

        if not self.enable_bucket:
            pixel_values = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous()
@@ -173,6 +176,7 @@ class VideoDataset(Dataset):
        video_length_drop_start=0.0,
        video_length_drop_end=1.0,
        text_drop_ratio=0.1,
+        random_start=True,
    ):
        # Loading annotations from files
        print(f"loading annotations from {ann_path} ...")
@@ -189,6 +193,7 @@ class VideoDataset(Dataset):
        self.video_length_drop_start = video_length_drop_start
        self.video_length_drop_end = video_length_drop_end
        self.text_drop_ratio = text_drop_ratio
+        self.random_start = random_start
        
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size)
        self.pixel_transforms = transforms.Compose(
@@ -212,17 +217,23 @@ class VideoDataset(Dataset):

        with VideoReader_contextmanager(video_dir, num_threads=2) as video_reader:
            # Calculate frame sampling range with length dropout
+            sample_stride = (random.choice(self.sample_stride)
+                             if isinstance(self.sample_stride, (tuple, list)) else self.sample_stride)
            min_sample_n_frames = min(
                self.sample_n_frames, 
-                int(len(video_reader) * (self.video_length_drop_end - self.video_length_drop_start) // self.sample_stride)
+                int((len(video_reader) * (self.video_length_drop_end - self.video_length_drop_start) - 1)
+                    // sample_stride) + 1
            )
            if min_sample_n_frames == 0:
                raise ValueError(f"No Frames in video.")

-            # Select contiguous clip with random start position
+            # Select contiguous clip; the start is random unless this is an eval dataset
            video_length = int(self.video_length_drop_end * len(video_reader))
-            clip_length = min(video_length, (min_sample_n_frames - 1) * self.sample_stride + 1)
-            start_idx   = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
+            clip_length = min(video_length, (min_sample_n_frames - 1) * sample_stride + 1)
+            start_idx = sample_clip_start(
+                video_length, clip_length,
+                drop_start=self.video_length_drop_start, random_start=self.random_start,
+            )
            batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)

            try:
@@ -303,6 +314,7 @@ class VideoSpeechDataset(Dataset):
        audio_stereo=False,
        audio_span_includes_last_frame=False,
        enable_ref2va=False,
+        random_start=True,
    ):
        # Loading annotations from files
        print(f"loading annotations from {ann_path} ...")
@@ -323,6 +335,8 @@ class VideoSpeechDataset(Dataset):
        # Video params: resize, center crop, normalize to [-1, 1]
        self.video_sample_stride = video_sample_stride
        self.video_sample_n_frames = video_sample_n_frames
+        # Whether the clip start is drawn randomly (training) or pinned (validation/eval); see `sample_clip_start`.
+        self.random_start = random_start
        # Fewest sampled frames a clip has to yield to be usable. A model whose VAE only encodes certain frame
        # counts (MiniMax-H3 needs `17n + 5`, so at least 5) raises this so that clips which cannot fill one chunk
        # are skipped by `__getitem__`'s retry instead of reaching the collate as a short batch. The default of 1 is
@@ -437,10 +451,11 @@ class VideoSpeechDataset(Dataset):
                    "run needs; skipping."
                )

-            # Randomly select the starting frame
+            # Select the starting frame; pinned when random_start=False so eval clips are reproducible
            frame_span = (actual_n_frames - 1) * frame_step
-            max_start = total_frames - 1 - int(math.ceil(frame_span))
-            start_frame = random.randint(0, max_start) if max_start > 0 else 0
+            start_frame = sample_clip_start(
+                total_frames, int(math.ceil(frame_span)) + 1, random_start=self.random_start
+            )
            frame_indices = [
                min(total_frames - 1, int(round(start_frame + index * frame_step)))
                for index in range(actual_n_frames)
@@ -745,6 +760,7 @@ class VideoSpeechControlDataset(Dataset):
        audio_native_sr_resample=False,
        audio_stereo=False,
        audio_span_includes_last_frame=False,
+        random_start=True,
    ):
        # Loading annotations from files
        print(f"loading annotations from {ann_path} ...")
@@ -765,6 +781,8 @@ class VideoSpeechControlDataset(Dataset):
        # Video params: resize, center crop, normalize to [-1, 1]
        self.video_sample_stride = video_sample_stride
        self.video_sample_n_frames = video_sample_n_frames
+        # Whether the clip start is drawn randomly (training) or pinned (validation/eval); see `sample_clip_start`.
+        self.random_start = random_start
        # Fewest sampled frames a clip has to yield to be usable; see `VideoSpeechDataset` for the rationale.
        self.min_video_sample_n_frames = max(1, int(min_video_sample_n_frames))
        # Frame rate the sampled clip has to land on; see `VideoSpeechDataset` for the rationale.
@@ -871,10 +889,11 @@ class VideoSpeechControlDataset(Dataset):
                    "run needs; skipping."
                )

-            # Randomly select the starting frame
+            # Select the starting frame; pinned when random_start=False so eval clips are reproducible
            frame_span = (actual_n_frames - 1) * frame_step
-            max_start = total_frames - 1 - int(math.ceil(frame_span))
-            start_frame = random.randint(0, max_start) if max_start > 0 else 0
+            start_frame = sample_clip_start(
+                total_frames, int(math.ceil(frame_span)) + 1, random_start=self.random_start
+            )
            frame_indices = [
                min(total_frames - 1, int(round(start_frame + index * frame_step)))
                for index in range(actual_n_frames)
@@ -1140,6 +1159,7 @@ class VideoAnimateDataset(Dataset):
        video_length_drop_start=0.1, 
        video_length_drop_end=0.9,
        return_file_name=False,
+        random_start=True,
    ):
        # Loading annotations from files
        print(f"loading annotations from {ann_path} ...")
@@ -1175,6 +1195,8 @@ class VideoAnimateDataset(Dataset):

        self.video_length_drop_start = video_length_drop_start
        self.video_length_drop_end = video_length_drop_end
+        # Whether the clip start is drawn randomly (training) or pinned (validation/eval); see `sample_clip_start`.
+        self.random_start = random_start

        # Video params: resize, center crop, normalize to [-1, 1]
        self.video_sample_stride = video_sample_stride
@@ -1210,10 +1232,13 @@ class VideoAnimateDataset(Dataset):
            if min_sample_n_frames == 0:
                raise ValueError(f"No Frames in video.")

-            # Select contiguous clip with random start position
+            # Select contiguous clip; the start is random unless this is an eval dataset
            video_length = int(self.video_length_drop_end * len(video_reader))
            clip_length = min(video_length, (min_sample_n_frames - 1) * self.video_sample_stride + 1)
-            start_idx = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
+            start_idx = sample_clip_start(
+                video_length, clip_length,
+                drop_start=self.video_length_drop_start, random_start=self.random_start,
+            )
            batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)

            try:
@@ -82,6 +82,28 @@ class _AVBatchFrames:
    def asnumpy(self):
        return np.stack(self._frames)

+
+def sample_clip_start(video_length, clip_length, drop_start=0.0, random_start=True):
+    """Pick the first frame index of a contiguous clip.
+
+    `video_length` is the usable tail boundary (already scaled by `video_length_drop_end`) and
+    `clip_length` the span the clip covers, so `video_length - clip_length` is the last valid start.
+    With `random_start` the start is drawn uniformly over every valid position, which is the temporal
+    augmentation training sets want; with `random_start=False` it is pinned to the
+    `video_length_drop_start` boundary, so validation/eval datasets rebuild the very same clip on
+    every pass and their metrics stay comparable across steps. The head boundary is clamped to the
+    last valid start, hence a clip that already fills the usable range can never run past the video.
+
+    Note this only decides *where* a clip starts; how many frames it yields stays the caller's
+    business, since the dataset classes differ in that formula.
+    """
+    last_valid_start = max(video_length - clip_length, 0)
+    start_idx = min(int(drop_start * video_length), last_valid_start)
+    if random_start and last_valid_start > start_idx:
+        start_idx = random.randint(start_idx, last_valid_start)
+    return start_idx
+
+
 def get_random_mask(shape, image_start_only=False):
    f, c, h, w = shape
    mask = torch.zeros((f, 1, h, w), dtype=torch.uint8)
@@ -21,9 +21,13 @@ from .longcatvideo_xfuser import (usp_attn_longcatvideo_avatar_forward,
                                  usp_rope_longcatvideo_forward)
 from .ltx2_xfuser import (LTX2MultiGPUsAttnProcessor,
                          LTX2PerturbedMultiGPUsAttnProcessor)
-from .minimax_h3_xfuser import MiniMaxH3MultiGPUsAttnProcessor
+from .minimax_h3_xfuser import (MiniMaxH3MultiGPUsAttnProcessor,
+                                MiniMaxH3MultiGPUsVSAAttnProcessor)
 from .qwen_xfuser import QwenImageMultiGPUsAttnProcessor2_0
 from .qwenimage21_xfuser import QwenImage21MultiGPUsAttnProcessor
-from .wan_xfuser import (usp_attn_forward, usp_attn_s2v_forward,
-                         usp_attn_self_forcing_forward)
+from .wan_xfuser import (usp_attn_forward, usp_attn_fastvideo_forward,
+                         usp_attn_s2v_forward,
+                         usp_attn_self_forcing_forward,
+                         usp_attn_sparkwan_forward,
+                         usp_attn_turbowan_forward)
 from .z_image_xfuser import ZMultiGPUsSingleStreamAttnProcessor
@@ -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


@@ -41,7 +41,7 @@ def set_multi_gpus_devices(ulysses_degree, ring_degree, classifier_free_guidance
        if get_sp_group is None:
            raise RuntimeError("xfuser is not installed.")
        dist.init_process_group("nccl")
-        print('parallel inference enabled: ulysses_degree=%d ring_degree=%d classifier_free_guidance_degree=% rank=%d world_size=%d' % (
+        print('parallel inference enabled: ulysses_degree=%d ring_degree=%d classifier_free_guidance_degree=%d rank=%d world_size=%d' % (
            ulysses_degree, ring_degree, classifier_free_guidance_degree, dist.get_rank(),
            dist.get_world_size()))
        assert dist.get_world_size() == ring_degree * ulysses_degree * classifier_free_guidance_degree, \
@@ -53,6 +53,14 @@ def set_multi_gpus_devices(ulysses_degree, ring_degree, classifier_free_guidance
                ulysses_degree=ulysses_degree)
        # device = torch.device("cuda:%d" % dist.get_rank())
        device = torch.device(f"cuda:{get_world_group().local_rank}")
+        # Bind this process to its own GPU. ``init_process_group`` above runs before xfuser's
+        # ``init_distributed_environment``, whose ``set_device`` is guarded by
+        # ``if not torch.distributed.is_initialized()`` and so never fires; without this call the
+        # current device stays 0 on every rank while the model lives on ``cuda:{local_rank}``.
+        # PyTorch ops dispatch on ``tensor.device`` and are unaffected, but a Triton kernel launches
+        # on ``torch.cuda.current_device()`` and rejects the rank's own pointers as inaccessible,
+        # so any Triton sparse attention (e.g. VSA) would fail on rank > 0.
+        torch.cuda.set_device(device)
        print('rank=%d device=%s' % (get_world_group().rank, str(device)))
    else:
        device = "cuda"
@@ -1,5 +1,7 @@
 import torch
+import torch.nn.functional as F

+from .fuser import get_sp_group, ulysses_all_to_all
 from .fuser import xFuserLongContextAttention


@@ -80,3 +82,99 @@ class MiniMaxH3MultiGPUsAttnProcessor:
        hidden_states = attn.to_out[0](hidden_states)
        hidden_states = attn.to_out[1](hidden_states)
        return hidden_states
+
+
+class MiniMaxH3MultiGPUsVSAAttnProcessor:
+    """Sequence-parallel attention processor for the FastVideo VSA student on MiniMax-H3.
+
+    Mirrors ``MiniMaxH3MultiGPUsAttnProcessor``'s row-split -> head-split Ulysses scheme, but swaps
+    the dense ``xFuserLongContextAttention`` core for the VSA block-sparse op (``attn.attn_op``) so
+    the trained coarse-to-fine selection and the ``to_gate_compress`` compensation branch are kept
+    instead of being silently dropped by a dense fallback.
+
+    Two VSA-specific points under head-parallelism:
+      * ``to_gate_compress`` is a ``Linear(dim, heads*head_dim)`` that MIXES heads, so it cannot run
+        in the head-shard layout. It is strictly per-token and SP shards rows while keeping the full
+        hidden dim, so the gate is evaluated here on this rank's ROW shard (yielding all heads), then
+        all-to-all'd alongside q/k/v -- element-wise identical to computing it on the full sequence
+        and slicing heads, and it lands on the same contiguous head block the all-to-all hands back.
+      * the parent forward pads the packed sequence up to a multiple of the group size, so the
+        gathered sequence carries ``<= sp_world_size - 1`` zero rows at its tail. VSA builds its tile
+        geometry from ``vsa_segments``, whose lengths sum to the PADLESS ``valid_length``; the padding
+        tail is therefore sliced off before the op and the output padded back afterwards. Unlike the
+        dense processor (which lets the pad rows dilute the softmax by ~1e-5), this keeps VSA
+        bitwise-equal to the single-GPU run -- the pad rows never enter the kernel.
+
+    Requires a padless packed sequence (``attention_mask is None``) and a pure-Ulysses split
+    (``num_heads`` divisible by the group size, checked by the model's ``enable_multi_gpus_inference``).
+    A dense student (``topk >= 1.0``) carries no ``attn_op`` and must keep the dense processor.
+    """
+
+    def __call__(
+        self,
+        attn,
+        hidden_states: torch.Tensor,
+        rotary_emb=None,
+        attention_mask=None,
+        valid_length=None,
+        vsa_segments=None,
+    ) -> torch.Tensor:
+        from ..models.minimax_h3_transformer3d import \
+            apply_minimax_h3_rotary_emb
+
+        if attention_mask is not None:
+            raise ValueError(
+                "MiniMaxH3MultiGPUsVSAAttnProcessor does not support a masked (padded) packed sequence.")
+        if not hasattr(attn, "attn_op"):
+            raise RuntimeError(
+                "The VSA sequence-parallel processor needs attn.attn_op; a dense student (topk>=1.0) "
+                "must bind MiniMaxH3MultiGPUsAttnProcessor instead.")
+
+        query = attn.to_q(hidden_states).unflatten(-1, (attn.heads, -1))
+        key = attn.to_k(hidden_states).unflatten(-1, (attn.heads, -1))
+        value = attn.to_v(hidden_states).unflatten(-1, (attn.heads, -1))
+
+        query = attn.norm_q(query)
+        key = attn.norm_k(key)
+
+        # The rotary coordinates are this rank's rows, so apply them before the all-to-all moves rows.
+        if rotary_emb is not None:
+            query = apply_minimax_h3_rotary_emb(query, *rotary_emb)
+            key = apply_minimax_h3_rotary_emb(key, *rotary_emb)
+
+        # Cross-head compress gate on this rank's ROW shard (full hidden dim present): [B, chunk, H, D].
+        gate = attn.attn_op.to_gate_compress(hidden_states).unflatten(-1, (attn.heads, -1))
+
+        half_dtypes = (torch.float16, torch.bfloat16)
+
+        def half(x):
+            return x if x.dtype in half_dtypes else x.to(torch.bfloat16)
+
+        # row-split -> head-split over the whole (padded) packed sequence, gate alongside q/k/v.
+        query = ulysses_all_to_all(half(query), scatter_dim=2, gather_dim=1)
+        key = ulysses_all_to_all(half(key), scatter_dim=2, gather_dim=1)
+        value = ulysses_all_to_all(half(value), scatter_dim=2, gather_dim=1)
+        gate = ulysses_all_to_all(half(gate), scatter_dim=2, gather_dim=1)
+
+        # Drop the SP padding tail so VSA sees exactly the padless length its segments describe.
+        padded_length = query.shape[1]
+        trim = valid_length is not None and valid_length < padded_length
+        if trim:
+            query = query[:, :valid_length]
+            key = key[:, :valid_length]
+            value = value[:, :valid_length]
+            gate = gate[:, :valid_length]
+
+        out = attn.attn_op(query, key, value, None, segments=vsa_segments, gate=gate)
+
+        # Pad the output back to the gathered length before the inverse all-to-all; the zero rows are
+        # dropped by the model's final all_gather and never touch a real row's attention.
+        if trim:
+            out = F.pad(out, (0, 0, 0, 0, 0, padded_length - valid_length))
+
+        # head-split -> row-split, back to this rank's rows.
+        out = ulysses_all_to_all(out.to(hidden_states.dtype), scatter_dim=1, gather_dim=2)
+        out = out.flatten(2, 3).type_as(query)
+        out = attn.to_out[0](out)
+        out = attn.to_out[1](out)
+        return out
@@ -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,248 @@ 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)
+        # heads_total/head_start keep the SLA linear branch's fp32 normalizer sum bitwise-equal to
+        # the single-GPU result despite this rank holding only the H/P head block.
+        h_per = n // sp_size
+        x = self.attn_op(q, k, v, heads_total=n, head_start=sp_rank * h_per)
+        x = ulysses_all_to_all(x, scatter_dim=1, gather_dim=2)
+        x = x.flatten(2)
+        x = self.o(x)
+        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
+
+def usp_attn_turbowan_forward(self,
+                     x,
+                     seq_lens,
+                     grid_sizes,
+                     freqs,
+                     dtype=torch.bfloat16,
+                     t=0):
+    """Sequence-parallel self-attention for the TurboDiffusion SLA student.
+
+    ``TurboWanTransformer3DModel`` swaps ``WanSelfAttention`` for a local variant whose kernel
+    is ``self.attn_op`` (``FakeLocalAttention`` -> ``SparseLinearAttention``, block-sparse exact
+    softmax + a ``proj_l`` feature-map linear branch). The stock ``usp_attn_forward`` binds
+    ``xFuserLongContextAttention``'s dense full-softmax and only touches the base-class
+    ``q/k/v/norm_q/norm_k/o`` attributes, so patching it onto this model would silently drop the
+    SLA kernel -- every needed attribute is inherited, hence no AttributeError, no shape error,
+    just dense numbers. This variant keeps ``self.attn_op`` live under Ulysses head-parallelism.
+
+    Equivalence: ``SparseLinearAttention.forward`` takes and returns ``[B, L, H, D]``, transposes
+    to BHLD internally and does every reduction over L/D only (``get_block_map`` pooling and
+    top-k, the sparse kernel, the ``k^T v`` linear branch), and ``proj_l`` is a per-head
+    ``head_dim -> head_dim`` projection -- no cross-head mixing and no head-indexed parameter,
+    so an all-to-all from the sequence shard ``[B, S/P, H, D]`` to the head shard
+    ``[B, S, H/P, D]`` reproduces the single-GPU result exactly (top-k is a per-query key-block
+    ratio). RoPE is applied BEFORE the all-to-all because ``rope_apply`` slices the frequency
+    table by sequence-parallel rank.
+
+    Requires a pure-Ulysses split (``ring_degree`` must be 1: the SLA kernel exposes no partial
+    output / logsumexp interface, so ring's online rescale-concat cannot be wired to it) and
+    ``num_heads`` divisible by the SP degree.
+    """
+    # No sparse kernel on this instance: keep the stock dense SP path.
+    if not hasattr(self, "attn_op"):
+        return usp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=dtype, t=t)
+
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    half_dtypes = (torch.float16, torch.bfloat16)
+
+    def half(t_):
+        return t_ if t_.dtype in half_dtypes else t_.to(dtype)
+
+    # query, key, value function
+    def qkv_fn(x):
+        q = self.norm_q(self.q(x.to(dtype))).view(b, s, n, d)
+        k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d)
+        v = self.v(x.to(dtype)).view(b, s, n, d)
+        return q, k, v
+
+    q, k, v = qkv_fn(x)
+    q, k = rope_apply_qk(q, k, grid_sizes, freqs)
+
+    # seq-shard -> head-shard, SLA over the FULL sequence on a subset of heads, then back.
+    q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
+    k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
+    v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
+    # heads_total/head_start keep the SLA linear branch's fp32 normalizer sum bitwise-equal to the
+    # single-GPU result despite this rank holding only the H/P head block.
+    sp_size = get_sequence_parallel_world_size()
+    sp_rank = get_sequence_parallel_rank()
+    h_per = n // sp_size
+    x = self.attn_op(q, k, v, heads_total=n, head_start=sp_rank * h_per)
+    x = ulysses_all_to_all(x.to(dtype), scatter_dim=1, gather_dim=2)
+
+    # output
+    x = x.flatten(2)
+    x = self.o(x)
+    return x
+
+def usp_attn_fastvideo_forward(self,
+                     x,
+                     seq_lens,
+                     grid_sizes,
+                     freqs,
+                     dtype=torch.bfloat16,
+                     t=0):
+    """Sequence-parallel self-attention for the FastVideo VSA student.
+
+    Same Ulysses head-parallel scheme as ``usp_attn_sparkwan_forward``, with one VSA-specific
+    wrinkle: ``VSASparseAttention.to_gate_compress`` is a ``Linear(dim, num_heads*head_dim)``
+    that MIXES heads, so it cannot run inside the head-shard layout the way the rest of the VSA
+    kernel can (``sparse`` / ``compressed`` only reduce over L/D and are head-independent).
+    It is, however, strictly per-token -- and sequence parallelism shards the sequence while
+    leaving the hidden dim intact, so every rank still holds the full ``C`` for its own tokens.
+    The gate is therefore evaluated here on the SEQUENCE shard (yielding all ``H`` heads), then
+    all-to-all'd alongside q/k/v, which is element-wise identical to computing it on the full
+    sequence and slicing heads -- and lands on exactly the contiguous head block
+    ``[sp_rank*H/P, (sp_rank+1)*H/P)`` the all-to-all hands this rank.
+
+    The gathered sequence is ``s * sp_size``, which the parent forward pads to a multiple of the
+    SP degree, so it can exceed ``grid_sizes.prod()``; ``_vsa_dispatch`` routes that case to its
+    per-sample branch (attend over the valid grid tokens, pad the output back), keeping the tile
+    geometry built from the true grid.
+
+    Requires a pure-Ulysses split (``ring_degree`` must be 1) and ``num_heads`` divisible by the
+    SP degree. ``topk >= 1.0`` (the dense fallback) defers to ``usp_attn_forward``.
+    """
+    # Dense student (topk >= 1.0) or a plain Wan self-attention: keep the stock SP path.
+    if getattr(self, "topk", 1.0) >= 1.0 or not hasattr(self, "attn_op"):
+        return usp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=dtype, t=t)
+
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    half_dtypes = (torch.float16, torch.bfloat16)
+
+    def half(t_):
+        return t_ if t_.dtype in half_dtypes else t_.to(dtype)
+
+    # query, key, value function
+    def qkv_fn(x):
+        q = self.norm_q(self.q(x.to(dtype))).view(b, s, n, d)
+        k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d)
+        v = self.v(x.to(dtype)).view(b, s, n, d)
+        return q, k, v
+
+    q, k, v = qkv_fn(x)
+    q, k = rope_apply_qk(q, k, grid_sizes, freqs)
+
+    # Cross-head compress gate, evaluated on this rank's sequence shard where the full hidden
+    # dim is available: [B, S/P, C] -> [B, S/P, H, D] (all heads).
+    gate = self.attn_op.to_gate_compress(x.to(dtype)).view(b, s, n, d)
+
+    # seq-shard -> head-shard for q/k/v and the gate alike.
+    q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
+    k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
+    v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
+    gate = ulysses_all_to_all(half(gate), scatter_dim=2, gather_dim=1)
+
+    # VSA over the FULL sequence on a head shard; x is unused once gate is supplied.
+    out = self._vsa_dispatch(q, k, v, None, seq_lens, grid_sizes, gate=gate)
+
+    # head-shard -> seq-shard, back to this rank's sequence chunk.
+    out = ulysses_all_to_all(out.to(dtype), scatter_dim=1, gather_dim=2)
+    out = out.flatten(2)
+    out = self.o(out)
+    return out
+
@amp.autocast(enabled=False)
@torch.compiler.disable()
 def s2v_rope_apply(x, grid_sizes, freqs):
@@ -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,562 @@ 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).
+#
+# 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()
+        # The forward consumes the official (q2k_index, q2k_num) pair map_to_index produced:
+        # the first q2k_num[row] entries of each row are that query tile's kv-tile ids and the
+        # -1 pads are never read. It produces M as [B,H,L_BUF] base-2 lse (m_i + log2 l_i),
+        # which is exactly what the varlen backward kernels below consume.
+        out, lse = triton_block_sparse_attn_forward(
+            q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes)
+        ctx.save_for_backward(q_tile, k_tile, v_tile, q2k_idx, q2k_num,
+                              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(
@@ -60,7 +612,13 @@ def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
    pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)

    K = pooled_score.shape[-1]
-    topk = min(K, int(topk_ratio * K))
+    # Clamp to at least one key block. A truncating int() selects ZERO blocks whenever
+    # topk_ratio * K < 1 (short sequence and/or aggressive ratio -- e.g. RoLa's default 0.03
+    # below 2176 tokens, SLA's 0.1 below 640), and the sparse kernel then attends to nothing and
+    # emits all-NaN. Keeps the truncating count everywhere else, so no existing configuration
+    # changes its selection; matches the lower bound the VSA paths already apply
+    # (VSASparseAttention.sparse / _h3_compute_topk).
+    topk = max(1, min(K, int(topk_ratio * K)))
    lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices

    sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)
@@ -1,8 +1,11 @@
+import functools
+import math
 import os
 import warnings

 import torch
 import torch.nn as nn
+import torch.nn.functional as F

 try:
    import flash_attn_interface
@@ -42,7 +45,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):
@@ -341,6 +345,32 @@ def attention(
    return out


+def _head_shard_seq_sum(x, heads_total, head_start):
+    """``torch.sum(x, dim=-2, keepdim=True)`` that is bitwise-invariant to Ulysses head-sharding.
+
+    ``x`` is ``[B, h_local, L, D]``. The pipeline runs under autocast and ``torch.sum`` is in
+    PyTorch's float32 promotion list, so this normalizer accumulates in fp32 over ``L`` and
+    returns fp32; the CUDA ReduceConfig picks its accumulation tree from the output count
+    ``B*H*D``. When head-parallelism shrinks ``H`` (e.g. 12 -> 6 -> 3) the tree changes and the
+    fp32 sum moves in its last bits (~1e-7), which the linear branch's denominator passes through
+    undamped -- a bf16 output would absorb it, fp32 does not. Padding the head dim back to
+    ``heads_total`` makes every rank reduce over the full ``B*heads_total*D`` outputs, i.e. the
+    single-GPU tree, then this shard's own head block is sliced back out, so multi-GPU reproduces
+    the single-GPU result bitwise.
+
+    Uninitialized padding is safe: each output element reads only its own head's rows, so the
+    garbage heads (sliced away) never touch the real ones. ``heads_total=None`` or an already-full
+    head count falls back to the plain reduction, leaving the single-GPU path unchanged.
+    """
+    h_local = x.shape[1]
+    if heads_total is None or h_local == heads_total:
+        return torch.sum(x, dim=-2, keepdim=True)
+    buf = torch.empty(x.shape[0], heads_total, x.shape[2], x.shape[3],
+                      device=x.device, dtype=x.dtype)
+    buf[:, head_start:head_start + h_local].copy_(x)
+    return torch.sum(buf, dim=-2, keepdim=True)[:, head_start:head_start + h_local]
+
+
 class SparseLinearAttention(nn.Module):
    # Modified from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/kernel.py
    def __init__(self, head_dim, topk, feature_map='softmax', BLKQ=64, BLKK=64, use_bf16=True, tie_feature_map_qk=True):
@@ -387,13 +417,18 @@ class SparseLinearAttention(nn.Module):
            nn.init.zeros_(self.proj_l.weight)
            nn.init.zeros_(self.proj_l.bias)

-    def forward(self, q, k, v, return_sparsity=False):
+    def forward(self, q, k, v, return_sparsity=False, heads_total=None, head_start=0):
        R'''
        Args:
            q: queries  [B, L, H, D]
            k: keys     [B, L, H, D].
            v: values   [B, L, H, D].
            return_sparsity: whether to return the actual sparsity.
+            heads_total: full head count before Ulysses head-sharding, or None when unsharded.
+                Kept so the linear branch's fp32 normalizer sum stays bitwise-equal to the
+                single-GPU result under sequence parallelism (see _head_shard_seq_sum).
+            head_start: this shard's offset into the full head range; used only when heads_total
+                is not None.
        '''
        dtype = q.dtype
        
@@ -413,7 +448,7 @@ class SparseLinearAttention(nn.Module):

        def calc_linear(q, k, v):
            kvsum = k.transpose(-1, -2) @ v
-            ksum = torch.sum(k, dim=-2, keepdim=True)
+            ksum = _head_shard_seq_sum(k, heads_total, head_start)
            return (q @ kvsum) / (1e-5 + (q * ksum).sum(dim=-1, keepdim=True))
        o_l = calc_linear(q, k, v)

@@ -424,4 +459,636 @@ 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, gate=None):
+        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.
+            gate: [B, L, H, D] pre-computed ``to_gate_compress`` output; when given, ``x`` is
+                ignored. ``to_gate_compress`` is a ``Linear(dim, num_heads*head_dim)`` and thus
+                mixes heads, so under Ulysses head-parallelism the caller must evaluate it on
+                its own sequence shard (where the FULL hidden dim is still present) and
+                all-to-all the result -- passing the head shard in here keeps this module
+                head-local. ``H`` is then the local head count.
+        Returns:
+            [B, L, H, D] attention output (before the output projection).
+
+        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) if gate is None else gate
+        gate = gate.view(B, L, H, D)
+        gb = _tile(gate, gate.dtype, 'g')
+
+        pq = _h3_pool_tiles(qb, vbs, tile_elems, n_tiles)     # [B,H,n_tiles,D] fp32
+        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, gate=None):
+        R'''
+        Args:
+            q, k, v: [B, L, H, D] POST-RoPE queries/keys/values.
+            x:       [B, L, C] attention-input hidden states (drives the compress gate).
+            return_sparsity: also return the sparse-branch keep ratio.
+            gate: [B, L, H, D] pre-computed ``to_gate_compress`` output; see ``forward_h3``
+                (required under Ulysses head-parallelism, where ``x`` is unavailable in the
+                head-shard layout because the gate Linear mixes heads).
+        Returns:
+            [B, L, H, D] attention output (before the output projection).
+        '''
+        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,
+                                   gate=gate)
+        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) if gate is None else gate
+        gate = gate.view(B, L, H, D).transpose(1, 2)          # [B,H,L,D]
+        o_c = self.compressed(qb, kb, vb)
+        out = o_s + gate.to(o_s.dtype) * o_c.to(o_s.dtype)
+        out = out.transpose(1, 2)                             # [B,L,H,D]
+        if return_sparsity:
+            return out, self.topk
+        return out
@@ -0,0 +1,263 @@
+# 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,):
+        """Enable sequence-parallel inference for the VSA sparse student.
+
+        Reuses the parent's packed-sequence SP assembly -- it sizes ``sp_world_size`` /
+        ``all_gather``, pads the sequence to a group multiple, hands every rank a contiguous row
+        slice and checks ``num_heads`` is divisible by the group -- then re-binds the VSA
+        sequence-parallel processor onto every block that carries an ``attn_op``. The parent
+        binding (``MiniMaxH3MultiGPUsAttnProcessor``) is a DENSE core; leaving it on a sparse
+        block would silently bypass VSA and drop the ``to_gate_compress`` params, so the sparse
+        blocks are swapped to ``MiniMaxH3MultiGPUsVSAAttnProcessor``, which turns the row split
+        into a head split around the VSA op. A dense student (``topk >= 1.0``) has no ``attn_op``
+        and keeps the parent's dense processor untouched.
+        """
+        super().enable_multi_gpus_inference()
+        from ..dist import MiniMaxH3MultiGPUsVSAAttnProcessor
+        for block in self.transformer_blocks:
+            if hasattr(block.attn, "attn_op"):
+                block.attn.set_processor(MiniMaxH3MultiGPUsVSAAttnProcessor())
+
+    def sparse_trainable_params(self):
+        """Names of the extra VSA-branch parameters, for logging / freezing."""
+        return [name for name, _ in self.named_parameters() if "attn_op." in name]
@@ -0,0 +1,251 @@
+# 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)
+
+        out = self._vsa_dispatch(q.to(dtype), k.to(dtype), v.to(dtype), x.to(dtype),
+                                 seq_lens, grid_sizes)
+        x = out.to(dtype).flatten(2)
+        return self.o(x)
+
+    def _vsa_dispatch(self, q, k, v, x, seq_lens, grid_sizes, gate=None):
+        """Run the VSA kernel over POST-RoPE ``q/k/v`` ([B, L, H, D]) -> [B, L, H, D].
+
+        ``x`` drives ``attn_op.to_gate_compress`` unless a pre-computed ``gate`` ([B, L, H, D])
+        is handed in: that Linear mixes heads, so the sequence-parallel path
+        (``usp_attn_fastvideo_forward``) evaluates it on its own sequence shard and passes the
+        head-sharded result, with ``x`` left as None.
+
+        The segment-aware tile geometry needs the sequence length to equal the true video token
+        count. Under sequence parallelism the gathered length is padded to a multiple of the SP
+        degree and can exceed it, so that case takes the per-sample branch below, which attends
+        over the valid grid tokens only and pads the output back to ``L``.
+        """
+        s = q.shape[1]
+        g0 = grid_sizes[0]
+        if (bool((grid_sizes == g0).all()) and s == int(g0.prod())
+                and bool((seq_lens == s).all())):
+            segments = ((int(g0[0]), int(g0[1]), int(g0[2])),)
+            return self.attn_op(q, k, v, x, segments=segments, gate=gate)
+
+        outputs = []
+        for i, (grid, length) in enumerate(zip(grid_sizes.tolist(), seq_lens.tolist())):
+            grid = tuple(int(axis) for axis in grid)
+            length = int(length)
+            if min(grid) <= 0 or length != grid[0] * grid[1] * grid[2] or not 0 < length <= s:
+                raise ValueError("Wan VSA 的 seq_lens 必须等于有效网格 token 数且不超过序列长度。")
+            gate_i = None if gate is None else gate[i:i + 1, :length]
+            # x is only consumed by to_gate_compress, so it is dropped once gate is supplied.
+            x_i = None if gate is not None else x[i:i + 1, :length]
+            oi = self.attn_op(q[i:i + 1, :length], k[i:i + 1, :length],
+                              v[i:i + 1, :length], x_i,
+                              segments=(grid,), gate=gate_i)
+            outputs.append(torch.nn.functional.pad(oi, (0, 0, 0, 0, 0, s - length)))
+        return torch.cat(outputs, dim=0)
+
+
+class FastVideoWanTransformer3DModel(WanTransformer3DModel):
+    """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,):
+        """Bind the VSA-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 ``FastVideoWanSelfAttention.attn_op``
+        and would silently drop the VSA kernel plus the trained ``to_gate_compress`` gate.
+        Rebind to ``usp_attn_fastvideo_forward``: Ulysses head-parallel, with the cross-head
+        compress gate evaluated on each rank's sequence shard (where the full hidden dim still
+        lives) and all-to-all'd next to q/k/v.
+
+        Constraint: the VSA kernel has no ring-attention path (no partial-output/logsumexp
+        interface to rescale across KV chunks), so this is pure Ulysses -- ``ring_degree`` must
+        be 1 and ``num_heads`` must divide the SP degree.
+        """
+        import types
+
+        from ..dist import (get_sequence_parallel_rank,
+                            get_sequence_parallel_world_size, get_sp_group,
+                            usp_attn_fastvideo_forward)
+        self.sp_world_size = get_sequence_parallel_world_size()
+        self.sp_world_rank = get_sequence_parallel_rank()
+        self.all_gather = get_sp_group().all_gather
+
+        # Ulysses head-split needs num_heads divisible by the SP degree.
+        n_heads = self.blocks[0].self_attn.num_heads
+        if self.sp_world_size > 1 and n_heads % self.sp_world_size != 0:
+            raise ValueError(
+                f"VSA multi-GPU (ulysses) needs num_heads ({n_heads}) divisible by the "
+                f"sequence-parallel degree ({self.sp_world_size}); set ulysses_degree "
+                f"accordingly and keep ring_degree=1 (the VSA kernel has no ring path).")
+
+        for block in self.blocks:
+            block.self_attn.forward = types.MethodType(
+                usp_attn_fastvideo_forward, block.self_attn)
+
+        if hasattr(self, 'vace_blocks'):
+            for block in self.vace_blocks:
+                block.self_attn.forward = types.MethodType(
+                    usp_attn_fastvideo_forward, block.self_attn)
+
+    def sparse_trainable_params(self):
+        """Names of the extra VSA-branch parameters, for logging / freezing."""
+        return [name for name, _ in self.named_parameters() if "attn_op." in name]
@@ -111,10 +111,13 @@ class MiniMaxH3MixedPrecisionLoaderMixin:
            kwargs["device_map"] = device_map
        if max_memory is not None:
            kwargs["max_memory"] = max_memory
-        model = super().from_pretrained(pretrained_model_name_or_path, torch_dtype=torch_dtype, **kwargs)
+        result = super().from_pretrained(pretrained_model_name_or_path, torch_dtype=torch_dtype, **kwargs)
+        # ModelMixin returns (model, loading_info) when output_loading_info=True.
+        # Restore precision on the model while preserving that return contract.
+        model = result[0] if isinstance(result, tuple) else result
        if torch_dtype not in (None, torch.float32) and cls._keep_in_fp32_modules:
            restore_fp32_modules(model, pretrained_model_name_or_path, kwargs.get("subfolder", None))
-        return model
+        return result

    @classmethod
    def _from_pretrained_with_missing_keys(cls, pretrained_model_name_or_path, torch_dtype, **kwargs):
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Tuple, Union

@@ -264,7 +264,11 @@ class MiniMaxH3Attention(nn.Module):
        rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        valid_length: Optional[int] = None,
+        vsa_segments=None,
    ) -> torch.Tensor:
+        if vsa_segments is not None:
+            return self.processor(self, hidden_states, rotary_emb, attention_mask, valid_length,
+                                  vsa_segments=vsa_segments)
        return self.processor(self, hidden_states, rotary_emb, attention_mask, valid_length)


@@ -374,6 +378,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
        rotary_emb: Tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        valid_length: Optional[int] = None,
+        vsa_segments=None,
    ) -> torch.Tensor:
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(temb)

@@ -382,7 +387,8 @@ class MiniMaxH3TransformerBlock(nn.Module):
        norm_hidden_states = norm_hidden_states * (
            1.0 + _index_select_rows(scale_msa, adaln_indices)
        ) + _index_select_rows(shift_msa, adaln_indices)
-        attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask, valid_length)
+        attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask, valid_length,
+                                vsa_segments=vsa_segments)
        hidden_states = residual + _index_select_rows(gate_msa, adaln_indices) * attn_output

        residual = hidden_states
@@ -546,6 +552,9 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin

        self.gradient_checkpointing = False
        self.gradient_checkpointing_save_on_cpu = False
+        # Optional non-reentrant checkpoint `context_fn` for the main blocks; the streaming KV cache installs one so
+        # backward recomputes see the cache view of the original forward.
+        self.gradient_checkpointing_context_fn = None
        self.sp_world_size = 1
        self.sp_world_rank = 0
        self.all_gather = None
@@ -747,10 +756,14 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin
            adaln_indices = adaln_indices[rows]
            rotary_emb = (rotary_emb[0][rows], rotary_emb[1][rows])

+        vsa_segments = attention_kwargs.get("vsa_segments") if attention_kwargs else None
+        block_kwargs = {"vsa_segments": vsa_segments} if vsa_segments is not None else {}
        for block in self.transformer_blocks:
            with torch.autograd.graph.save_on_cpu() if self.gradient_checkpointing_save_on_cpu else nullcontext():
                if torch.is_grad_enabled() and self.gradient_checkpointing:
                    ckpt_kwargs = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    if self.gradient_checkpointing_context_fn is not None:
+                        ckpt_kwargs["context_fn"] = self.gradient_checkpointing_context_fn
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        block,
                        hidden_states,
@@ -759,10 +772,12 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin
                        rotary_emb,
                        attention_mask,
                        valid_length,
+                        **block_kwargs,
                        **ckpt_kwargs,
                    )
                else:
-                    hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb, attention_mask, valid_length)
+                    hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb, attention_mask,
+                                          valid_length, **block_kwargs)

        if self.sp_world_size > 1:
            hidden_states = self.all_gather(hidden_states.contiguous(), dim=1)[:, :sequence_length]
@@ -822,11 +837,20 @@ class MiniMaxH3StreamingKVCache:
    Args:
        num_layers (`int`): The number of main transformer blocks (50 for the released H3).
        dtype (`torch.dtype`, defaults to `torch.bfloat16`): The dtype the clean K/V is stored at.
+        offload_to_cpu (`bool`, defaults to `False`): Whether the clean K/V is parked on the CPU.
+        commit_attend_history (`bool`, defaults to `False`): Whether a clean commit pass runs the
+            live attention pattern (media rows attend `[text K/V; persistent history; current chunk
+            K/V]`, text rows self-attend) instead of the legacy standalone pass over
+            `[text | current chunk]`. Enable it when the history must condition the committed K/V,
+            as in the Self-Forcing clean refresh pass.
    """

-    def __init__(self, num_layers: int, dtype: torch.dtype = torch.bfloat16):
+    def __init__(self, num_layers: int, dtype: torch.dtype = torch.bfloat16, offload_to_cpu: bool = False,
+                 commit_attend_history: bool = False):
        self.layer_names = tuple(f"transformer_blocks.{index}.attn" for index in range(num_layers))
        self.dtype = dtype
+        self.offload_to_cpu = offload_to_cpu
+        self.commit_attend_history = commit_attend_history
        self._history: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
        self._staged: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
        self._commit_active = False
@@ -838,6 +862,10 @@ class MiniMaxH3StreamingKVCache:
        # Per-forward attention mode: None ("plain"), "live" or "commit".
        self.mode: Optional[str] = None
        self.token_tags: Optional[torch.Tensor] = None
+        # Activation-checkpoint replay: while set, `history_pair` exposes only this many leading history rows (the
+        # history a checkpointed forward saw); `_epoch` changes whenever rows are dropped, invalidating pending replays.
+        self._history_view_rows: Optional[int] = None
+        self._epoch = 0

    # ------------------------------------------------------------------
    # Per-forward setup consumed by the streaming processor.
@@ -902,11 +930,51 @@ class MiniMaxH3StreamingKVCache:
            self._staged_row_count = int(key.shape[0])
        elif int(key.shape[0]) != self._staged_row_count:
            raise RuntimeError(f"{layer_name}: staged row count differs across layers")
-        self._staged[layer_name] = (key, value)
+        device = torch.device("cpu") if self.offload_to_cpu else key.device
+        self._staged[layer_name] = (
+            key.detach().to(device=device, dtype=self.dtype).contiguous(),
+            value.detach().to(device=device, dtype=self.dtype).contiguous(),
+        )

    def history_pair(self, layer_name: str) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
        r"""Return the persistent `(key, value)` of one layer, or `None` when empty."""
-        return self._history.get(layer_name)
+        pair = self._history.get(layer_name)
+        if pair is None or self._history_view_rows is None:
+            return pair
+        if self._history_view_rows == 0:
+            return None
+        return pair[0][: self._history_view_rows], pair[1][: self._history_view_rows]
+
+    # ------------------------------------------------------------------
+    # Activation checkpointing.
+    # ------------------------------------------------------------------
+    def checkpoint_context(self):
+        r"""
+        `context_fn` for non-reentrant `torch.utils.checkpoint` over a main block.
+
+        Called when the checkpointed forward runs, it snapshots the cache view that forward sees — the mode, the
+        token tags and the committed history length — and returns `(forward_context, recompute_context)`; the
+        recompute context restores that view while backward replays the block, after `end_forward()` cleared the
+        mode and later commits appended more history. The history is append-only (`commit` concatenates), so the
+        snapshot is a row count, not a tensor copy: the replay reads a prefix view of the current history.
+        """
+        history_rows = self.history_tokens if self._history_view_rows is None else self._history_view_rows
+        snapshot = (self.mode, self.token_tags, history_rows, self._epoch)
+        return nullcontext(), self._replay_view(*snapshot)
+
+    @contextmanager
+    def _replay_view(self, mode, token_tags, history_rows, epoch):
+        if epoch != self._epoch:
+            raise RuntimeError(
+                "the streaming KV history was cleared or trimmed between a checkpointed forward and its backward "
+                "recompute; clear the cache only after backward"
+            )
+        saved = (self.mode, self.token_tags, self._history_view_rows)
+        self.mode, self.token_tags, self._history_view_rows = mode, token_tags, history_rows
+        try:
+            yield
+        finally:
+            self.mode, self.token_tags, self._history_view_rows = saved

    # ------------------------------------------------------------------
    # Transaction control.
@@ -931,6 +999,7 @@ class MiniMaxH3StreamingKVCache:
                    torch.cat((previous[1], current[1]), dim=0),
                )
            self._history[layer_name] = combined
+            del current, previous, combined
        self._commit_token_counts.append(len(self._commit_tags))
        self._commit_token_tags.append(self._commit_tags)
        self._commit_active = False
@@ -1005,6 +1074,7 @@ class MiniMaxH3StreamingKVCache:
    def clear(self) -> None:
        r"""Clear persistent and staged KV and reset the logical chunk count."""
        self._history.clear()
+        self._epoch += 1
        self._commit_token_counts.clear()
        self._commit_token_tags.clear()
        self._staged.clear()
@@ -1014,6 +1084,7 @@ class MiniMaxH3StreamingKVCache:
        self.end_forward()

    def _retain_commit_rows(self, selection) -> None:
+        self._epoch += 1
        offsets = [0]
        for count in self._commit_token_counts:
            offsets.append(offsets[-1] + count)
@@ -1068,9 +1139,11 @@ class MiniMaxH3StreamingKVCacheAttnProcessor:
      [`MiniMaxH3AttnProcessor`];
    * *live*: text rows self-attend, while media rows attend to
      `[text K/V; persistent clean history; current chunk K/V]`;
-    * *commit*: a plain pass whose output is discarded, staging the media rows' K/V of
-      every layer into the cache (the K/V of a row depends only on that row's input, so
-      the commit pass's attention output is irrelevant).
+    * *commit*: the pass whose output is discarded, staging the media rows' K/V of every
+      layer into the cache. With `commit_attend_history` off (legacy) it is a plain pass over
+      `[text | current chunk]`; with it on, the commit pass runs the live attention pattern so
+      the staged K/V is conditioned on the persistent history, matching the Self-Forcing clean
+      refresh pass.

    Live and commit passes must be padless (`attention_mask is None`); the streaming
    packed sequences are single documents by construction.
@@ -1107,19 +1180,19 @@ class MiniMaxH3StreamingKVCacheAttnProcessor:
            query = apply_minimax_h3_rotary_emb(query, *rotary_emb)
            key = apply_minimax_h3_rotary_emb(key, *rotary_emb)

-        if cache.mode == "commit":
-            # The committed rows' K/V depend only on their own inputs; the attention output
-            # of a commit pass is discarded by the caller.
+        if cache.mode == "commit" and not cache.commit_attend_history:
+            # Legacy commit: a plain pass over [text | current chunk]; the attention output is
+            # discarded by the caller and only the staged media K/V rows matter.
            output = attention(query, key, value, causal=False)
            rows = cache.commit_row_indices.to(key.device)
            staged_key = key.index_select(1, rows).squeeze(0)
            staged_value = value.index_select(1, rows).squeeze(0)
            cache.stage(
                self.layer_name,
-                staged_key.detach().to(cache.dtype).contiguous(),
-                staged_value.detach().to(cache.dtype).contiguous(),
+                staged_key,
+                staged_value,
            )
-        else:  # live
+        else:  # live, or a commit with commit_attend_history: media rows attend the history.
            tags = cache.token_tags
            text_rows = tags == MINIMAX_H3_STREAMING_TEXT_TAG
            media_rows = ~text_rows
@@ -1153,6 +1226,16 @@ class MiniMaxH3StreamingKVCacheAttnProcessor:
                    values.transpose(1, 2),
                ).transpose(1, 2)

+        if cache.mode == "commit" and cache.commit_attend_history:
+            # The history-conditioned commit pass output is still discarded; only the staged K/V
+            # rows matter, now computed with the history in context.
+            rows = cache.commit_row_indices.to(key.device)
+            cache.stage(
+                self.layer_name,
+                key.index_select(1, rows).squeeze(0),
+                value.index_select(1, rows).squeeze(0),
+            )
+
        output = output.flatten(2, 3).type_as(query)
        output = attn.to_out[0](output)
        output = attn.to_out[1](output)
@@ -1201,21 +1284,26 @@ def install_minimax_h3_streaming_kv_cache(
    transformer: "MiniMaxH3Transformer3DModel",
    *,
    dtype: torch.dtype = torch.bfloat16,
+    offload_to_cpu: bool = False,
+    commit_attend_history: bool = False,
 ) -> MiniMaxH3StreamingKVCache:
    r"""
    Swap every main block's attention processor for the streaming one and return the shared
    cache.

    The token refiner keeps its standard processor: the streaming contract only concerns the
-    main block stack over the media sequence.
+    main block stack over the media sequence. `commit_attend_history` is forwarded to
+    [`MiniMaxH3StreamingKVCache`].
    """
    cache = MiniMaxH3StreamingKVCache(
-        num_layers=len(transformer.transformer_blocks), dtype=dtype
+        num_layers=len(transformer.transformer_blocks), dtype=dtype, offload_to_cpu=offload_to_cpu,
+        commit_attend_history=commit_attend_history,
    )
    for index, block in enumerate(transformer.transformer_blocks):
        block.attn.set_processor(
            MiniMaxH3StreamingKVCacheAttnProcessor(cache, cache.layer_names[index])
        )
+    transformer.gradient_checkpointing_context_fn = cache.checkpoint_context
    return cache


@@ -1223,3 +1311,4 @@ def uninstall_minimax_h3_streaming_kv_cache(transformer: "MiniMaxH3Transformer3D
    r"""Restore the standard attention processor on every main block."""
    for block in transformer.transformer_blocks:
        block.attn.set_processor(MiniMaxH3AttnProcessor())
+    transformer.gradient_checkpointing_context_fn = None
@@ -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]
@@ -51,8 +51,8 @@ class FakeLocalAttention(nn.Module):
        super().__init__()
        self.local_attn = SparseLinearAttention(head_dim, topk)

-    def forward(self, q, k, v):
-        x = self.local_attn(q, k, v)
+    def forward(self, q, k, v, heads_total=None, head_start=0):
+        x = self.local_attn(q, k, v, heads_total=heads_total, head_start=head_start)
        return x


@@ -310,3 +310,45 @@ class TurboWanTransformer3DModel(WanTransformer3DModel):

        if hasattr(self, "img_emb"):
            del self.img_emb
+
+    def enable_multi_gpus_inference(self,):
+        """Bind the SLA-aware USP forward instead of the parent's dense one.
+
+        ``WanTransformer3DModel.enable_multi_gpus_inference`` monkey-patches every
+        ``blocks[i].self_attn.forward`` to ``usp_attn_forward`` (dense full-softmax via
+        ``xFuserLongContextAttention``). That path only touches base-class attributes
+        (``q/k/v/norm_q/norm_k/o/window_size/num_heads/head_dim``), all of which this module's
+        ``WanSelfAttention`` inherits, so the inherited patch used to degrade the SLA student to
+        dense attention SILENTLY -- no AttributeError, no shape error, just ``self.attn_op``
+        (``FakeLocalAttention`` -> ``SparseLinearAttention``) never being reached. Rebind to
+        ``usp_attn_turbowan_forward``: Ulysses head-parallel all-to-all around the SLA kernel,
+        which is head-independent (all reductions over L/D, ``proj_l`` is per-head), so the
+        multi-GPU result matches single-GPU.
+
+        Constraint: the SLA kernel exposes no partial-output/logsumexp interface, so ring
+        attention cannot be wired to it -- pure Ulysses only, ``ring_degree`` must be 1 and
+        ``num_heads`` must divide the SP degree.
+        """
+        from ..dist import (get_sequence_parallel_rank,
+                            get_sequence_parallel_world_size, get_sp_group,
+                            usp_attn_turbowan_forward)
+        self.sp_world_size = get_sequence_parallel_world_size()
+        self.sp_world_rank = get_sequence_parallel_rank()
+        self.all_gather = get_sp_group().all_gather
+
+        # Ulysses head-split needs num_heads divisible by the SP degree.
+        n_heads = self.blocks[0].self_attn.num_heads
+        if self.sp_world_size > 1 and n_heads % self.sp_world_size != 0:
+            raise ValueError(
+                f"SLA multi-GPU (ulysses) needs num_heads ({n_heads}) divisible by the "
+                f"sequence-parallel degree ({self.sp_world_size}); set ulysses_degree "
+                f"accordingly and keep ring_degree=1 (the SLA kernel has no ring path).")
+
+        for block in self.blocks:
+            block.self_attn.forward = types.MethodType(
+                usp_attn_turbowan_forward, block.self_attn)
+
+        if hasattr(self, 'vace_blocks'):
+            for block in self.vace_blocks:
+                block.self_attn.forward = types.MethodType(
+                    usp_attn_turbowan_forward, block.self_attn)
@@ -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)
@@ -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
@@ -27,6 +26,29 @@ DETERMINISTIC_LOGVAR = -30.0
 TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))


+def _load_raw_state_dict(path):
+    """Load a flat tensor state_dict from a .safetensors/.pth/.bin checkpoint.
+
+    Mirrors the normalization done in AutoencoderTinyWan.from_pretrained so that a
+    companion encoder checkpoint can be borrowed the same way.
+    """
+    path = os.fspath(path)
+    if path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+
+        sd = load_file(path)
+    else:
+        sd = torch.load(path, map_location="cpu", weights_only=True)
+    if isinstance(sd, dict):
+        if isinstance(sd.get("state_dict"), dict):
+            sd = sd["state_dict"]
+        elif isinstance(sd.get("decoder"), dict):
+            sd = sd["decoder"]
+    if isinstance(sd, dict) and all(k.startswith("model.") for k in sd):
+        sd = {k[len("model.") :]: v for k, v in sd.items()}
+    return sd
+
+
 def conv(n_in, n_out, **kwargs):
    return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)

@@ -39,7 +61,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 +73,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 +234,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 +260,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 +351,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 +376,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 +398,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):
@@ -329,25 +414,148 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
    - taew2_2*.pth / .safetensors: Wan 2.2 5B (48ch latent, patch_size=2)
    """

+    # The TAE has no checkpointed block path, but the flag is declared so that
+    # `enable_gradient_checkpointing()` (a shared training-script call) is a
+    # no-op instead of raising, mirroring AutoencoderKLWan.
+    _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__(
@@ -361,6 +569,7 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
        temporal_compression_ratio=4,
        spatial_compression_ratio=None,
        parallel=True,
+        decoder_only=False,
    ):
        super().__init__()
        if latent_channels == 16:
@@ -368,7 +577,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)
@@ -378,6 +589,10 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
        if spatial_compression_ratio is None:
            spatial_compression_ratio = 8 * patch_size
        self.spatial_compression_ratio = spatial_compression_ratio
+        # Persist the derived value into the frozen config as well: some pipelines
+        # read `vae.config.spatial_compression_ratio` (not the attribute), and the
+        # register_to_config decorator would otherwise keep the `None` default.
+        self.register_to_config(spatial_compression_ratio=spatial_compression_ratio)

        self.model = TAEWan(
            encoder_time_downscale=tuple(encoder_time_downscale),
@@ -387,6 +602,8 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
            latent_channels=latent_channels,
            arch_variant=arch_variant,
        )
+        if decoder_only:
+            self.model.encoder = nn.Identity()
        self.parallel = parallel
        self.gradient_checkpointing = False

@@ -406,12 +623,12 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
        # latents_mean=zeros/latents_std=ones and passes pipeline latents
        # through unchanged), so no mean/std scaling is applied here: the
        # encoder output is already interchangeable with full-VAE latents.
-        x = (x.add(1).div_(2)).clamp_(0, 1)     # [-1, 1] -> [0, 1]
-        x = x.permute(0, 2, 1, 3, 4)            # BCTHW -> NTCHW
+        x = (x.add(1).div_(2)).clamp_(0, 1)  # [-1, 1] -> [0, 1]
+        x = x.permute(0, 2, 1, 3, 4)  # BCTHW -> NTCHW
        zs = []
        for u in x:
            z = self.model.encode_video(u.unsqueeze(0), parallel=self.parallel).squeeze(0)
-            zs.append(z.permute(1, 0, 2, 3))    # TCHW -> CTHW
+            zs.append(z.permute(1, 0, 2, 3))  # TCHW -> CTHW
        return torch.stack(zs)

    @apply_forward_hook
@@ -432,7 +649,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)
@@ -449,49 +666,131 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
            return (decoded,)
        return DecoderOutput(sample=decoded)

+    def forward(self, z: torch.Tensor, return_dict: bool = True):
+        # Use forward so decoder training runs through DDP hooks.
+        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()}
+
+        # Borrow an encoder when the checkpoint is decoder-only. Some taew2_2_super
+        # releases ship only the Super decoder; without an encoder, encode() would
+        # raise. We merge encoder.* weights from a companion base checkpoint so the
+        # model becomes a full encoder(base)+decoder(super) TAE.
+        # Priority: explicit `encoder_path` (via additional_kwargs / vae_kwargs or
+        # kwargs) > auto sibling found by stripping "_super" from the filename.
+        encoder_path = None
+        if isinstance(additional_kwargs, dict):
+            encoder_path = additional_kwargs.get("encoder_path")
+        if encoder_path is None:
+            encoder_path = kwargs.get("encoder_path")
+        if not any(k.startswith("encoder.") for k in state_dict):
+            if encoder_path is None:
+                root, ext = os.path.splitext(pretrained_model_path)
+                sibling = root.replace("_super", "") + ext
+                if sibling != pretrained_model_path and os.path.exists(sibling):
+                    encoder_path = sibling
+            if encoder_path is not None and os.path.exists(encoder_path):
+                enc_sd = {k: v for k, v in _load_raw_state_dict(encoder_path).items() if k.startswith("encoder.")}
+                if enc_sd:
+                    state_dict = {**enc_sd, **state_dict}
+                    print(f"[TAE] decoder-only 权重：已从 {encoder_path} 借用 {len(enc_sd)} 个 encoder 键以支持 encode。")
+            if not any(k.startswith("encoder.") for k in state_dict):
+                print("[TAE][warn] 未找到 encoder 权重，将以 decoder_only 加载，encode() 不可用。"
+                      "可通过 vae_kwargs.encoder_path=<base taew2_2 权重路径> 指定，"
+                      "或把 base 权重放在与 super 同名（去掉 _super）的路径下。")
+
+        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
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
bubbliiiing	d6024c0473	Update one forcing	2026-11-02 00:11:06 +08:00
bubbliiiing	be4fa825cd	Update multigpus in vsa, sla, spark	2026-11-01 10:00:44 +08:00
bubbliiiing	f00a7d480e	Update one-forcing and causal forcing params	2026-10-30 16:33:15 +08:00
bubbliiiing	a15b293128	Update ema and update taehv	2026-10-29 10:59:15 +08:00
Bubbliiiing	e8c314fbfa	Update Sparkwan, Reward Model Training, FastVideo (#525 )	2026-10-28 11:34:50 +08:00
hkz	9dc6f1d49a	Add the training script for taew2_2_super (#526 )	2026-10-27 23:16:10 +08:00
li-lizhe	43681aebea	fix(offload): resolve the group offloading transfer stream from the onload device (#528 ) The vendored group-offloading helper (see the "Modified from .../diffusers/v0.36.0" header) resolved its data-transfer stream by probing for CUDA/XPU only, so `gpu_memory_mode="model_group_offload"`, which calls `apply_group_offloading(..., use_stream=True)`, raised on a host whose accelerator is neither of those: ValueError: Using streams for data transfer requires a CUDA device, or an Intel XPU device. even when the accelerator's backend implements streams (for example Ascend NPU, where `torch.cuda.is_available()` is False while `torch.npu` is available and provides `Stream`). Resolve the stream from the onload device's own backend instead, and keep rejecting devices whose backend cannot stream, with a message that names the device type. This is the same policy upstream diffusers applies in the function this file is derived from. Signed-off-by: li-lizhe <147392333@qq.com>	2026-10-27 23:15:38 +08:00
li-lizhe	13bb3c625e	fix(z-image): strip the ComfyUI `model.diffusion_model.` key prefix before converting (#522 ) LoadZImageTransformerModel.loadmodel decides whether to run convert_state_dict by looking for a bare `x_embedder.weight` key, and the converter itself assumes bare keys. ComfyUI-native checkpoints (e.g. CivitAI Z-Image fine-tunes) store every tensor under a `model.diffusion_model.` prefix, so the gate never fired, the prefix survived, and the strict load_state_dict failed with "Missing key(s): x_pad_token ... / Unexpected key(s): model.diffusion_model...". Strip the prefix when every key carries it, then convert as before. Bare-key checkpoints are untouched. Signed-off-by: li-lizhe <147392333@qq.com>	2026-10-14 09:43:58 +08:00
Bubbliiiing	4b7b6402a1	Upgrade GPU Memory Mode architecture (#521 )	2026-09-29 19:18:17 +08:00