Update sparkwan

examples/sparkdiffusion/predict_t2v.py

@@ -0,0 +1,369 @@
+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, 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     = "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.
+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 by default: 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.
+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"
+# model path
+model_name          = "models/Diffusion_Transformer/Wan2.2-T2V-A14B"
+
+# SparkDiffusion sparse-student config.
+# student_topk is the ratio of key blocks each query keeps in the sparse (RoLa) attention and
+# MUST match the value used in training (scripts/sparkdiffusion); sparsity = 1 - student_topk.
+student_topk        = 0.03
+
+# 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. 
+# Used when the sampler is in "Flow_Unipc", "Flow_DPM++".
+shift               = 12
+
+# Load pretrained model if need
+# The transformer_path is used for low noise model, the transformer_high_path is used for high noise model.
+# Each path doubles as a per-expert SparkDiffusion sparse-student checkpoint: an expert whose path is
+# set is rebuilt as SparkWan2_2Transformer3DModel (sparse RoLa self-attention) and its checkpoint is
+# loaded after set_student_topk, while a None path keeps that expert dense. Wan2.2 A14B is a two-expert
+# MoE (low_noise -> transformer, high_noise -> transformer_2), so the full student sets both.
+transformer_path        = "output_dir_sparkdiffusion_wan2.2_t2v_low_sparse/checkpoint-2200/diffusion_pytorch_model.safetensors"
+transformer_high_path   = "output_dir_sparkdiffusion_wan2.2_t2v_high_sparse/checkpoint-17000/diffusion_pytorch_model.safetensors"
+vae_path                = None
+# Load lora model if need
+# The lora_path is used for low noise model, the lora_high_path is used for high noise model.
+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 distilled student is CFG-free
+seed                = 43
+num_inference_steps = 4
+# The lora_weight is used for low noise model, the lora_high_weight is used for high noise model.
+lora_weight         = 0.55
+lora_high_weight    = 0.55
+save_path           = "samples/wan-videos-sparkdiffusion-t2v"
+
+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. An expert becomes a SparkWan2_2Transformer3DModel (sparse
+# RoLa self-attention) iff its own sparse checkpoint path is set; otherwise it stays dense.
+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:
+    # low-noise expert is a sparse student: rebuild the sparse kernels at the training top-k.
+    # Call order matters: set_student_topk recreates attn_op, so its checkpoint is loaded below.
+    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:
+        # high-noise expert is a sparse student (see the low-noise branch above).
+        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 checkpoints AFTER set_student_topk so a sparse student's trained RoLa params
+# overwrite the freshly-init ones; a missing attn_op.* key means the checkpoint is not a sparse student.
+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 "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)}")
+    if len([k for k in m if "attn_op." in k]) > 0:
+        print("WARNING: sparse attn_op keys are missing on the low-noise expert -- check student_topk / ckpt.")
+
+if transformer_2 is not None:
+    if transformer_high_path is not None:
+        print(f"From checkpoint: {transformer_high_path}")
+        if transformer_high_path.endswith("safetensors"):
+            from safetensors.torch import load_file
+            state_dict = load_file(transformer_high_path)
+        else:
+            state_dict = torch.load(transformer_high_path, map_location="cpu")
+        state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+
+        m, u = transformer_2.load_state_dict(state_dict, strict=False)
+        print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+        if len([k for k in m if "attn_op." in k]) > 0:
+            print("WARNING: sparse attn_op keys are missing on the high-noise expert -- check student_topk / ckpt.")
+
+# 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 attention only supports ulysses head-parallel (no ring-attention path); a sparse
+    # expert under ring_degree>1 would silently mis-attend, so require ring_degree=1.
+    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.
+# The order lives inside the helper: quantization has to happen before the offload hooks
+# are installed, and both transformers of this MoE setup are handled in one call.
+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()

examples/wan2.1_flex_forcing/predict_t2v.py

@@ -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

examples/wan2.1_flex_forcing/predict_t2v_edit.py

@@ -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)

scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1.md

@@ -0,0 +1,414 @@
+# 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 CrossDistill), 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 **CrossDistill** few-step 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 — CrossDistill** ([`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 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 — plus a PCM-style diversity term.
+>
+> | 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 — 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.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 CrossDistill**. 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`, which the distill launcher mirrors via `--teacher_timestep_shift=5.0`.
+
+### 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=400`. 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 --main_process_port=0 --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=400 \
+  --checkpointing_steps=100 \
+  --learning_rate=2e-05 \
+  --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.03 \
+  --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.03` ≈ 97% 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-400/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` 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 — while the CrossDistill diversity term (`--div_mode="pcm_random" --loss_scale_div=10.0`) is kept.
+
+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-400/diffusion_pytorch_model.safetensors"
+NCCL_DEBUG=INFO
+
+accelerate launch --main_process_port=0 --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=4e-07 \
+  --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" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --teacher_timestep_shift=5.0 \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --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.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.03` |
+| `--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-400/diffusion_pytorch_model.safetensors` |
+| `--distill_loss` | Student distribution-matching loss (`dmd` / `tdm`, mutually exclusive) | `dmd` |
+| `--max_simulation_steps_fake` | Student Euler steps in the backward simulation; `>= 2` keeps the first segment `no_grad` | `2` |
+| `--div_mode` / `--loss_scale_div` | PCM-style diversity term | `pcm_random` / `10.0` |
+| `--teacher_timestep_shift` | Frozen dense teacher timestep shift; matches the Wan2.1 scheduler `shift` | `5.0` |
+| `--real_guidance_scale` | CFG scale of the frozen dense teacher | `5.0` |
+| `--learning_rate` / `--learning_rate_critic` | Generator / critic learning rate | `2e-06` / `4e-07` |
+| `--checkpointing_steps` | Save checkpoint every N steps | `200` (distill) / `100` (warm-up) |
+| `--max_grad_norm` | Gradient clipping threshold | `0.05` |
+| `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` |
+
+> 💡 The rf time-grid parameters (`--rf_t_min`, `--rf_t_max`, `--backward_timesteps`) span the full range `[0, 1]` by default for a single-expert model; leave them at their `-1` / empty defaults.
+
+### 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-400/diffusion_pytorch_model.safetensors"
+NCCL_DEBUG=INFO
+
+accelerate launch --main_process_port=0 --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=4e-07 \
+  --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" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --teacher_timestep_shift=5.0 \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --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 --main_process_port=0 --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.03` |
+| `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)

scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1_zh-CN.md

@@ -0,0 +1,414 @@
+# SparkDiffusion 稀疏蒸馏训练指南（Wan2.1-T2V-1.3B）
+
+本文档给出 Wan2.1-T2V-1.3B 上 **SparkDiffusion** 的完整流程：一条两阶段管线，先暖启一个**稀疏（RoLa）学生**，再把它**蒸馏**成少步、CFG-free 的生成器。内容涵盖环境配置、数据准备、两个启动脚本（Stage-1 暖启 + Stage-2 CrossDistill）、以及稀疏少步推理。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）与 **CrossDistill** 少步配方结合起来。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 — CrossDistill**（[`train_distill_wan2.1.py`](./train_distill_wan2.1.py)）：对着 **dense** teacher 把稀疏学生蒸馏成少步生成器。它是分布匹配蒸馏（DMD），含三个同驻网络——可训练的稀疏 `generator`、冻结的 dense `real_score` teacher、可训练的 dense `fake_score` critic——并额外加一项 PCM 式多样性损失。
+>
+> | 阶段 | 启动脚本 | 产物 |
+> |-------|----------|--------|
+> | **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 — 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.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 CrossDistill**。请手动按此顺序运行两个启动脚本。
+
+### 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`，蒸馏启动脚本用 `--teacher_timestep_shift=5.0` 与之对齐。
+
+### 3.2 Stage 1 — 稀疏暖启
+
+Stage 1 插入 RoLa 稀疏自注意力，并用 step-local flow-matching 目标短暂暖启。这是一个**短**阶段——短暂暖启即足以为蒸馏建立可用先验，因此启动脚本把它限制在 `--max_train_steps=400`。**不要**用该阶段自身的采样质量来评判它；清晰度是在 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 --main_process_port=0 --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=400 \
+  --checkpointing_steps=100 \
+  --learning_rate=2e-05 \
+  --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.03 \
+  --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.03` ≈ 97% 稀疏。**必须与** 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-400/diffusion_pytorch_model.safetensors`。
+
+### 3.3 Stage 2 — CrossDistill
+
+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 接力——同时保留 CrossDistill 多样性项（`--div_mode="pcm_random" --loss_scale_div=10.0`）。
+
+用 [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-400/diffusion_pytorch_model.safetensors"
+NCCL_DEBUG=INFO
+
+accelerate launch --main_process_port=0 --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=4e-07 \
+  --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" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --teacher_timestep_shift=5.0 \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --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.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.03` |
+| `--sparse_student` | （Stage 2）把 generator 建成稀疏学生 | - |
+| `--generator_transformer_path` | （Stage 2）载入 generator 的 Stage-1 暖启 checkpoint | `.../checkpoint-400/diffusion_pytorch_model.safetensors` |
+| `--distill_loss` | 学生分布匹配损失（`dmd` / `tdm`，互斥） | `dmd` |
+| `--max_simulation_steps_fake` | backward simulation 中学生的 Euler 步数；`>= 2` 使首段保持 `no_grad` | `2` |
+| `--div_mode` / `--loss_scale_div` | PCM 式多样性项 | `pcm_random` / `10.0` |
+| `--teacher_timestep_shift` | 冻结 dense teacher 的 timestep shift；与 Wan2.1 调度器 `shift` 对齐 | `5.0` |
+| `--real_guidance_scale` | 冻结 dense teacher 的 CFG 尺度 | `5.0` |
+| `--learning_rate` / `--learning_rate_critic` | 生成器 / critic 学习率 | `2e-06` / `4e-07` |
+| `--checkpointing_steps` | 每 N 步保存 checkpoint | `200`（蒸馏）/ `100`（暖启） |
+| `--max_grad_norm` | 梯度裁剪阈值 | `0.05` |
+| `--resume_from_checkpoint` | 续训路径或 `"latest"` | `latest` |
+
+> 💡 rf 时间网格参数（`--rf_t_min`、`--rf_t_max`、`--backward_timesteps`）对单专家模型默认覆盖整个区间 `[0, 1]`；请保持它们的 `-1` / 空默认值。
+
+### 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-400/diffusion_pytorch_model.safetensors"
+NCCL_DEBUG=INFO
+
+accelerate launch --main_process_port=0 --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=4e-07 \
+  --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" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --teacher_timestep_shift=5.0 \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --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 --main_process_port=0 --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.03` |
+| `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)

scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2.md

@@ -0,0 +1,451 @@
+# 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=400` (≈2,000 videos at 8 GPUs × batch 1). 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 --main_process_port=0 --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=400 \
+  --checkpointing_steps=100 \
+  --learning_rate=2e-05 \
+  --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.03 \
+  --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.03` ≈ 97% 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-400/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.
+
+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-400/diffusion_pytorch_model.safetensors"
+NCCL_DEBUG=INFO
+
+accelerate launch --main_process_port=0 --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.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=4e-07 \
+  --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" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --boundary_type="high" \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --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 backward-simulation 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.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.03` |
+| `--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-400/diffusion_pytorch_model.safetensors` |
+| `--distill_loss` | Student distribution-matching loss (`dmd` / `tdm`, mutually exclusive) | `dmd` |
+| `--max_simulation_steps_fake` | Student Euler steps in the backward simulation; per-expert this equals the expert's trajectory-segment count | `2` |
+| `--div_mode` / `--loss_scale_div` | PCM-style diversity term (high expert only) | `pcm_random` / `10.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 backward-simulation start latent | high-distill ckpt |
+| `--learning_rate` / `--learning_rate_critic` | Generator / critic learning rate | `2e-06` / `4e-07` |
+| `--checkpointing_steps` | Save checkpoint every N steps | `200` (distill) / `100` (warm-up) |
+| `--max_grad_norm` | Gradient clipping threshold | `0.05` |
+| `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` |
+
+> 💡 The rf time-grid parameters (`--rf_t_min`, `--rf_t_max`, `--backward_timesteps`, `--joint_backward_timesteps`, `--init_rollout_timesteps`) are **auto-derived** from `--boundary_type` and the config boundary; leave them at their `-1` / empty defaults.
+
+### 3.5 FSDP Backend and Memory Notes
+
+Stage 2 keeps **three** 14B transformers co-resident (sparse generator + dense teacher + dense critic), 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.** The backward recomputes two grad-carrying 14B forwards, 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`.
+
+> 💡 If lowering the resolution is not enough, 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 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 `--div_mode` / `--loss_scale_div` flags, 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-400/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 --main_process_port=0 --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=4e-07 \
+  --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" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --boundary_type="high" \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --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-400/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 --main_process_port=0 --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=4e-07 \
+  --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" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --boundary_type="high" \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --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.03` |
+| `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`

scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2_zh-CN.md

@@ -0,0 +1,451 @@
+# 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=400`（8 卡 × batch 1 时约 2,000 个视频）。**不要**用该阶段自身的采样质量来评判它；清晰度是在 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 --main_process_port=0 --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=400 \
+  --checkpointing_steps=100 \
+  --learning_rate=2e-05 \
+  --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.03 \
+  --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.03` ≈ 97% 稀疏。**必须与** 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-400/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。
+
+用 [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-400/diffusion_pytorch_model.safetensors"
+NCCL_DEBUG=INFO
+
+accelerate launch --main_process_port=0 --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.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=4e-07 \
+  --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" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --boundary_type="high" \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --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` | 不使用 | 用训练好的**高噪**学生暖启 backward-simulation 起点 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.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.03` |
+| `--sparse_student` | （Stage 2）把 generator + init_rollout 建成稀疏学生 | - |
+| `--generator_transformer_path` | （Stage 2）载入 generator 的 Stage-1 暖启 checkpoint | `.../checkpoint-400/diffusion_pytorch_model.safetensors` |
+| `--distill_loss` | 学生分布匹配损失（`dmd` / `tdm`，互斥） | `dmd` |
+| `--max_simulation_steps_fake` | backward simulation 中学生的 Euler 步数；按专家拆分时等于该专家的轨迹段数 | `2` |
+| `--div_mode` / `--loss_scale_div` | PCM 式多样性项（仅高噪专家） | `pcm_random` / `10.0` |
+| `--real_guidance_scale` | 冻结 dense teacher 的 CFG 尺度 | `5.0` |
+| `--init_rollout_path` | （低噪专家）用于滚动 backward-simulation 起点 latent 的高噪学生 | 高噪蒸馏 ckpt |
+| `--learning_rate` / `--learning_rate_critic` | 生成器 / critic 学习率 | `2e-06` / `4e-07` |
+| `--checkpointing_steps` | 每 N 步保存 checkpoint | `200`（蒸馏）/ `100`（暖启） |
+| `--max_grad_norm` | 梯度裁剪阈值 | `0.05` |
+| `--resume_from_checkpoint` | 续训路径或 `"latest"` | `latest` |
+
+> 💡 rf 时间网格参数（`--rf_t_min`、`--rf_t_max`、`--backward_timesteps`、`--joint_backward_timesteps`、`--init_rollout_timesteps`）会由 `--boundary_type` 与 config 分界**自动推导**；请保持它们的 `-1` / 空默认值。
+
+### 3.5 FSDP Backend 与显存说明
+
+Stage 2 需让**三个** 14B 网络同驻（稀疏 generator + dense teacher + dense critic），且都以优化器态为主，因此必须使用分片后端——此处为 FSDP `FULL_SHARD`。蒸馏启动脚本组合了下列显存杠杆：
+
+- `--fsdp_sharding_strategy "FULL_SHARD"`——在多卡间分片 param/grad/optimizer。
+- `--fsdp_backward_prefetch "BACKWARD_POST"`——降低 backward 期间的中位驻留。
+- `--gradient_checkpointing`——在 backward 中重算激活。
+- **OOM 的杠杆是分辨率、而非帧数。** backward 会重算两段带梯度的 14B 前向，因此显存随 `seq_len` 增长；启动脚本使用 `--fix_sample_size 480 832` 配 `--video_sample_n_frames=81`。若 OOM，请降低 `--fix_sample_size`（例如降到 `320 512`），而不是减少帧数。
+- 帧数须满足 `(N − 1) % 4 == 0`。
+
+> 💡 若降低分辨率仍不够，可在 `accelerate launch` 行加上 `--fsdp_offload_params true`，把 param/grad/optimizer offload 到 CPU（更慢）。`--low_vram` 只把冻结侧停到 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` / `--loss_scale_div` 旗标、加上 `--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-400/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 --main_process_port=0 --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=4e-07 \
+  --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" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --boundary_type="high" \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --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-400/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 --main_process_port=0 --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=4e-07 \
+  --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" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --boundary_type="high" \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --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.03` |
+| `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`

scripts/sparkdiffusion/train_distill_wan2.1.sh

@@ -0,0 +1,67 @@
+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).
+# Env-overridable. --student_topk below MUST match the Stage-1 value (0.03).
+export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-400/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 CrossDistill recipe adds a PCM-style diversity term (--div_mode pcm_random --loss_scale_div 10.0), which needs
+# --max_simulation_steps_fake >= 2 so the first backward-simulation segment stays no_grad; --teacher_timestep_shift
+# matches the Wan2.1 scheduler shift (5.0). Memory: --fsdp_offload_params true + --gradient_checkpointing keep the
+# three nets (student/teacher/critic) off the optimizer-state OOM path. No --max_train_steps: this phase runs until
+# you stop it (kill the train_distill_wan2.1.py processes).
+accelerate launch --main_process_port=0 --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=4e-07 \
+  --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 \
+  --trainable_modules "." \
+  --train_mode="normal" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --teacher_timestep_shift=5.0 \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --generator_transformer_path=$SPARSE_STUDENT_CKPT \
+  --resume_from_checkpoint="latest"

scripts/sparkdiffusion/train_distill_wan2.2_high.sh

@@ -0,0 +1,68 @@
+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 (train_finetune_sparse_wan2.2_high.sh -> _high_sparse_rola/checkpoint-400).
+# --student_topk below MUST match the Stage-1 value (0.03).
+export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-400/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: 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 range. The CrossDistill
+# recipe adds a PCM-style diversity term (--div_mode pcm_random --loss_scale_div 10.0) on the high expert
+# only, which needs --max_simulation_steps_fake >= 2 so the first backward-simulation segment stays
+# no_grad. 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: this phase runs until you stop it (kill the train_distill_wan2.2.py processes).
+accelerate launch --main_process_port=0 --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.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=4e-07 \
+  --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 \
+  --trainable_modules "." \
+  --train_mode="normal" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --div_mode="pcm_random" \
+  --loss_scale_div=10.0 \
+  --boundary_type="high" \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --generator_transformer_path=$SPARSE_STUDENT_CKPT \
+  --resume_from_checkpoint="latest"

scripts/sparkdiffusion/train_distill_wan2.2_low.sh

@@ -0,0 +1,76 @@
+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-400/diffusion_pytorch_model.safetensors"
+# Optional warm-start of the backward-simulation starting latent 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: same sparse-student DMD as
+# train_distill_wan2.2_high.sh, with two low-specific differences. (1) --boundary_type="low" starts the few-step
+# trajectory at the MoE boundary and keeps only the low-segment knots, while --rf_t_max=1.0 lets the
+# DMD/critic comparison time s_D span the FULL noise range -- mirrors official wan2.2_joint.sh's
+# joint_low_d_time_full_range=true (the low expert still emits a complete x0, so its fake/teacher scores
+# are matched at all levels, not only s < boundary). (2) The low expert is div-free (no --div_mode /
+# --loss_scale_div): the CrossDistill
+# recipe keeps the PCM diversity term on the high expert only. The low segment of the few-step trajectory
+# starts at the MoE boundary rather than pure noise, so --init_rollout_path warm-starts it from the frozen
+# trained HIGH student (INIT_ROLLOUT_CKPT); leaving it empty makes the untrained low student roll the
+# high-noise segment itself and the samples degrade. 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).
+accelerate launch --main_process_port=0 --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.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=4e-07 \
+  --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 \
+  --trainable_modules "." \
+  --train_mode="normal" \
+  --distill_loss="dmd" \
+  --max_simulation_steps_fake=2 \
+  --boundary_type="low" \
+  --rf_t_max=1.0 \
+  --real_guidance_scale=5.0 \
+  --sparse_student \
+  --student_topk=0.03 \
+  --generator_transformer_path=$SPARSE_STUDENT_CKPT \
+  --init_rollout_path=$INIT_ROLLOUT_CKPT \
+  --resume_from_checkpoint="latest"

scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh

@@ -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.03 (~97% 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 --main_process_port=0 --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=400 \
+  --checkpointing_steps=100 \
+  --learning_rate=2e-05 \
+  --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=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 \
+  --trainable_modules "attn_op" \
+  --student_topk=0.03 \
+  --train_mode="normal" \
+  --resume_from_checkpoint="latest"

scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh

@@ -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.03 (~97% 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 --main_process_port=0 --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=400 \
+  --checkpointing_steps=100 \
+  --learning_rate=2e-05 \
+  --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=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 \
+  --trainable_modules "attn_op" \
+  --student_topk=0.03 \
+  --train_mode="normal" \
+  --boundary_type="high" \
+  --resume_from_checkpoint="latest"

scripts/sparkdiffusion/train_finetune_sparse_wan2.2_low.sh

@@ -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.03
+# must match train_distill_wan2.2_low.sh. Its checkpoint feeds train_distill_wan2.2_low.sh as the low sparse student.
+accelerate launch --main_process_port=0 --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=400 \
+  --checkpointing_steps=100 \
+  --learning_rate=2e-05 \
+  --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=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 \
+  --trainable_modules "attn_op" \
+  --student_topk=0.03 \
+  --train_mode="normal" \
+  --boundary_type="low" \
+  --resume_from_checkpoint="latest"

scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL.md

@@ -12,11 +12,9 @@ This document provides the complete workflow for **Flex-Forcing Stage 2 — Dist
 > 2. **§3.2 Pyramid timestep chunking** — after finishing step `t+1`, the partition for step `t` is *refined* by inserting boundaries while preserving the existing ones. The step-`t` result is buffered over the whole original chunk, then each sub-chunk resumes autoregressively once the KV it needs is available. Coarse planning first, fine detail later.
 > 3. **§3.3 K-Projection** — clean-context keys live in a different space from noisy queries at timestep `t`, so an identity-initialised, timestep-conditioned projection maps the cache into the current noisy latent space. It is applied on the fly, never mutates the cache, and trains together with the generator.
 >
-> Training draws a **fresh random partition every iteration** (chunk sizes 2–10), which is what lets one set of weights cover the whole causal↔bidirectional spectrum. §4.2 of the paper then reuses the same machinery for **any-order editing**.
->
 > The Flex-Forcing pipeline has two stages:
 >
-> 1. **Stage 1 — ODE Regression** (`train_ode.py`): regress the teacher's ODE trajectory so the model becomes a competent few-step generator. This is where the random partitions are first seen.
+> 1. **Stage 1 — ODE Regression** (`train_ode.py`): regress the teacher's ODE trajectory so the model becomes a competent few-step generator. 
 > 2. **Stage 2 — DMD Distillation** (`train_distill.py`, this README): compress to a **4-step** generator with a **Wan2.1-T2V-14B** real-score teacher, and add the pyramid timestep schedule.
 >
 > This README covers **Stage 2 only**. See [README_TRAIN_ODE.md](./README_TRAIN_ODE.md) for Stage 1.
@@ -210,7 +208,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
 
-The launcher pins: 81 pixel frames = **21 latent frames** (a 5 s clip) at **480×832**, a **4-step** model (`--denoising_step_indices_list 1000 750 500 250`, jittered per iteration by `--randomize_step_indices`), `--flex_pyramid_levels=4` - which over 21 latent frames reproduces the paper's ladder `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` - partitions drawn in **2..10** (plus, on half the iterations, the whole-clip coarse layout of §3.2 - see below), `--flow_euler_rollout`, `diag_rank1` K-Projection, batch 8 (8 GPUs × `--gradient_accumulation_steps=1`).
+The launcher pins: 81 pixel frames = **21 latent frames** (a 5 s clip) at **480×832**, a **4-step** model (`--denoising_step_indices_list 1000 750 500 250`, jittered per iteration by `--randomize_step_indices`), `--flex_pyramid_levels=4` - which over 21 latent frames reproduces the paper's ladder `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` - the *random* arm's partitions drawn in **2..10** (that arm is only 10% of iterations; another 60% start from the whole-clip coarse layout of §3.2 - see below), `--flow_euler_rollout`, `diag_rank1` K-Projection, batch 8 (8 GPUs × `--gradient_accumulation_steps=1`).
 
 Output: `output_dir_wan2.1_flex_forcing_distill/`.
 
@@ -243,7 +241,7 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
 | `--denoising_step_indices_list` | Denoising step indices (DMD core param). Mapped through `timesteps[train_sampling_steps - idx]`, so `1000 750 500 250` becomes t = `[1000, 937.5, 833.33, 625]` at `shift=5.0` - bit-identical to what `stochastic_sampling_timesteps(4, shift)` feeds the pipeline. | `1000 750 500 250` |
 | `--randomize_step_indices` | Jitter every index but the first symmetrically within `--index_jitter_ratio` of its neighbouring gap each iteration. The expected schedule is unchanged and monotonicity is enforced by construction, so inference's fixed t stays inside every step's trained envelope. | off |
 | `--index_jitter_ratio` | Jitter budget as a fraction of the neighbouring gap. At `1000 750 500 250` with 0.3, 20 000 draws give idx `[675,825] / [425,575] / [175,325]`, i.e. t `[912.2,959.3] / [787.0,871.2] / [514.7,706.5]`; the inference points 937.5 / 833.33 / 625 all fall inside. | 0.3 |
-| `--flow_euler_rollout` | Advance the generator's self-rollout with a deterministic Euler ODE step in flow space (`x_next = x_t + (sigma_next - sigma_t) * v`, fp32) instead of converting to x0 and re-noising with fresh noise. Algebraically this *is* re-noising with the same implied noise - measured `8.9e-16` apart in fp64, versus up to `6.4` for a fresh draw - so it removes the rollout's stochasticity and matches normal flow-matching inference. The final step still converts to x0, which is what DMD is defined on. The critic re-rolls the generator and follows the same switch. | off |
+| `--flow_euler_rollout` | Advance the generator's self-rollout with a deterministic Euler ODE step in flow space (`x_next = x_t + (sigma_next - sigma_t) * v`, fp32) instead of converting to x0 and re-noising with fresh noise. Algebraically this *is* re-noising with the same implied noise - identical to it in exact arithmetic, whereas drawing fresh noise is an independent sample - so it removes the rollout's stochasticity and matches normal flow-matching inference. The final step still converts to x0, which is what DMD is defined on. The critic re-rolls the generator and follows the same switch. | off |
 | `--output_dir` | Output directory | `output_dir_wan2.1_flex_forcing_distill` |
 | `--gradient_checkpointing` | Trade compute for memory | - |
 | `--max_grad_norm` | Gradient clipping threshold | 0.05 |
@@ -256,24 +254,14 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
 | Parameter | Paper | Description | Default |
 |-----------|-------|-------------|---------|
 | `--flex_forcing` | §3.1 | Instantiate `WanTransformer3DModel_FlexForcing` and validate with `WanFlexForcingPipeline`. Off = the inherited path, unchanged. | off |
-| `--flex_chunk_min` | §3.1 | Smallest chunk drawn per iteration. `1` is legal but adds no coverage under a pyramid - refining a 2-frame chunk already yields 1-frame leaves - and it would let single-frame chunks appear mid-clip. | 2 |
-| `--flex_chunk_max` | §3.1 | Largest chunk drawn per iteration. Set equal to `--flex_chunk_min` to pin one fixed layout. **Keep it below the latent frame count**: the whole-clip layout is already covered by the mixture below, so raising this to the frame count only spends draws on a layout you already have and thins out §3.1 (at 21 latent frames, `max=21` leaves 3649 distinct partitions versus 4882 at `max=10`). | 10 |
-| `--flex_pyramid_levels` | §3.2 | `1` = one partition per iteration. `>1` = nest it into a coarse-to-fine ladder, one level per denoising step; both training paths roll it out. **Stage 2 only.** | 1 |
+| `--flex_chunk_min` | §5 | Smallest chunk drawn by the *random* arm (10% of iterations, regardless of `--flex_pyramid_levels`); the other three arms pin their own layouts and split the remaining 90% equally. `1` is legal but adds no coverage under a pyramid - refining a 2-frame chunk already yields 1-frame leaves - and it would let single-frame chunks appear mid-clip. | 2 |
+| `--flex_chunk_max` | §5 | Largest chunk drawn by the *random* arm. Set equal to `--flex_chunk_min` to pin that arm to one fixed layout. **Keep it below the latent frame count**: the whole-clip layout is its own arm, so raising this to the frame count only spends random-arm draws on a layout you already have and thins out every other one (exhaustive enumeration at 21 latent frames: 6441 distinct partitions at `max=10`, 6765 at `max=21` - but `max=21` makes `[21]` alone 1/19 of that arm and cuts `P([3]*7)` from `1/183708` to `1/1106560`). | 10 |
+| `--flex_pyramid_levels` | §3.2 | `1` = one partition per iteration. `>1` = nest it into a coarse-to-fine ladder, one level per denoising step; both training paths roll it out. The ladder itself is **Stage 2 only** (Stage 1 regresses `x0` in one forward per chunk, so it has no ladder to build), but Stage 1 **must still carry the same value**: it decides whether the shared level-0 mixture's coarse arm exists. | 1 |
 | `--flex_min_num_frame_per_block` | §3.2 | Block size the refinement stops at; `1` is fully causal at the leaves. | 1 |
-| `--flex_self_generated_context` | §3.3 | Feeds the K-Projection the model's own previous-step `x0` prediction as context, so it enters the autograd graph under prompt-only training. Costs a doubled attention sequence; from denoising step 1 on. Inert under `--use_kv_cache_training` - that path passes no `flex_state`, so the K-Projection is never invoked and stays at its identity init. | off |
-| `--num_frame_per_block` | — | The uniform fallback for the inherited (non-Flex) path; with Flex on it is also one of the three level-0 bands (`UNIFORM_BLOCK_PROB = 0.1` of iterations pin level 0 to its uniform partition), and it is what `log_validation` samples with at `--flex_pyramid_levels=1` - see the two bullets below. | 3 |
+| `--flex_self_generated_context` | §3.3 | Feeds the K-Projection the model's own previous-step `x0` prediction as context, so it enters the autograd graph under prompt-only training. Costs a doubled attention sequence; from denoising step 1 on. Inert under `--use_kv_cache_training`, but **not because the K-Projection is idle** - that path already commits its own clean context into the KV cache and needs no extra `clean_x` half; Π itself **is** invoked there. | off |
+| `--num_frame_per_block` | — | The uniform fallback for the inherited (non-Flex) path; with Flex on it is also reached inside level 0's four-arm mixture (the uniform arm is 30% of iterations and one of its three equal sub-arms pins level 0 to its uniform partition, so `[nfpb x k]` alone is 10%), and it is what `log_validation` samples with at `--flex_pyramid_levels=1`. | 3 |
 | `--independent_first_frame` | — | Prepend a length-1 chunk (`[1, N, N, ...]`); honoured by the partition sampler. | off |
 
-Seven interactions are worth stating explicitly:
-
-- **Level 0 is a three-band mixture decided by one draw.** `denoise_mode="pyramid"` at inference builds its level 0 as the whole clip in one chunk (fully bidirectional planning), so a partition drawn only from 2..10 would leave that step untrained - and the uniform `--num_frame_per_block` layout is just as unreachable by chance alone (at 21 latent frames `[3]*7` has probability `1/183708`, i.e. 0.11 expected hits in 20 000 iterations). `train_distill.py` therefore splits level 0 with a single uniform draw into three bands, neither constant a CLI flag: `COARSE_GLOBAL_PROB = 0.5` for the whole clip as one chunk, `UNIFORM_BLOCK_PROB = 0.1` pinned to the uniform `--num_frame_per_block` partition, and the remaining `0.4` for §3.1's random draw; all three refine into the same ladder. Measured over 200 000 iterations at 21 latent frames with `--flex_pyramid_levels=4`: whole clip `49.999%`, `[3]*7` `10.034%`, random `39.968%`, **4078** distinct level-0 layouts, and the inference ladder `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` at `49.999%`. At `--flex_pyramid_levels=1` the whole-clip band is empty and the split becomes `10% / 90%`, which keeps that mode as the paper describes it. Peak memory is set by the densest case (level 0 = whole clip), so budget for it.
-- The ladder is **clamped to the number of trained denoising steps**. A 4-level ladder under a 2-step schedule can never reach its last two levels, so `sample_flex_partitions` drops the unreachable tail rather than silently reporting it as trained. Keep `len(--denoising_step_indices_list) >= --flex_pyramid_levels`; the launcher's 4 indices and `--flex_pyramid_levels=4` are matched on purpose. The converse also wastes compute: more steps than levels leaves `install_flex_partition` clamped on its finest level for the extra steps.
-- **`--flow_euler_rollout` keeps the prediction in flow space between steps, so §3.3's context needs its own x0.** `--flex_self_generated_context` feeds the K-Projection the previous step's *clean* prediction; under an Euler rollout that variable is a velocity, so both rollout loops convert a detached copy for the context while the rollout state itself advances by Euler. The generator and the critic convert identically, otherwise the critic would score a rollout that is not the one being trained.
-- **`--randomize_step_indices` does not break the pyramid.** The jitter moves the *timesteps*, not the partitions, and `sample_flex_partitions` keys off `len(denoising_step_list)`, which the jitter preserves. Every level of the ladder still lines up with its denoising step.
-- **`log_validation` renders the layout training renders.** With `--flex_pyramid_levels > 1` it passes `chunk_spec=None` + `denoise_mode="pyramid"`, so its ladder is identical to `predict_t2v.py`'s and measured in `49.999%` of iterations. It used to pin `chunk_spec` to `--num_frame_per_block`, which nested `[[3]*7, [2,1]*7, [1]*21]` - a **3-level** ladder under a 4-step schedule (refinement hits the all-ones fixed point) that the trainer drew **0 times in 20 000** iterations, so those clips said nothing about the model being trained. At `--flex_pyramid_levels=1` the conclusion reverses: the whole-clip band is empty there and `[21]` is never drawn, so validation uses the uniform `--num_frame_per_block` partition - the one band that is both fixed across checkpoints and genuinely trained (`10.066%` of iterations, against ~`1.3%` for the most common random partition). On that 3-level ladder, now covered `10.034%` of the time: step 4 reuses the finest level `[1]*21` through `install_flex_partition`'s index clamp (`min(step_index, len(partitions) - 1)`), so nothing runs off the end.
-- §3.3 has **no launcher flag**: the config yaml chooses the variant through `transformer_additional_kwargs.flex_kproj_mode` (default `diag_rank1`; set it to `none` there to ablate Pi, `diag` for the cheaper head-only map). Its parameters are named `flex_kproj.proj_out.*` on purpose - that name is what lets `initialize_missing_parameters()` zero-fill them when a checkpoint predates §3.3, which is how identity init survives `low_cpu_mem_usage=True` - and they otherwise train like any other layer, following `--learning_rate`. The trainer forces them `requires_grad` whenever the built model carries them, so a narrowed `--trainable_modules` cannot silently freeze Pi at its identity. Under `--use_kv_cache_training` it is never invoked and stays at its identity init as well.
-- Every rank must train the same layout, otherwise the FlexAttention mask and the `num_frame_per_block` derived from it disagree across the SP/FSDP group. The trainer therefore passes the drawn partition through `broadcast_chunk_sizes` before using it, and `randomize_denoising_step_indices` broadcasts its jittered indices from rank 0 for the same reason.
-
 ### 3.4 Training Validation
 
 You can configure validation parameters to periodically generate test videos during training, allowing you to monitor training progress and model quality.
@@ -757,7 +745,7 @@ num_frame_per_block = 7          # clean-context commit granularity; None = one
 
 This `num_frame_per_block` is the same knob as Self-Forcing's: the clean context is committed in uniform blocks of that many frames, in temporal order. It differs in exactly two ways - the trailing chunk absorbs the remainder so divisibility is not required (the inherited rollout asserts it), and the transformer's block width becomes `max(it, edit span width)`.
 
-Any-order editing needs the full KV cache: both entries set `local_attn_size = -1`. A pyramid needs the same, plus `stochastic_sampling = True` — the buffered step between two levels *is* the schedule's own re-noising. The pipeline raises rather than silently degrading if either is violated.
+Any-order editing needs the full KV cache: both entries set `local_attn_size = -1`. Any per-step ladder (`denoise_mode="pyramid"` or `"full_then_blocks"`) needs the same, plus `stochastic_sampling = True` — the buffered step between two levels *is* the schedule's own re-noising. The pipeline raises rather than silently degrading if either is violated.
 
 ---
 
@@ -765,7 +753,7 @@ Any-order editing needs the full KV cache: both entries set `local_attn_size = -
 
 - **Flex-Forcing paper**: https://arxiv.org/abs/2607.03509
 - **Stage 1**: [README_TRAIN_ODE.md](./README_TRAIN_ODE.md)
-- **Shared chunking helpers**: `videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`, `sample_flexible_chunks`, `refine_partition`, `build_pyramid_partitions`, `validate_nested_partitions`, `broadcast_chunk_sizes`
+- **Shared chunking helpers**: `videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`, `sample_flexible_chunks`, `sample_flex_level0` (the level-0 four-arm mixture, drawn by **both** trainers), `refine_partition`, `build_pyramid_partitions`, `validate_nested_partitions`, `broadcast_chunk_sizes`
 - **Model / pipeline**: `videox_fun/models/wan_transformer3d_flex_forcing.py`, `videox_fun/pipeline/pipeline_wan_flex_forcing.py`
 - **Upstream stages**: [README_TRAIN_ODE.md](../wan2.1_self_forcing/README_TRAIN_ODE.md), [README_TRAIN_DISTILL.md](../wan2.1_self_forcing/README_TRAIN_DISTILL.md)
 - **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun

scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL_zh-CN.md

@@ -12,11 +12,9 @@
 > 2. **§3.2 金字塔时间步分块** —— 完成步 `t+1` 后，步 `t` 的划分是在**保留已有边界的前提下插入新边界**细化的。步 `t` 的结果先在整个原 chunk 上缓冲，随后每个子 chunk 在其所需 KV 可用后逐个自回归续写。先粗规划、后细刻画。
 > 3. **§3.3 K-Projection** —— 干净上下文的 key 与时间步 `t` 的含噪 query 处在不同的空间，因此用一个恒等初始化、以时间步为条件的投影把 cache 映射进当前含噪 latent 空间。它即时施加、从不改写 cache，并与生成器一起训练。
 >
-> 训练时每次迭代都抽取**全新的随机划分**（chunk 大小 2–10），这正是让同一套权重覆盖整个"因果↔双向"谱系的原因。论文 §4.2 随后把同一套机制复用为**任意序编辑**。
->
 > Flex-Forcing 流程分两个阶段：
 >
-> 1. **阶段 1——ODE 回归**（`train_ode.py`）：对教师模型的 ODE 轨迹做回归，让模型先成为一个合格的小步数生成器。随机分块在这一阶段就首次出现。
+> 1. **阶段 1——ODE 回归**（`train_ode.py`）：对教师模型的 ODE 轨迹做回归，让模型先成为一个合格的小步数生成器。
 > 2. **阶段 2——DMD 蒸馏**（`train_distill.py`，即本 README）：用 **Wan2.1-T2V-14B** 真实分数教师蒸馏出 **4 步**生成器，并加入金字塔时间步调度。
 >
 > 本 README 只覆盖**阶段 2**。阶段 1 见 [README_TRAIN_ODE_zh-CN.md](./README_TRAIN_ODE_zh-CN.md)。
@@ -210,7 +208,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
 
-launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒片段）、分辨率 **480×832**、**4 步**模型（`--denoising_step_indices_list 1000 750 500 250`，并由 `--randomize_step_indices` 每次迭代抖动）、`--flex_pyramid_levels=4` —— 在 21 个 latent 帧上正好复现论文的阶梯 `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` —— 划分抽取范围 **2..10**（另有半数的迭代以 §3.2 的整段粗粒度布局起步，见下）、`--flow_euler_rollout`、`diag_rank1` K-Projection、batch 8（8 卡 × `--gradient_accumulation_steps=1`）。
+launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒片段）、分辨率 **480×832**、**4 步**模型（`--denoising_step_indices_list 1000 750 500 250`，并由 `--randomize_step_indices` 每次迭代抖动）、`--flex_pyramid_levels=4` —— 在 21 个 latent 帧上正好复现论文的阶梯 `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` —— 随机档的划分抽取范围 **2..10**（该档只占 10% 的迭代；另有 60% 以 §3.2 的整段粗粒度布局起步，见下）、`--flow_euler_rollout`、`diag_rank1` K-Projection、batch 8（8 卡 × `--gradient_accumulation_steps=1`）。
 
 输出目录：`output_dir_wan2.1_flex_forcing_distill/`。
 
@@ -243,7 +241,7 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
 | `--denoising_step_indices_list` | 去噪步索引（DMD 核心参数）。经 `timesteps[train_sampling_steps - idx]` 映射，所以 `shift=5.0` 下 `1000 750 500 250` 变成 t = `[1000, 937.5, 833.33, 625]` —— 与 `stochastic_sampling_timesteps(4, shift)` 喂给 pipeline 的值逐位一致。 | `1000 750 500 250` |
 | `--randomize_step_indices` | 每次迭代对除首个索引以外的所有索引做对称抖动，幅度为相邻间隔的 `--index_jitter_ratio`。期望调度不变，单调性由构造保证，所以推理用的固定 t 始终落在每一步的训练包络内。 | 关 |
 | `--index_jitter_ratio` | 抖动预算（相邻间隔的比例）。在 `1000 750 500 250` 且取 0.3 时，20000 次抽取给出 idx `[675,825] / [425,575] / [175,325]`，即 t `[912.2,959.3] / [787.0,871.2] / [514.7,706.5]`；推理点 937.5 / 833.33 / 625 全部落在其中。 | 0.3 |
-| `--flow_euler_rollout` | 生成器自滚用 flow 空间里的确定性 Euler ODE 步推进（`x_next = x_t + (sigma_next - sigma_t) * v`，fp32），而不是转成 x0 再用新噪声重加噪。代数上这**就等于**用同一个隐含噪声重加噪 —— fp64 下实测相差 `8.9e-16`，而重采新噪声最大差 `6.4` —— 因此它消掉了 rollout 的随机性，与正常的 flow-matching 推理一致。最后一步仍然转 x0，因为 DMD 目标定义在 x0 上。critic 重滚生成器，跟随同一开关。 | 关 |
+| `--flow_euler_rollout` | 生成器自滚用 flow 空间里的确定性 Euler ODE 步推进（`x_next = x_t + (sigma_next - sigma_t) * v`，fp32），而不是转成 x0 再用新噪声重加噪。代数上这**就等于**用同一个隐含噪声重加噪 —— 在精确算术下与之完全相同，而重采新噪声是一次独立抽样 —— 因此它消掉了 rollout 的随机性，与正常的 flow-matching 推理一致。最后一步仍然转 x0，因为 DMD 目标定义在 x0 上。critic 重滚生成器，跟随同一开关。 | 关 |
 | `--output_dir` | 输出目录 | `output_dir_wan2.1_flex_forcing_distill` |
 | `--gradient_checkpointing` | 用算力换显存 | - |
 | `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
@@ -256,24 +254,14 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
 | 参数 | 论文 | 说明 | 默认值 |
 |-----------|-------|-------------|---------|
 | `--flex_forcing` | §3.1 | 实例化 `WanTransformer3DModel_FlexForcing` 并用 `WanFlexForcingPipeline` 校验。关闭 = 走继承的原路径，行为不变。 | 关 |
-| `--flex_chunk_min` | §3.1 | 每次迭代抽取的最小 chunk 大小。金字塔下 `1` 合法但不增加覆盖：2 帧 chunk 细化后本来就会产生 1 帧叶子，而它会让单帧 chunk 出现在片段中间。 | 2 |
-| `--flex_chunk_max` | §3.1 | 每次迭代抽取的最大 chunk 大小。设为与 `--flex_chunk_min` 相等即固定一种布局。**必须小于 latent 帧数**：“整段一个 chunk”这个布局已经由下面的混合机制覆盖，把这个值抬到帧数只是把抽取次数花在一个你已经有的布局上，反而稀释 §3.1（21 个 latent 帧时，`max=21` 只剩 3649 种不同划分，而 `max=10` 有 4882 种）。 | 10 |
-| `--flex_pyramid_levels` | §3.2 | `1` = 每次迭代一个划分。`>1` = 把它嵌套成由粗到细的阶梯，每个去噪步一层；两条训练路径都会展开。**仅阶段 2 有。** | 1 |
+| `--flex_chunk_min` | §5 | **random 档**（占 10% 迭代，与 `--flex_pyramid_levels` 无关）抽取的最小 chunk 大小；其余三档各自钉定自己的布局、平分剩下的 90%。金字塔下 `1` 合法但不增加覆盖：2 帧 chunk 细化后本来就会产生 1 帧叶子，而它会让单帧 chunk 出现在片段中间。 | 2 |
+| `--flex_chunk_max` | §5 | **random 档**抽取的最大 chunk 大小。设为与 `--flex_chunk_min` 相等即把该档钉成一种固定布局。**必须小于 latent 帧数**：“整段一个 chunk”本身就是独立的一档，把这个值抬到帧数只是把 random 档的抽取次数花在一个你已经有的布局上，反而稀释其他所有布局（21 个 latent 帧下穷举计数：`max=10` 有 6441 种不同划分、`max=21` 有 6765 种，但 `max=21` 会让单是 `[21]` 就占该档 1/19 的抽取，并把 `P([3]*7)` 从 1/183708 压到 1/1106560）。 | 10 |
+| `--flex_pyramid_levels` | §3.2 | `1` = 每次迭代一个划分。`>1` = 把它嵌套成由粗到细的阶梯，每个去噪步一层；两条训练路径都会展开。阶梯本身**仅阶段 2 有**（阶段 1 每 chunk 一次前向回归 x0，没有阶梯可搭），但阶段 1 **仍然要带同一个值**：它决定共享的 level-0 混合里 coarse 档是否存在。 | 1 |
 | `--flex_min_num_frame_per_block` | §3.2 | 细化终止的块大小（每块二分直到不超过它）；`1` 表示叶子级完全因果。 | 1 |
-| `--flex_self_generated_context` | §3.3 | 把模型自己上一步的 `x0` 预测喂给 K-Projection 当上下文，使其在纯 prompt 训练下也能进入 autograd 图。代价是注意力序列翻倍；从去噪步 1 起生效。在 `--use_kv_cache_training` 下不生效 —— 该路径不传 `flex_state`，K-Projection 不会被调用，停在恒等初始化。 | 关 |
-| `--num_frame_per_block` | — | 继承路径（未开 Flex）的均匀 fallback；开了 Flex 后它还是 level 0 三档混合中的一档（`UNIFORM_BLOCK_PROB = 0.1` 的迭代把 level 0 钉成它的均匀划分），且 `--flex_pyramid_levels=1` 时 `log_validation` 就用它采样 —— 见下面两条 bullet。 | 3 |
+| `--flex_self_generated_context` | §3.3 | 把模型自己上一步的 `x0` 预测作为上下文喂给 K-Projection，因此在纯 prompt 训练下它会进入自动求导图。代价是注意力序列翻倍；从去噪第 1 步起生效。在 `--use_kv_cache_training` 下它不起作用，但**并不是因为 K-Projection 空闲** —— 那条路径已经把自己的干净上下文写进 KV cache，不需要额外的 `clean_x` 一半；Π 本身在那里**是**被调用的。 | 关 |
+| `--num_frame_per_block` | — | 继承路径（未开 Flex）的均匀 fallback；开了 Flex 后它还是 level 0 四档混合中 uniform 那一档（该档占 30% 的迭代、内部再三等分，其一把 level 0 钉成它的均匀划分，故 `[nfpb x k]` 单独占 10%），且 `--flex_pyramid_levels=1` 时 `log_validation` 就用它采样。 | 3 |
 | `--independent_first_frame` | — | 在划分前插入一个长度为 1 的 chunk（`[1, N, N, ...]`）；划分采样器会遵守它。 | 关 |
 
-有七个交互值得明说：
-
-- **level 0 是三档混合，一次抽样定档。** 推理时 `denoise_mode="pyramid"` 的 level 0 是整段一个 chunk（全双向的全局规划），只从 2..10 抽划分会让那一步从未被训练；而 `--num_frame_per_block` 那个均匀布局光靠随机抽样又几乎碰不到（21 个 latent 帧下 `[3]*7` 的概率是 `1/183708`，20000 次迭代期望 0.11 次）。所以 `train_distill.py` 用一个均匀抽样把 level 0 分成三档，两个常量都故意不做成 CLI 参数：`COARSE_GLOBAL_PROB = 0.5` 整段一个 chunk、`UNIFORM_BLOCK_PROB = 0.1` 钉成 `--num_frame_per_block` 的均匀划分、剩下 `0.4` 是 §3.1 的随机划分；三档都细化成同一条阶梯。实测 200000 次（21 个 latent 帧、`--flex_pyramid_levels=4`）：整段 `49.999%`、`[3]*7` `10.034%`、随机 `39.968%`，共 **4078** 种不同的 level-0 布局，推理那条阶梯 `[[21], [11, 10], [6, 5, 5, 5], [3, 3, 3, 2, 3, 2, 3, 2]]` 占 `49.999%`。`--flex_pyramid_levels=1` 时整段那一档为空，比例变成 `10% / 90%`，与论文描述一致。峰值显存由最密的那种情形（level 0 = 整段）决定，请据此预留。
-- 阶梯会被**截断到训练的去噪步数**。2 步调度下配 4 级阶梯永远到不了最后两级，所以 `sample_flex_partitions` 会丢弃到不了的那部分，而不是默默把它们当作已训练。请保持 `len(--denoising_step_indices_list) >= --flex_pyramid_levels`；launcher 里的 4 个索引与 `--flex_pyramid_levels=4` 就是故意对齐的。反过来也浪费算力：步数多于级数时，多出来的那些步里 `install_flex_partition` 会被钉在最细的一级上。
-- **`--flow_euler_rollout` 会把预测留在 flow 空间，所以 §3.3 的上下文需要自己那份 x0。** `--flex_self_generated_context` 喂给 K-Projection 的是上一步的**干净**预测；在 Euler rollout 下那个变量是速度，所以两个 rollout 循环都会另外转换一份 detached 副本给上下文用，而 rollout 状态本身仍按 Euler 推进。生成器与 critic 的转换必须完全一致，否则 critic 评的就不是正在被训的那个 rollout。
-- **`--randomize_step_indices` 不会破坏金字塔。** 抖动动的是**时间步**而不是划分，而 `sample_flex_partitions` 只看 `len(denoising_step_list)`，抖动不会改变它。阶梯的每一级仍然对得上自己的去噪步。
-- **`log_validation` 渲染的就是训练渲染的那个布局。** `--flex_pyramid_levels > 1` 时它传 `chunk_spec=None` + `denoise_mode="pyramid"`，阶梯与 `predict_t2v.py` 的完全相同，实测占 `49.999%` 的迭代。它原本把 `chunk_spec` 钉成 `--num_frame_per_block`，那会嵌套出 `[[3]*7, [2,1]*7, [1]*21]` —— 在 4 步调度下只有 **3 级**（细化到全 1 就撞上固定点），而当时那条阶梯 trainer 在 20000 次迭代里抽出 **0 次**，所以那些校验视频根本说明不了正在训练的模型。`--flex_pyramid_levels=1` 时结论反过来：整段那一档为空、`[21]` 一次也抽不到，所以校验用 `--num_frame_per_block` 的均匀划分 —— 它是唯一既跨 checkpoint 固定、又真被训练的一档（`10.066%` 的迭代，而最高频的随机划分只有约 `1.3%`）。顺带说明那条 3 级阶梯：它现在有 `10.034%` 的覆盖率，第 4 步由 `install_flex_partition` 的索引 clamp（`min(step_index, len(partitions) - 1)`）复用最细一级 `[1]*21`，不会越界。
-- §3.3 **没有 launcher 参数**：变体由 config yaml 里的 `transformer_additional_kwargs.flex_kproj_mode` 决定（默认 `diag_rank1`；要消融就在那里改成 `none`，想更省就改 `diag`）。它的参数故意取名 `flex_kproj.proj_out.*` —— 正是这个名字让 `initialize_missing_parameters()` 在 checkpoint 早于 §3.3 时将其置零，从而在 `low_cpu_mem_usage=True` 下保住恒等初始化；除此之外它和其他层一样训，跟 `--learning_rate`。只要构建出的模型带 Pi，trainer 就强制其 `requires_grad`，所以收窄 `--trainable_modules` 也不会把 Pi 悄悄冻在恒等点。在 `--use_kv_cache_training` 下它同样不会被调用，始终停在恒等初始化。
-- 每张卡必须训练同一种布局，否则 FlexAttention 掩码与其派生的 `num_frame_per_block` 在 SP/FSDP 组之间会不一致。trainer 在使用抽取的划分前会先经 `broadcast_chunk_sizes` 广播对齐；`randomize_denoising_step_indices` 同样从 rank 0 广播它抖动后的索引，道理一样。
-
 ---
 
 ### 3.4 训练验证
@@ -746,15 +734,11 @@ num_frame_per_block = 7          # 干净上下文的提交粒度；None = 整
 
 这里的 `num_frame_per_block` 与 Self-Forcing 那个同名参数是一回事：干净上下文按这么多帧一块、依时间序提交。区别只有两点——尾块吸收余数所以不要求整除（继承的 rollout 会断言整除），且 transformer 上的块宽取 `max(它, 编辑区间宽度)`。
 
-任意序编辑需要完整 KV cache：两个入口都设 `local_attn_size = -1`。金字塔还需要 `stochastic_sampling = True`——两级之间的缓冲步本身就是调度自己的重加噪。任何一项不满足时 pipeline 会直接报错，而不是默默降级。
+任意序编辑需要完整 KV cache：两个入口都设 `local_attn_size = -1`。任何逐步换划分的 ladder（`denoise_mode="pyramid"` 或 `"full_then_blocks"`）受同一条约束，并额外要求 `stochastic_sampling = True`——两级之间的缓冲步本身就是调度自己的重加噪。任何一项不满足时 pipeline 会直接报错，而不是默默降级。
 
 ---
 
 ## 五、更多资源
 
 - **Flex-Forcing 论文**：https://arxiv.org/abs/2607.03509
-- **阶段 1**：[README_TRAIN_ODE_zh-CN.md](./README_TRAIN_ODE_zh-CN.md)
-- **共享分块 helper**：`videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`、`sample_flexible_chunks`、`refine_partition`、`build_pyramid_partitions`、`validate_nested_partitions`、`broadcast_chunk_sizes`
-- **模型 / pipeline**：`videox_fun/models/wan_transformer3d_flex_forcing.py`、`videox_fun/pipeline/pipeline_wan_flex_forcing.py`
-- **上游阶段**：[README_TRAIN_ODE.md](../wan2.1_self_forcing/README_TRAIN_ODE.md)、[README_TRAIN_DISTILL.md](../wan2.1_self_forcing/README_TRAIN_DISTILL.md)
 - **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun

scripts/wan2.1_flex_forcing/README_TRAIN_ODE.md

@@ -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
 ```
 

scripts/wan2.1_flex_forcing/README_TRAIN_ODE_zh-CN.md

@@ -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
 ```
 

scripts/wan2.1_flex_forcing/train_distill.py

@@ -77,16 +77,12 @@ from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
                                WanTransformer3DModel,
                                WanTransformer3DModel_FlexForcing,
                                WanTransformer3DModel_SelfForcing)
-from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
-                                 WanFlexForcingPipeline,
-                                 WanSelfForcingPipeline)
+from videox_fun.pipeline import (WanFlexForcingPipeline, WanI2VPipeline,
+                                 WanPipeline, WanSelfForcingPipeline)
 from videox_fun.utils.discrete_sampler import DiscreteSampling
-from videox_fun.utils.flex_chunking import (broadcast_chunk_sizes,
-                                            build_full_then_blocks_partitions,
-                                            build_pyramid_partitions,
+from videox_fun.utils.flex_chunking import (build_pyramid_partitions,
                                             chunk_boundaries,
-                                            sample_flexible_chunks,
-                                            uniform_chunks,
+                                            sample_flex_level0,
                                             validate_nested_partitions)
 from videox_fun.utils.tqdm_bar import PauseAwareTqdm
 from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
@@ -299,6 +295,7 @@ def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer
                     # the layout inference builds - otherwise the samples say
                     # nothing about the model being trained.
                     flex_kwargs = {}
+                    val_stochastic_sampling = not args.flow_euler_rollout
                     if args.flex_forcing:
                         if args.flex_pyramid_levels > 1:
                             # 3.2: exactly what predict_t2v.py passes. `chunk_spec=None`
@@ -311,14 +308,6 @@ def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer
                                 min_num_frame_per_block=args.flex_min_num_frame_per_block,
                             )
                         else:
-                            # No ladder in training, and with no pyramid the
-                            # whole-clip band is empty - `[21]` is never drawn at
-                            # `flex_chunk_max` below the frame count, so "bidir"
-                            # would validate out of distribution too. The uniform
-                            # block layout is the one band that is both fixed
-                            # across checkpoints and actually trained:
-                            # FLEX_ARM_PROB of iterations use it, against
-                            # ~1.3% for the most common random partition.
                             flex_kwargs = dict(
                                 chunk_spec=args.num_frame_per_block,
                                 denoise_mode="fixed",
@@ -336,6 +325,7 @@ def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer
                         num_frame_per_block = args.num_frame_per_block,
                         independent_first_frame = args.independent_first_frame,
                         context_noise = args.context_noise,
+                        stochastic_sampling = val_stochastic_sampling,
                         **flex_kwargs,
                     ).videos
                     os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
@@ -379,21 +369,20 @@ def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=Non
     return torch.clip(t.to(torch.int32), low, high - 1)
 
 
-# Each training iteration draws its level-0 layout from four equally likely arms
-# (FLEX_ARM_PROB = 25% of iterations each): the whole-clip planning chunk (3.2's
-# coarse step, `[[F], ...]` at inference), the launcher's own uniform
-# `num_frame_per_block`, the fixed "first step full, then block-major" ladder
-# (`denoise_mode="full_then_blocks"`), and a random 3.1 partition. Equal shares
-# keep the coarse, causal and every in-between layout covered while giving the
-# full_then_blocks ladder enough iterations to train the block-major-at-high-noise
-# trajectory the binary pyramid never reaches. Deliberately not a CLI flag - it is
-# a property of the schedule mixture, not a knob the launcher keeps in sync.
-FLEX_ARM_PROB = 0.25
+# The level-0 four-arm mixture (30% of iterations each for the whole-clip
+# planning chunk, the launcher's own uniform `num_frame_per_block` and the fixed
+# "first step full, then block-major" ladder, plus 10% for a random §3.3/§5
+# partition - see `FLEX_RANDOM_PROB`) now lives in
+# `videox_fun.utils.flex_chunking.sample_flex_level0`, because `train_ode.py`
+# draws from the SAME law: its checkpoint is this stage's `--ode_transformer_path`
+# initialisation, so a layout this stage spends most of its iterations on must
+# not be one stage 1 never regressed. What stays here is only the ladder built on
+# top of the draw.
 
 # How many steps of a launch report the partition they drew. Counted from this
 # process rather than from `global_step`, so a resumed run still gets its own
 # look at the layout. Two lines come out per step (the generator and the critic
-# each draw one), which is enough to watch the three arms show up in their
+# each draw one), which is enough to watch the arms show up in their
 # stated shares; after that they are only noise in a multi-day run.
 FLEX_LAYOUT_LOG_STEPS = 20
 
@@ -402,134 +391,74 @@ def sample_flex_partitions(args, num_frames, num_denoising_steps, torch_rng,
                            device, verbose=False):
     """Draw the Flex-Forcing partition ladder for one training iteration.
 
-    arXiv 2607.03509 3.1 asks for a *random* partition per rollout (chunk sizes
-    2..10) so a single model covers the whole causal-to-bidirectional spectrum,
-    and 3.2 then nests that partition into a coarse-to-fine ladder with one level
-    per denoising step. Returns ``None`` when Flex-Forcing is off, which leaves
-    the inherited uniform ``num_frame_per_block`` masks untouched.
+    arXiv 2607.03509 §3.3 asks for a *random* partition per rollout so a single
+    model covers the whole causal-to-bidirectional spectrum, §5 fixes that draw's
+    chunk sizes at 2..10, and §3.2 then nests the partition into a coarse-to-fine
+    ladder with one level per denoising step. Returns ``None`` when Flex-Forcing
+    is off, which leaves the inherited uniform ``num_frame_per_block`` masks
+    untouched.
 
-    The level-0 layout is a four-way mixture decided by a single uniform draw,
-    each arm taking ``FLEX_ARM_PROB`` (25%) of iterations: one chunk over the
-    whole clip (inference's coarse planning step), a single-level uniform block
-    grid, the fixed "first step full, then block-major" ladder
-    (`denoise_mode="full_then_blocks"`), and a random 2..10 partition. The coarse
-    and random arms are refined into the same binary pyramid, so a single set of
-    weights covers `denoise_mode="pyramid"` whether or not the caller pins
-    `chunk_spec`. The uniform and full_then_blocks arms bypass that refinement:
-    uniform emits a pure-Self-Forcing ladder in one of three layouts drawn from
+    The level-0 layout is a four-way mixture decided by a single uniform draw:
+    one chunk over the whole clip (inference's coarse planning step), a
+    single-level uniform block grid, the fixed "first step full, then
+    block-major" ladder (`denoise_mode="full_then_blocks"`), and a random 2..10
+    partition. The random arm takes ``FLEX_RANDOM_PROB`` (10%) and the other
+    three split what is left equally, i.e. 30% each. The coarse and random arms
+    are refined into the same binary pyramid, so a single set of weights covers
+    `denoise_mode="pyramid"` whether or not the caller pins `chunk_spec`. The
+    uniform and full_then_blocks arms bypass that refinement: uniform emits a
+    pure-Self-Forcing ladder in one of three layouts drawn from
     `num_frame_per_block` (`[3x7]`) and fully causal (`[1x21]`) -- `[3x7]` alone,
     `[1x21]` alone, or `[3x7]` planning followed by `[1x21]` for every remaining
-    step; full_then_blocks emits its two-level ladder whole, drawing the same
-    block width. With no pyramid the coarse arm is empty (a `[F]` level 0 would be
-    a degenerate single-chunk rollout), so the split is 25% uniform / 25%
-    full_then_blocks / 50% random.
+    step, 10% of iterations each; full_then_blocks emits its two-level ladder
+    whole, drawing the same block width. With no pyramid the coarse arm is empty
+    (a `[F]` level 0 would be a degenerate single-chunk rollout), so the two
+    survivors widen and the split becomes 45% uniform / 45% full_then_blocks /
+    10% random.
+
+    The draw itself is :func:`videox_fun.utils.flex_chunking.sample_flex_level0`,
+    shared with ``train_ode.py`` so that stage 1 - whose checkpoint this stage
+    loads via ``--ode_transformer_path`` - regresses the same level-0 law it is
+    then distilled on. What this function adds is only the *ladder*: how the drawn
+    layout nests into one level per denoising step.
     """
     if not args.flex_forcing:
         return None
-    # The arm selector must be IDENTICAL on every rank, not just its broadcast
-    # count. Only the coarse/random arms funnel through the shared refine tail
-    # below; the uniform and full_then_blocks arms each set the ladder directly
-    # from a different tensor, so a per-rank `u` would leave ranks in different
-    # arms building different ladders
-    # -> different walk forward counts -> FSDP all-gather desync (NCCL hang).
-    # Broadcast rank 0's draw so all ranks take the same arm; the arm's existing
-    # single broadcast then reconciles its within-arm randomness (the random base,
-    # the block width) to rank 0, making every ladder byte-identical.
-    u = torch.rand((), generator=torch_rng, device=device)
-    if dist.is_available() and dist.is_initialized():
-        dist.broadcast(u.reshape(1), src=0)
-    u = u.item()
     pyramid = args.flex_pyramid_levels > 1
-    # Four equally likely arms; the coarse (whole-clip planning) one only when a
-    # pyramid is on, else its 25% folds into the random arm.
-    coarse_hi = FLEX_ARM_PROB if pyramid else 0.0
-    uniform_hi = coarse_hi + FLEX_ARM_PROB
-    ftb_hi = uniform_hi + FLEX_ARM_PROB
-    ladder = None
-    if u < coarse_hi:
-        # Coarse end of 3.1: reuse the ladder builder's own level-0 rule so the
-        # `independent_first_frame` handling cannot drift from inference's.
-        base = build_pyramid_partitions(
-            num_frames, num_levels=1, base_chunks=None,
-            independent_first_frame=args.independent_first_frame)[0]
-        arm = "whole clip, the 3.2 coarse planning layout"
-    elif u < uniform_hi:
-        # Pure Self-Forcing (no binary-pyramid refinement), three layouts drawn
-        # from `num_frame_per_block` (-> [3x7]) and fully causal (1 frame/block,
-        # -> [1x21]): [3x7] alone, [1x21] alone, and [3x7] planning followed by
-        # [1x21] for every remaining step (the coarse-block -> per-frame descent
-        # the pyramid never samples because it halves through [2,1x14] first).
-        # The three change the ladder DEPTH (hence the walk's forward count), so
-        # the choice must be rank-consistent; it is free here because `u` was
-        # already broadcast to rank 0, so slicing its position within the uniform
-        # band gives every rank the same sub-arm with no extra collective. Like
-        # every arm this issues EXACTLY ONE broadcast (the level-0 layout; the
-        # trailing [1x21] levels are deterministic from `num_frames`). Setting
-        # `ladder` here skips the refine tail. `uniform_chunks`, not
-        # `normalize_chunk_spec`: the latter encodes `independent_first_frame` as
-        # a leading 1 and would silently disagree with the bands around it.
-        band = (u - coarse_hi) / FLEX_ARM_PROB          # [0, 1) in the uniform arm
-        sub = min(int(band * 3), 2)                     # 0 / 1 / 2, equiprobable
-        nfpb = max(1, int(args.num_frame_per_block))
-        iff = args.independent_first_frame
-        causal = uniform_chunks(num_frames, 1, independent_first_frame=iff)
-        if sub == 0:
-            base = broadcast_chunk_sizes(causal, device=device)
-            ladder = [base]
-            arm = "uniform 1-frame blocks, single level (fully causal)"
-        elif sub == 1:
-            base = broadcast_chunk_sizes(
-                uniform_chunks(num_frames, nfpb, independent_first_frame=iff),
-                device=device)
-            ladder = [base]
-            arm = f"uniform {nfpb}-frame blocks, single level"
-        else:
-            base = broadcast_chunk_sizes(
-                uniform_chunks(num_frames, nfpb, independent_first_frame=iff),
-                device=device)
-            ladder = ([base] + [list(causal)
-                                for _ in range(max(0, num_denoising_steps - 1))])
-            arm = (f"uniform {nfpb}-frame plan -> fully causal [1x" f"{num_frames}] "
-                   "for the remaining steps")
-    elif u < ftb_hi:
-        # "First step full, every later step block-major": a fixed two-level
-        # ladder that does NOT go through the binary pyramid refinement, so the
-        # block-major level is reached at high noise (the 2nd step) instead of as
-        # a deep refinement - the one trajectory the pyramid arms never sample.
-        # The block-major width is drawn between fully causal (1 frame / block)
-        # and the launcher's `num_frame_per_block`, so both the tight-AR and the
-        # coarse-block refinement of the whole-clip plan get trained. The block
-        # width is a per-rank draw, so it has to be reconciled with EXACTLY ONE
-        # broadcast - the same collective count every other arm issues (they each
-        # broadcast their single base) - or the ranks desync and NCCL deadlocks.
-        # Level 0 is just [F], identical on every rank and needing no sync, so we
-        # broadcast only the block-major level; rank 0's draw wins and every rank
-        # rebuilds the same two-level ladder locally.
-        block_choices = sorted({1, max(1, int(args.num_frame_per_block))})
-        pick = int(torch.rand((), generator=torch_rng, device=device).item()
-                   * len(block_choices))
-        block = block_choices[min(pick, len(block_choices) - 1)]
-        block = min(block, int(num_frames))
-        ftb = build_full_then_blocks_partitions(
-            num_frames, block,
-            independent_first_frame=args.independent_first_frame)
-        ftb[1] = broadcast_chunk_sizes(ftb[1], device=device)
-        ladder = ftb[:num_denoising_steps]
-        arm = (f"full clip -> uniform {max(ftb[1])}-frame blocks, "
-               "the full_then_blocks ladder")
+    draw = sample_flex_level0(
+        num_frames,
+        num_frame_per_block=args.num_frame_per_block,
+        chunk_min=args.flex_chunk_min,
+        chunk_max=args.flex_chunk_max,
+        pyramid=pyramid,
+        generator=torch_rng,
+        device=device,
+        independent_first_frame=args.independent_first_frame)
+    base = draw["sizes"]
+    arm = draw["arm_text"]
+    if draw["arm"] in ("uniform_causal", "uniform_blocks"):
+        # Pure Self-Forcing grids: a single level, and they deliberately bypass
+        # the binary refinement below - refining `[3x7]` would halve it through
+        # `[2, 1x14]`, which is the pyramid's own trajectory, not this arm's.
+        ladder = [base]
+    elif draw["arm"] == "uniform_plan_then_causal":
+        # `[3x7]` planning, then fully causal for every remaining step: the
+        # coarse-block -> per-frame descent the pyramid never samples. The
+        # trailing `[1x21]` levels are deterministic from `num_frames`, so the
+        # arm's single broadcast (already issued) covered everything random.
+        ladder = ([base] + [list(draw["causal"])
+                            for _ in range(max(0, num_denoising_steps - 1))])
+    elif draw["arm"] == "full_then_blocks":
+        # Two levels, taken whole. The walk clamps to the finest level for every
+        # later step, so a ladder shorter than the denoising schedule is the
+        # intended shape here, not wasted steps.
+        ladder = [list(base), list(draw["block_level"])][:num_denoising_steps]
     else:
-        base = sample_flexible_chunks(
-            num_frames, min_chunk=args.flex_chunk_min, max_chunk=args.flex_chunk_max,
-            generator=torch_rng, device=device,
-            independent_first_frame=args.independent_first_frame)
-        arm = f"random {args.flex_chunk_min}..{args.flex_chunk_max}-frame blocks, the 3.1 spectrum"
-    if ladder is None:
-        # Every rank has to train the same layout: the FlexAttention mask, and the
+        # `coarse` / `random`: the two arms the binary pyramid refines. Every rank
+        # has to train the same layout - the FlexAttention mask, and the
         # `num_frame_per_block` derived from it, must agree across the SP/FSDP
-        # group. Only the coarse and random arms reach here; the uniform and
-        # full_then_blocks arms above already set `ladder` (each with its single
-        # broadcast), so they skip this refine path.
-        base = broadcast_chunk_sizes(base, device=device)
+        # group - which `sample_flex_level0` already guaranteed by broadcasting
+        # rank 0's draw.
         ladder = [base]
         if pyramid:
             ladder = build_pyramid_partitions(
@@ -601,16 +530,24 @@ def flex_self_context_enabled(args, clean_x, final_step_index):
     """Whether 3.3's K-Projection gets a self-generated clean context this round.
 
     This is a block-mask-path concern only, and this flag is inert under
-    ``--use_kv_cache_training``. The KV-cache rollout passes no ``flex_state``,
-    so ``_flex_project`` short-circuits and :math:`\\Pi` stays at its identity
-    initialisation there - which matches KV inference, which never supplies a
-    timestep either, so the two sides agree rather than one of them drifting.
-    The block-mask path is the one that exercises :math:`\\Pi`, and it builds no
-    KV cache, so with prompt-only data - ``clean_latents is None``, which the
-    prompt-only branch sets unconditionally - no clean context exists anywhere
-    in the sequence, :math:`\\Pi_{t\\leftarrow 0}` has nothing to project and
-    never enters the autograd graph: its parameters sit in an optimizer group
-    and still stay at the identity initialisation for the entire run.
+    ``--use_kv_cache_training`` - but *not* because :math:`\\Pi` is idle there.
+    :class:`WanTransformer3DModel_FlexForcing` arms :math:`\\Pi` by itself
+    whenever a projection was built and the forward has a clean context to read
+    (``kv_cache is not None or clean_x is not None``), taking the noise level
+    from ``t``; a caller that passes no ``flex_state`` still projects. So
+    :math:`\\Pi` **is** invoked on the KV-cache path, in training and in
+    inference alike, and the exit step that carries gradient gives it real
+    gradient. All this flag does *not* do there is supply a ``clean_x`` half:
+    the KV rollout already commits its own clean context into the cache ("feed
+    denoised_pred directly"), so it has no use for one and the switch has no
+    lever to pull.
+    On the block-mask path the switch is what decides whether :math:`\\Pi` has
+    anything to project at all: that path builds no KV cache, so with
+    prompt-only data - ``clean_latents is None``, which the prompt-only branch
+    sets unconditionally - no clean context exists anywhere in the sequence,
+    :math:`\\Pi_{t\\leftarrow 0}` has nothing to project and never enters the
+    autograd graph: its parameters sit in an optimizer group and still stay at
+    the identity initialisation for the entire run.
 
     The denoising loop already re-rolls the model under ``no_grad`` and converts
     every non-final step to an x0 prediction, so the previous step's prediction
@@ -1137,21 +1074,23 @@ def parse_args():
         "--flex_chunk_min",
         type=int,
         default=2,
-        help="Smallest chunk drawn when sampling a partition (paper 3.1 uses 2). "
-             "1 is allowed but adds no coverage: refining a 2-frame chunk already "
-             "yields 1-frame leaves."
+        help="Smallest chunk drawn by the level-0 mixture's *random* arm (paper "
+             "§5 uses 2). 1 is allowed but adds no coverage: refining a 2-frame "
+             "chunk already yields 1-frame leaves."
     )
     parser.add_argument(
         "--flex_chunk_max",
         type=int,
         default=10,
-        help="Largest chunk drawn when sampling a partition (paper 3.1 uses 10; "
-             "setting it equal to --flex_chunk_min pins one fixed layout). Keep it "
+        help="Largest chunk drawn by the level-0 mixture's *random* arm (paper §5 "
+             "uses 10; setting it equal to --flex_chunk_min pins that arm to one "
+             "fixed layout). Keep it "
              "below the latent frame count: the whole-clip layout is already "
-             "covered by the pyramid mixture, so raising max to the frame count "
-             "only spends draws on a layout you already have and thins out 3.1 - "
-             "at 21 latent frames max=21 leaves 3649 distinct partitions versus "
-             "4882 at max=10."
+             "covered by the coarse arm, so raising max to the frame count spends "
+             "draws on a layout you already have and thins every other one - at 21 "
+             "latent frames exhaustive enumeration gives 6441 distinct partitions "
+             "at max=10 and 6765 at max=21, but max=21 makes [21] alone 1/19 of "
+             "draws and cuts P([3]*7) from 1/183708 to 1/1106560."
     )
     parser.add_argument(
         "--flex_pyramid_levels",
@@ -1178,10 +1117,14 @@ def parse_args():
              "block-mask path trains on prompts alone: it builds no KV cache, so "
              "without this there is no clean context at all and the K-Projection "
              "never enters the autograd graph. Doubles the attention sequence. "
-             "Inert under --use_kv_cache_training: that path passes no flex_state, "
-             "so the K-Projection is never invoked on either side and stays at its "
-             "identity init. Also ignored when real --use_teacher_forcing data is "
-             "present, and with --flex_forcing off."
+             "Inert under --use_kv_cache_training for a different reason than it "
+             "looks like: that path already commits its own clean context into the "
+             "KV cache, so it needs no clean_x half. The K-Projection itself IS "
+             "active there - the Flex backbone arms it automatically whenever a "
+             "projection was built and kv_cache or clean_x is present - so this "
+             "flag is not evidence that the projection sits untrained under "
+             "--use_kv_cache_training. Also ignored when real --use_teacher_forcing "
+             "data is present, and with --flex_forcing off."
     )
     parser.add_argument(
         "--independent_first_frame",

scripts/wan2.1_flex_forcing/train_ode.py

@@ -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(

scripts/wan2.1_flex_forcing/train_ode.sh

@@ -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

scripts/wan2.2/README_TRAIN_REWARD.md

@@ -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>

scripts/wan2.2/README_TRAIN_REWARD_zh-CN.md

@@ -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>

scripts/wan2.2/train_reward_lora.sh

@@ -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

scripts/wan2.2_fun/README_TRAIN_REWARD.md

@@ -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>

scripts/wan2.2_fun/README_TRAIN_REWARD_zh-CN.md

@@ -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>

scripts/wan2.2_fun/train_reward_lora.sh

@@ -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

videox_fun/dist/__init__.py

@@ -25,5 +25,6 @@ from .minimax_h3_xfuser import MiniMaxH3MultiGPUsAttnProcessor
 from .qwen_xfuser import QwenImageMultiGPUsAttnProcessor2_0
 from .qwenimage21_xfuser import QwenImage21MultiGPUsAttnProcessor
 from .wan_xfuser import (usp_attn_forward, usp_attn_s2v_forward,
-                         usp_attn_self_forcing_forward)
+                         usp_attn_self_forcing_forward,
+                         usp_attn_sparkwan_forward)
 from .z_image_xfuser import ZMultiGPUsSingleStreamAttnProcessor

videox_fun/dist/wan_xfuser.py

@@ -6,7 +6,7 @@ import torch.cuda.amp as amp
 from .fuser import (get_sequence_parallel_rank,
                     get_sequence_parallel_world_size, get_sp_group,
                     init_distributed_environment, initialize_model_parallel,
-                    xFuserLongContextAttention)
+                    ulysses_all_to_all, xFuserLongContextAttention)
 
 
 def pad_freqs(original_tensor, target_len):
@@ -162,6 +162,111 @@ def usp_attn_forward(self,
     x = self.o(x)
     return x
 
+def usp_attn_sparkwan_forward(self,
+                     x,
+                     seq_lens,
+                     grid_sizes,
+                     freqs,
+                     dtype=torch.bfloat16,
+                     t=0):
+    """Sequence-parallel self-attention that KEEPS the RoLa sparse kernel (``self.attn_op``).
+
+    The stock ``usp_attn_forward`` binds ``xFuserLongContextAttention``'s dense full-softmax
+    and never reaches ``self.attn_op``, so using it on a ``SparkWanSelfAttention`` would
+    silently drop the trained RoLa params (``attn_op.{proj_q,proj_k,gate_proj,gate_bias}``)
+    and shift the numbers. This variant uses Ulysses head-parallelism instead: an all-to-all
+    turns each rank's sequence shard ``[B, S/P, H, D]`` into a head shard ``[B, S, H/P, D]``,
+    the RoLa branches (block-sparse exact softmax + rank-truncated low-rank linear + token
+    gate) run over the FULL sequence on a subset of heads -- every branch is head-independent
+    and topk is a per-query key-block ratio, so this reproduces the single-GPU result --
+    then a second all-to-all returns to the sequence-shard layout.
+
+    Two RoLa-specific points under head-parallelism:
+      * ``gate_bias`` is per-head ``[1, H, 1, 1]``; each rank holds the contiguous head block
+        ``[sp_rank * H/P, (sp_rank + 1) * H/P)`` produced by the ulysses all-to-all, so the
+        gate slices ``gate_bias`` to that range.
+      * the low-rank branch's rank RoPE runs on the FULL (post-all-to-all) sequence with the
+        non-SP ``rope_apply_lowrank``, so positions are complete and need no sharding.
+
+    Requires a pure-Ulysses split: ``num_heads`` must divide the SP degree and ``ring_degree``
+    should be 1 (the SLA kernel has no ring-attention path). ``topk >= 1.0`` (the dense
+    fallback) defers to ``usp_attn_forward``.
+    """
+    # Dense student (topk >= 1.0) or a plain Wan self-attention: keep the stock SP path.
+    if getattr(self, "topk", 1.0) >= 1.0 or not hasattr(self, "attn_op"):
+        return usp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=dtype, t=t)
+
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    half_dtypes = (torch.float16, torch.bfloat16)
+
+    def half(t_):
+        return t_ if t_.dtype in half_dtypes else t_.to(dtype)
+
+    # query, key, value function
+    def qkv_fn(x):
+        q = self.norm_q(self.q(x)).view(b, s, n, d)
+        k = self.norm_k(self.k(x)).view(b, s, n, d)
+        v = self.v(x).view(b, s, n, d)
+        return q, k, v
+
+    q, k, v = qkv_fn(x)
+    # PRE-RoPE qk-normed q/k feed the RoLa low-rank branch (RoPE is applied on the low-rank
+    # bottleneck, outside the non-linear feature map). rope_apply_qk returns new tensors, so
+    # q_pre/k_pre stay valid after the rebind below.
+    q_pre, k_pre = q, k
+    q, k = rope_apply_qk(q, k, grid_sizes, freqs)
+
+    sp_size = get_sequence_parallel_world_size()
+    sp_rank = get_sequence_parallel_rank()
+
+    # PureSLA ablation: fused block-sparse + feature-map linear (proj_l), BSHD in/out.
+    if getattr(self, "attn_variant", "rola") == "pure_sla":
+        q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
+        k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
+        v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
+        x = self.attn_op(q, k, v)
+        x = ulysses_all_to_all(x, scatter_dim=1, gather_dim=2)
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+    # ---- RoLa: seq-shard -> head-shard all-to-all for every branch input ----
+    # post-RoPE q/k (sparse branch), pre-RoPE q/k (low-rank branch), v, and the raw hidden
+    # states x (token gate) are all gathered to [B, S_full, H/P, D].
+    q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
+    k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
+    v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
+    q_pre = ulysses_all_to_all(half(q_pre), scatter_dim=2, gather_dim=1)
+    k_pre = ulysses_all_to_all(half(k_pre), scatter_dim=2, gather_dim=1)
+    x_full = ulysses_all_to_all(half(x).view(b, s, n, d), scatter_dim=2, gather_dim=1)
+
+    # BSHD -> BHLD over the head shard.
+    q = q.transpose(1, 2)
+    k = k.transpose(1, 2)
+    v_bhld = v.transpose(1, 2)
+    q_pre = q_pre.transpose(1, 2)
+    k_pre = k_pre.transpose(1, 2)
+    x_bhld = x_full.transpose(1, 2)
+
+    o_s = self.attn_op.sparse(q, k, v_bhld)
+    o_l = self.attn_op.lowrank(q_pre, k_pre, v_bhld, grid_sizes, freqs)
+    # gate_bias is per-head; slice to this rank's contiguous head block from the all-to-all.
+    h_per = n // sp_size
+    head_slice = slice(sp_rank * h_per, (sp_rank + 1) * h_per)
+    gate = self.attn_op.gate(x_bhld, head_slice=head_slice)
+    if o_s is not None:
+        scale = (o_s.pow(2).mean(dim=-1, keepdim=True) + 1e-5).sqrt()
+        out = ((o_s.float() / scale.float() + gate * o_l.float()) * scale.float()).to(o_l.dtype)
+    else:
+        out = gate * o_l
+
+    # BHLD -> BSHD, head-shard -> seq-shard, back to this rank's sequence chunk.
+    out = out.transpose(1, 2)
+    out = ulysses_all_to_all(out, scatter_dim=1, gather_dim=2)
+    out = out.flatten(2)
+    out = self.o(out)
+    return out
+
 @amp.autocast(enabled=False)
 @torch.compiler.disable()
 def s2v_rope_apply(x, grid_sizes, freqs):

videox_fun/models/__init__.py

@@ -90,6 +90,9 @@ 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 .wan_audio_encoder import WanAudioEncoder
 from .wan_image_encoder import CLIPModel
@@ -108,5 +111,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

videox_fun/models/attention_utils.py

@@ -424,4 +424,130 @@ 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)

videox_fun/models/spark_transformer3d.py

@@ -0,0 +1,402 @@
+# 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, 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.
+
+    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; those pairs are budgeted across the temporal /
+    height / width axes exactly like ``rope_apply`` (``[c - 2*(c//3), c//3, c//3]``) and
+    each axis takes the FIRST ``n`` columns of its own block in the pre-trained table, so
+    the native t->h->w coordinate ordering and the per-axis band identity are preserved.
+
+    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
+    c_r = rank // 2                       # rotary pairs the bottleneck can carry
+    # Native ordering: split the pre-trained table into its t/h/w blocks, then keep the
+    # first n columns of each block for the rank budget (n mirrors rope_apply's split).
+    t_full, h_full, w_full = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+    t_n, h_n, w_n = c_r - 2 * (c_r // 3), c_r // 3, c_r // 3
+    freqs_r = torch.cat(
+        [t_full[:, :t_n], h_full[:, :h_n], w_full[:, :w_n]], dim=1)  # [M, c_r]
+    # rope_apply expects [B, L, n_channels, dim]; BHLD -> BSHD and back.
+    out = rope_apply(x.transpose(1, 2), grid_sizes, freqs_r)
+    return out.transpose(1, 2)
+
+
+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]

videox_fun/models/wan_transformer3d_flex_forcing.py

@@ -703,6 +703,24 @@ class WanTransformer3DModel_FlexForcing(WanTransformer3DModel_SelfForcing):
         optionally clipped to ``local_attn_size`` frames.
         """
         chunk_sizes = self._require_exact_partition(chunk_sizes, num_frames)
+        # Read the value the mask below actually uses - not the config mirror, so
+        # the check cannot pass against a window the builder then ignores.
+        local_attn_size = self.local_attn_size
+        # One chunk is one bidirectional unit, so every frame inside it has to see
+        # every other. The clipped window below is anchored at the chunk *end*
+        # (`kv_idx >= ends[q_idx] - local_attn_size * frame_seqlen`), so a chunk
+        # wider than the window quietly loses its own head-to-tail connectivity:
+        # the layout would be trained as if it were partially causal while every
+        # caller describes it as bidirectional. Reject it loudly instead. The same
+        # rule is enforced on the KV-cache read side in the pipelines.
+        if local_attn_size != -1 and max(chunk_sizes) > local_attn_size:
+            raise ValueError(
+                f"A chunk of {max(chunk_sizes)} latent frames exceeds "
+                f"`local_attn_size={local_attn_size}`: the window is anchored at "
+                f"the chunk end, so that chunk's first frames could not attend to "
+                f"its last ones and its intra-chunk bidirectionality (§3.1) would "
+                f"be lost. Use `local_attn_size=-1` or keep every chunk within the "
+                f"window.")
         total_length = num_frames * frame_seqlen
 
         # Padded query rows keep ends=0 exactly like the inherited uniform
@@ -712,8 +730,6 @@ class WanTransformer3DModel_FlexForcing(WanTransformer3DModel_SelfForcing):
             device=device, dtype=torch.long)
         ends[:total_length] = chunk_ends_tensor(chunk_sizes, frame_seqlen, device=device)
 
-        local_attn_size = self.local_attn_size
-
         def attention_mask(b, h, q_idx, kv_idx):
             if local_attn_size == -1:
                 # Global chunk-wise causal: bidirectional inside a chunk, all
@@ -723,6 +739,17 @@ class WanTransformer3DModel_FlexForcing(WanTransformer3DModel_SelfForcing):
             return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx)
 
         self._build_block_mask(attention_mask, total_length, device)
+        # Keep the parent's cached-mask descriptor in sync with the mask just
+        # installed: the guard reuses `self.block_mask` whenever
+        # `_block_mask_expected_len`/`_block_mask_is_teacher_forcing` match the
+        # next call, so a builder that replaced the mask without updating them
+        # would leave them describing a mask that no longer exists. The parent
+        # forward re-records the same two values right after dispatching here,
+        # so this is the invariant being stated at its source rather than a
+        # rebuild saving - `set_flex_chunk_sizes` is what forces one rebuild per
+        # new partition by clearing the length.
+        self._block_mask_expected_len = total_length
+        self._block_mask_is_teacher_forcing = False
 
         # Consumers that still expect the scalar Self-Forcing knobs (Forcing-KV
         # AR stride / rolling-cache budget) get the conservative largest chunk.
@@ -777,6 +804,11 @@ class WanTransformer3DModel_FlexForcing(WanTransformer3DModel_SelfForcing):
             return eye_mask | clean_mask | noise_mask
 
         self._build_block_mask(attention_mask, total_length, device)
+        # Same descriptor sync as the non-teacher-forcing builder above; here
+        # `total_length` already counts clean + noisy, which is exactly what the
+        # parent guard derives when `clean_x is not None`.
+        self._block_mask_expected_len = total_length
+        self._block_mask_is_teacher_forcing = True
 
         self.num_frame_per_block = max(chunk_sizes)
 

videox_fun/models/wan_vae3_8_1steps.py

@@ -0,0 +1,945 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+from typing import Tuple, Union
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.autoencoders.vae import (DecoderOutput,
+                                               DiagonalGaussianDistribution)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils.accelerate_utils import apply_forward_hook
+from einops import rearrange
+from torch.utils.checkpoint import checkpoint as torch_checkpoint
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x):
+        # Single-pass causal conv: the full temporal sequence is available at
+        # once, so the past is provided purely by the front temporal pad in
+        # ``self._padding`` (no feat_cache threading).
+        x = F.pad(x, self._padding)
+        x = super().forward(x)
+        return x
+
+
+class RMS_norm(nn.Module):
+
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        return (F.normalize(x, dim=(1 if self.channel_first else -1)) *
+                self.scale * self.gamma + self.bias)
+
+
+class Upsample(nn.Upsample):
+
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in (
+            "none",
+            "upsample2d",
+            "upsample3d",
+            "downsample2d",
+            "downsample3d",
+        )
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+                # nn.Conv2d(dim, dim//2, 3, padding=1)
+            )
+            self.time_conv = CausalConv3d(
+                dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(
+                dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x):
+        b, c, t, h, w = x.size()
+
+        if self.mode == 'upsample3d':
+            # Causal temporal upsample without cache: the first frame is kept
+            # as-is and only the remaining frames are temporally doubled. This
+            # reproduces the chunked path (frame 0 emitted with zero left
+            # context, later frames with 2 frames of context).
+            first_frame = x[:, :, :1, :, :]
+            rest_frames = x[:, :, 1:, :, :] if t > 1 else None
+            if rest_frames is not None:
+                rest_frames = self.time_conv(rest_frames)
+                rest_frames = rest_frames.reshape(b, 2, c, t - 1, h, w)
+                rest_frames = torch.stack((rest_frames[:, 0], rest_frames[:, 1]),
+                                          3)
+                rest_frames = rest_frames.reshape(b, c, (t - 1) * 2, h, w)
+                x = torch.cat([first_frame, rest_frames], dim=2)
+            else:
+                x = first_frame
+
+        t = x.shape[2]
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.resample(x)
+        x = rearrange(x, '(b t) c h w -> b c t h w', t=t)
+
+        if self.mode == 'downsample3d':
+            # Causal temporal downsample without cache: first frame passes
+            # through unchanged and the stride-2 causal conv runs over the full
+            # sequence (its own windows supply the previous-frame left context),
+            # exactly reproducing the chunked cache-threaded path.
+            if x.shape[2] > 1:
+                first_frame = x[:, :, :1, :, :]
+                rest_frames = self.time_conv(x)
+                x = torch.cat([first_frame, rest_frames], dim=2)
+
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix  # * 0.5
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = (
+            CausalConv3d(in_dim, out_dim, 1)
+            if in_dim != out_dim else nn.Identity())
+
+    def forward(self, x):
+        h = self.shortcut(x)
+        x = self.residual(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = (
+            self.to_qkv(x).reshape(b * t, 1, c * 3,
+                                   -1).permute(0, 1, 3,
+                                               2).contiguous().chunk(3, dim=-1))
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
+        return x + identity
+
+
+def patchify(x, patch_size):
+    if patch_size == 1:
+        return x
+    if x.dim() == 4:
+        x = rearrange(
+            x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(
+            x,
+            "b c f (h q) (w r) -> b (c r q) f h w",
+            q=patch_size,
+            r=patch_size,
+        )
+    else:
+        raise ValueError(f"Invalid input shape: {x.shape}")
+
+    return x
+
+
+def unpatchify(x, patch_size):
+    if patch_size == 1:
+        return x
+
+    if x.dim() == 4:
+        x = rearrange(
+            x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(
+            x,
+            "b (c r q) f h w -> b c f (h q) (w r)",
+            q=patch_size,
+            r=patch_size,
+        )
+    return x
+
+
+class AvgDown3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        factor_t,
+        factor_s=1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+
+        assert in_channels * self.factor % out_channels == 0
+        self.group_size = in_channels * self.factor // out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
+        pad = (0, 0, 0, 0, pad_t, 0)
+        x = F.pad(x, pad)
+        B, C, T, H, W = x.shape
+        x = x.view(
+            B,
+            C,
+            T // self.factor_t,
+            self.factor_t,
+            H // self.factor_s,
+            self.factor_s,
+            W // self.factor_s,
+            self.factor_s,
+        )
+        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
+        x = x.view(
+            B,
+            C * self.factor,
+            T // self.factor_t,
+            H // self.factor_s,
+            W // self.factor_s,
+        )
+        x = x.view(
+            B,
+            self.out_channels,
+            self.group_size,
+            T // self.factor_t,
+            H // self.factor_s,
+            W // self.factor_s,
+        )
+        x = x.mean(dim=2)
+        return x
+
+
+class DupUp3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        factor_t,
+        factor_s=1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+
+        assert out_channels * self.factor % in_channels == 0
+        self.repeats = out_channels * self.factor // in_channels
+
+    def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
+        x = x.repeat_interleave(self.repeats, dim=1)
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            self.factor_t,
+            self.factor_s,
+            self.factor_s,
+            x.size(2),
+            x.size(3),
+            x.size(4),
+        )
+        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            x.size(2) * self.factor_t,
+            x.size(4) * self.factor_s,
+            x.size(6) * self.factor_s,
+        )
+        if first_chunk:
+            x = x[:, :, self.factor_t - 1:, :, :]
+        return x
+
+
+class Down_ResidualBlock(nn.Module):
+
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 dropout,
+                 mult,
+                 temperal_downsample=False,
+                 down_flag=False):
+        super().__init__()
+
+        # Shortcut path with downsample
+        self.avg_shortcut = AvgDown3D(
+            in_dim,
+            out_dim,
+            factor_t=2 if temperal_downsample else 1,
+            factor_s=2 if down_flag else 1,
+        )
+
+        # Main path with residual blocks and downsample
+        downsamples = []
+        for _ in range(mult):
+            downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+
+        # Add the final downsample block
+        if down_flag:
+            mode = "downsample3d" if temperal_downsample else "downsample2d"
+            downsamples.append(Resample(out_dim, mode=mode))
+
+        self.downsamples = nn.Sequential(*downsamples)
+
+    def forward(self, x):
+        x_copy = x
+        for module in self.downsamples:
+            x = module(x)
+
+        return x + self.avg_shortcut(x_copy)
+
+
+class Up_ResidualBlock(nn.Module):
+
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 dropout,
+                 mult,
+                 temperal_upsample=False,
+                 up_flag=False):
+        super().__init__()
+        # Shortcut path with upsample
+        if up_flag:
+            self.avg_shortcut = DupUp3D(
+                in_dim,
+                out_dim,
+                factor_t=2 if temperal_upsample else 1,
+                factor_s=2 if up_flag else 1,
+            )
+        else:
+            self.avg_shortcut = None
+
+        # Main path with residual blocks and upsample
+        upsamples = []
+        for _ in range(mult):
+            upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+
+        # Add the final upsample block
+        if up_flag:
+            mode = "upsample3d" if temperal_upsample else "upsample2d"
+            upsamples.append(Resample(out_dim, mode=mode))
+
+        self.upsamples = nn.Sequential(*upsamples)
+
+    def forward(self, x):
+        # Single-pass: the whole latent sequence starts at frame 0, so the
+        # DupUp3D shortcut always applies the ``first_chunk`` front trim.
+        x_main = x
+        for module in self.upsamples:
+            x_main = module(x_main)
+        if self.avg_shortcut is not None:
+            x_shortcut = self.avg_shortcut(x, first_chunk=True)
+            return x_main + x_shortcut
+        else:
+            return x_main
+
+
+class Encoder3d(nn.Module):
+
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_down_flag = (
+                temperal_downsample[i]
+                if i < len(temperal_downsample) else False)
+            downsamples.append(
+                Down_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks,
+                    temperal_downsample=t_down_flag,
+                    down_flag=i != len(dim_mult) - 1,
+                ))
+            scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout),
+            AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout),
+        )
+
+        # # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1),
+        )
+
+    def forward(self, x):
+        # Single-pass encode: the whole (patchified) video goes through the
+        # encoder in one forward, no feat_cache / chunk loop.
+        x = self.conv1(x)
+        x = self.downsamples(x)
+        x = self.middle(x)
+        x = self.head(x)
+        return x
+
+
+class Decoder3d(nn.Module):
+
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[False, True, True],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2**(len(dim_mult) - 2)
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout),
+            AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_up_flag = temperal_upsample[i] if i < len(
+                temperal_upsample) else False
+            upsamples.append(
+                Up_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks + 1,
+                    temperal_upsample=t_up_flag,
+                    up_flag=i != len(dim_mult) - 1,
+                ))
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, 12, 3, padding=1),
+        )
+
+        # Toggle for single-pass per-block gradient checkpointing, propagated
+        # from AutoencoderKLWan2_2_.decode.
+        self.gradient_checkpointing = False
+
+    def forward(self, x):
+        # Single-pass decode. When gradient checkpointing is enabled each
+        # sub-block is recomputed during backward so only one block's
+        # activations stay alive at a time (vs. the whole upsampled video).
+        def _run(module, t):
+            if self.gradient_checkpointing and torch.is_grad_enabled():
+                return torch_checkpoint(module, t, use_reentrant=False)
+            return module(t)
+
+        x = self.conv1(x)
+        for layer in self.middle:
+            x = _run(layer, x)
+        for layer in self.upsamples:
+            x = _run(layer, x)
+        for layer in self.head:
+            x = _run(layer, x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class AutoencoderKLWan2_2_(nn.Module):
+
+    def __init__(
+        self,
+        dim=160,
+        dec_dim=256,
+        z_dim=16,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(
+            dim,
+            z_dim * 2,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_downsample,
+            dropout,
+        )
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(
+            dec_dim,
+            z_dim,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_upsample,
+            dropout,
+        )
+
+        # When True the decoder runs with per-block gradient checkpointing.
+        self.gradient_checkpointing = False
+
+    def forward(self, x, scale=[0, 1]):
+        mu = self.encode(x, scale)
+        x_recon = self.decode(mu, scale)
+        return x_recon, mu
+
+    def encode(self, x, scale):
+        # Single-pass encode: patchify the full video and run the encoder once
+        # (no chunk loop, no feat_cache).
+        scale = [item.to(x.device, x.dtype) for item in scale]
+        x = patchify(x, patch_size=2)
+        out = self.encoder(x)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        x = torch.cat([mu, log_var], dim=1)
+        return x
+
+    def decode(self, z, scale):
+        # Single-pass decode: run the whole latent sequence through the decoder
+        # in one forward. Because there is no feat_cache mutation, gradient
+        # checkpointing can be applied per decoder block (see Decoder3d.forward)
+        # instead of the previous per-frame cache-threading workaround.
+        scale = [item.to(z.device, z.dtype) for item in scale]
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        self.decoder.gradient_checkpointing = self.gradient_checkpointing
+        x = self.conv2(z)
+        out = self.decoder(x)
+        out = unpatchify(out, patch_size=2)
+        return out
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+
+def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs):
+    # params
+    cfg = dict(
+        dim=dim,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, True],
+        dropout=0.0,
+    )
+    cfg.update(**kwargs)
+
+    # init model
+    model = AutoencoderKLWan2_2_(**cfg)
+
+    return model
+
+
+class AutoencoderKLWan3_8_1Steps(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        latent_channels=48,
+        c_dim=160,
+        vae_pth=None,
+        dim_mult=[1, 2, 4, 4],
+        temperal_downsample=[False, True, True],
+        temporal_compression_ratio=4,
+        spatial_compression_ratio=16
+    ):
+        super().__init__()
+        mean = torch.tensor(
+            [
+                -0.2289,
+                -0.0052,
+                -0.1323,
+                -0.2339,
+                -0.2799,
+                0.0174,
+                0.1838,
+                0.1557,
+                -0.1382,
+                0.0542,
+                0.2813,
+                0.0891,
+                0.1570,
+                -0.0098,
+                0.0375,
+                -0.1825,
+                -0.2246,
+                -0.1207,
+                -0.0698,
+                0.5109,
+                0.2665,
+                -0.2108,
+                -0.2158,
+                0.2502,
+                -0.2055,
+                -0.0322,
+                0.1109,
+                0.1567,
+                -0.0729,
+                0.0899,
+                -0.2799,
+                -0.1230,
+                -0.0313,
+                -0.1649,
+                0.0117,
+                0.0723,
+                -0.2839,
+                -0.2083,
+                -0.0520,
+                0.3748,
+                0.0152,
+                0.1957,
+                0.1433,
+                -0.2944,
+                0.3573,
+                -0.0548,
+                -0.1681,
+                -0.0667,
+            ], dtype=torch.float32
+        )
+        std = torch.tensor(
+            [
+                0.4765,
+                1.0364,
+                0.4514,
+                1.1677,
+                0.5313,
+                0.4990,
+                0.4818,
+                0.5013,
+                0.8158,
+                1.0344,
+                0.5894,
+                1.0901,
+                0.6885,
+                0.6165,
+                0.8454,
+                0.4978,
+                0.5759,
+                0.3523,
+                0.7135,
+                0.6804,
+                0.5833,
+                1.4146,
+                0.8986,
+                0.5659,
+                0.7069,
+                0.5338,
+                0.4889,
+                0.4917,
+                0.4069,
+                0.4999,
+                0.6866,
+                0.4093,
+                0.5709,
+                0.6065,
+                0.6415,
+                0.4944,
+                0.5726,
+                1.2042,
+                0.5458,
+                1.6887,
+                0.3971,
+                1.0600,
+                0.3943,
+                0.5537,
+                0.5444,
+                0.4089,
+                0.7468,
+                0.7744,
+            ], dtype=torch.float32
+        )
+        self.scale = [mean, 1.0 / std]
+
+        # init model
+        self.model = _video_vae(
+            pretrained_path=vae_pth,
+            z_dim=latent_channels,
+            dim=c_dim,
+            dim_mult=dim_mult,
+            temperal_downsample=temperal_downsample,
+        ).eval().requires_grad_(False)
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, *args, **kwargs):
+        if "value" in kwargs:
+            self.gradient_checkpointing = kwargs["value"]
+        elif "enable" in kwargs:
+            self.gradient_checkpointing = kwargs["enable"]
+        else:
+            raise ValueError("Invalid set gradient checkpointing")
+        self.model.gradient_checkpointing = self.gradient_checkpointing
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        x = [
+            self.model.encode(u.unsqueeze(0), self.scale).squeeze(0)
+            for u in x
+        ]
+        x = torch.stack(x)
+        return x
+
+    @apply_forward_hook
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        h = self._encode(x)
+
+        posterior = DiagonalGaussianDistribution(h)
+
+        if not return_dict:
+            return (posterior,)
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(self, zs):
+        # Single-pass decode with per-block gradient checkpointing inside the
+        # decoder (see Decoder3d.forward), which lowers the backward peak to a
+        # single decoder block's activations while keeping the whole video in
+        # one forward (no feat_cache chunk loop).
+        self.model.gradient_checkpointing = self.gradient_checkpointing
+        dec = [
+            self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0)
+            for u in zs
+        ]
+        dec = torch.stack(dec)
+
+        return DecoderOutput(sample=dec)
+
+    @apply_forward_hook
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+        decoded = self._decode(z).sample
+
+        if not return_dict:
+            return (decoded,)
+        return DecoderOutput(sample=decoded)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_path, additional_kwargs={}):
+        def filter_kwargs(cls, kwargs):
+            import inspect
+            sig = inspect.signature(cls.__init__)
+            valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+            return filtered_kwargs
+
+        model = cls(**filter_kwargs(cls, additional_kwargs))
+        if pretrained_model_path.endswith(".safetensors"):
+            from safetensors.torch import load_file, safe_open
+            state_dict = load_file(pretrained_model_path)
+        else:
+            state_dict = torch.load(pretrained_model_path, map_location="cpu", weights_only=True)
+        tmp_state_dict = {} 
+        for key in state_dict:
+            if not key.startswith("model."):
+                tmp_state_dict["model." + key] = state_dict[key]
+            else:
+                tmp_state_dict[key] = state_dict[key]
+        state_dict = tmp_state_dict
+        m, u = model.load_state_dict(state_dict, strict=False)
+        print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
+        print(m, u)
+        return model

videox_fun/pipeline/pipeline_wan_flex_forcing.py

@@ -3,8 +3,7 @@
 import math
 from contextlib import nullcontext
 from dataclasses import dataclass
-from typing import (Any, Callable, Dict, List, Optional, Sequence, Tuple,
-                    Union)
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
@@ -13,14 +12,13 @@ from diffusers.utils import BaseOutput, logging, replace_example_docstring
 
 from ..models import (AutoencoderKLWan, AutoTokenizer, WanT5EncoderModel,
                       WanTransformer3DModel_FlexForcing)
+from ..utils.flex_chunking import (build_full_then_blocks_partitions,
+                                   build_pyramid_partitions, chunk_boundaries,
+                                   normalize_chunk_spec, uniform_chunks,
+                                   validate_nested_partitions)
 from ..utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
                                 get_sampling_sigmas)
 from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
-from ..utils.flex_chunking import (build_full_then_blocks_partitions,
-                                   build_pyramid_partitions,
-                                   chunk_boundaries, normalize_chunk_spec,
-                                   uniform_chunks,
-                                   validate_nested_partitions)
 from .pipeline_wan_self_forcing import (WanSelfForcingPipeline,
                                         WanSelfForcingPipelineOutput,
                                         retrieve_timesteps,
@@ -168,6 +166,11 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
                 (``[1] * F``), ``"bidir"`` (``[F]``) or an int for uniform
                 chunks. ``None`` keeps the inherited ``num_frame_per_block`` /
                 ``independent_first_frame`` schedule untouched.
+                The uniform forms (``None``, an int, ``"uniform[:N]"``) honour
+                ``independent_first_frame`` - which defaults to ``True`` here -
+                so ``"uniform:3"`` yields ``[1, 3, 3, ...]``, matching the layout
+                an IFF-trained checkpoint expects. Explicit size lists are taken
+                verbatim: write the leading ``1`` yourself if you want it.
             denoise_mode: Does the frame partition change with the noise level?
                 ``"fixed"`` (default) holds one partition across every denoising
                 step, i.e. the block-major Self-Forcing / 3.1 schedule.
@@ -302,7 +305,8 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
                 "generates and currently require T2V mode "
                 "(`initial_latent=None`).")
 
-        base = normalize_chunk_spec(chunk_spec, latent_frames) \
+        base = normalize_chunk_spec(chunk_spec, latent_frames,
+                                    independent_first_frame=independent_first_frame) \
             if chunk_spec is not None else None
         if depth > 1:
             if ladder_mode in FULL_THEN_BLOCKS_MODES:
@@ -310,6 +314,22 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
                 # two-level ladder whose fine level is the uniform
                 # `num_frame_per_block` partition rather than a binary split, so
                 # it does not go through `build_pyramid_partitions`.
+                # [NOTE] This ladder is fully determined by `num_frame_per_block`:
+                # it ignores both `chunk_spec` (level 0 is always the whole clip)
+                # and `min_num_frame_per_block` (the fine level is the uniform
+                # grid, not a binary split), so say so instead of dropping the
+                # caller's explicit arguments silently.
+                if chunk_spec is not None:
+                    logger.warning(
+                        "`denoise_mode=%r` ignores `chunk_spec=%r`: its level 0 is "
+                        "always one chunk over the whole clip and its fine level "
+                        "is the uniform `num_frame_per_block` partition.",
+                        ladder_mode, chunk_spec)
+                if min_num_frame_per_block != 1:
+                    logger.warning(
+                        "`denoise_mode=%r` ignores `min_num_frame_per_block=%r`: "
+                        "the finest level is fixed by `num_frame_per_block`.",
+                        ladder_mode, min_num_frame_per_block)
                 partitions = build_full_then_blocks_partitions(
                     latent_frames, num_frame_per_block,
                     independent_first_frame=independent_first_frame)
@@ -324,25 +344,71 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
             chunk_sizes = partitions[0]
         else:
             # `chunk_spec` cannot be None here: that case returned above.
-            chunk_sizes = normalize_chunk_spec(chunk_spec, latent_frames)
+            chunk_sizes = normalize_chunk_spec(
+                chunk_spec, latent_frames,
+                independent_first_frame=independent_first_frame)
             partitions = None
 
+        local_attn_size = getattr(self.transformer.config, 'local_attn_size', -1)
         if partitions is not None:
-            # The pyramid re-noises between levels, which only the stochastic
-            # schedule defines, and it moves `current_start` backwards whenever
-            # a chunk splits. Both are safe against a full KV cache and neither
-            # is against a rolling one.
-            if not stochastic_sampling:
+            # A level change hands the next level a buffer that sits exactly at
+            # `timesteps[step_idx + 1]`, and the finer level restarts at frame
+            # 0, so `current_start` moves backwards. Both are safe against a
+            # full KV cache and neither is against a rolling one, which evicts
+            # by `current_start` deltas.
+            #
+            # The handoff itself is not tied to `stochastic_sampling`. The
+            # re-noising branch defines the buffer as
+            # `(1 - t_i_1) * x0 + t_i_1 * eps`; the deterministic branch
+            # defines it as the Euler increment
+            # `x + (sigma_next - sigma_t) * flow_pred`, which lands on the
+            # very same sigma because `sigma_next` is read off the grid at
+            # `global_idx + 1` rather than off a scheduler counter. That
+            # increment is what `--flow_euler_rollout` trains, and the
+            # trainer's own KV-cache walk pops the same `(span, step_idx,
+            # commit)` stack, so a ladder under a deterministic schedule is
+            # legitimate - and is the rollout that actually matches a
+            # checkpoint distilled with `--flow_euler_rollout`.
+            if local_attn_size != -1:
                 raise ValueError(
-                    "Pyramid timestep chunking (`denoise_mode='pyramid'`) requires "
-                    "`stochastic_sampling=True`: the buffered step between two "
-                    "levels is the schedule's own re-noising.")
-            if getattr(self.transformer.config, 'local_attn_size', -1) != -1:
+                    "A per-step partition ladder (`denoise_mode='pyramid'` or "
+                    "'full_then_blocks') requires `local_attn_size=-1` (full KV "
+                    "cache). A rolling window evicts by `current_start` deltas, "
+                    "which a level change moves backwards.")
+            # Forcing-KV (arXiv 2605.09681) rewrites the cache by grouping heads
+            # into static/dynamic histories on a *clean* pass, and it anchors its
+            # `ar_step` on a fixed block size. A ladder breaks both premises: the
+            # partition changes between steps, so neither the per-block clean
+            # commit nor a fixed stride survives it.
+            if self._forcing_kv_enabled():
                 raise ValueError(
-                    "Pyramid timestep chunking (`denoise_mode='pyramid'`) requires "
-                    "`local_attn_size=-1` (full KV cache). A rolling window "
-                    "evicts by `current_start` deltas, which a splitting chunk "
-                    "moves backwards.")
+                    "A per-step partition ladder (`denoise_mode='pyramid'` or "
+                    "'full_then_blocks') cannot be combined with Forcing-KV head "
+                    "compression: compression commits a clean pass per block and "
+                    "derives its step from a fixed block size, neither of which "
+                    "survives a partition that changes between steps.")
+        # Applies to every mode, not just the pyramid. The partition defines one
+        # bidirectional chunk, so a frame must see the whole of its own chunk. A
+        # rolling window is anchored at the chunk *end* (`window_start =
+        # local_end_index - local_attn_size * frame_seqlen`), so once a chunk is
+        # wider than the window its earliest frames cannot attend to its latest
+        # ones and the chunk silently degrades to partial causality. `sink_size`
+        # does not rescue this: it pins the clip start, not the neighbourhood of
+        # the chunk being denoised.
+        # `partitions` is the ladder (a list of levels, each a partition) when
+        # one is on; otherwise the single `chunk_sizes` partition is the only
+        # level. Widest chunk over every level a rollout can reach.
+        levels = partitions if partitions is not None else [chunk_sizes]
+        widest = max(max(level) for level in levels)
+        if local_attn_size != -1 and widest > local_attn_size:
+            raise ValueError(
+                f"Partition breaks the local attention window: the widest chunk "
+                f"is {widest} latent frames but `local_attn_size="
+                f"{local_attn_size}`, so a chunk's first frames cannot attend to "
+                f"its last ones and the intra-chunk bidirectionality that "
+                f"Flex-Forcing 3.1 relies on is lost. Either use "
+                f"`local_attn_size=-1`, or keep every chunk at or below "
+                f"`local_attn_size` frames.")
 
         # Below is the Flex-Forcing rollout proper. It performs the same setup
         # the inherited `__call__` does - prompt encoding, timesteps, latents,
@@ -490,6 +556,11 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
         # copy of the initial noise in which every frame sits at its current
         # noise level, and which the walk both reads from and writes back into.
         state = noise.clone()
+        # The stochastic grid carries a trailing element that the re-noising
+        # reads as `timesteps[global_idx + 1]`, so it is not a step and gets
+        # sliced off. The deterministic grid needs no tail of its own: its
+        # transfer reads `scheduler.sigmas`, which `set_timesteps` already
+        # terminates with a 0.
         denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps
         last_step = len(denoise_timesteps) - 1
         block_idx = 0
@@ -531,8 +602,10 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
             # A splitting span takes a single buffered step; a leaf runs
             # whatever is left of the schedule to completion. `step_idx` is
             # where the slice starts inside `timesteps`, which is what lets the
-            # re-noising below read the right (t_i, t_i+1) pair when the walk
-            # resumes a span mid-schedule.
+            # transfer below pick up the right (t_i, t_i+1) pair when the walk
+            # resumes a span mid-schedule - indexed into `timesteps` directly
+            # under a stochastic rollout, resolved by the scheduler from the `t`
+            # it is handed under a deterministic one.
             schedule = (denoise_timesteps[step_idx:step_idx + 1] if splits
                         else denoise_timesteps[step_idx:])
 
@@ -650,6 +723,26 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
                         # Delegate to the scheduler's own step() so each sampler
                         # (Flow Euler, UniPC, DPM++) applies its multi-step
                         # formula, exactly as the inherited rollout does.
+                        #
+                        # This holds under a ladder as well, which is not
+                        # obvious: the walk revisits the schedule out of order,
+                        # but the per-span `_step_index = None` reset above
+                        # makes step() re-derive its position by looking the
+                        # handed-over `t` up in `scheduler.timesteps` instead of
+                        # counting calls, and inside one span the counter then
+                        # advances in step with `schedule`. `sigmas` carries the
+                        # trailing 0 `set_timesteps` appends, so the leaf's last
+                        # step lands exactly on x0.
+                        #
+                        # Known and accepted: this grid is
+                        # `set_timesteps(num_inference_steps)`, which at shift=5
+                        # gives [1000, 909.7, 717.3, 24.4], whereas the trainer
+                        # distils on [1000, 937.5, 833.3, 625] - the grid
+                        # `stochastic_sampling_timesteps` reproduces entry for
+                        # entry. A deterministic ladder therefore samples noise
+                        # levels this checkpoint was not distilled at. Use
+                        # `stochastic_sampling=True` when the samples have to
+                        # sit on the trained trajectory.
                         noisy_input = self.scheduler.step(
                             flow_pred, t, noisy_input, return_dict=False)[0]
                         denoised_pred = noisy_input
@@ -787,6 +880,7 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
         edit_steps: int = 1,
         shift: float = 5.0,
         context_noise: int = 0,
+        stochastic_sampling: bool = True,
         num_frame_per_block: Optional[int] = None,
         max_sequence_length: int = 512,
         generator: Optional[torch.Generator] = None,
@@ -818,6 +912,12 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
             edit_steps: How many *trailing* steps of that schedule to run. Keep
                 it small - editing at a planning timestep would restructure the
                 clip instead of refining it.
+            stochastic_sampling: Which rollout the refinement runs on - the same
+                knob :meth:`__call__` takes. ``True`` re-noises between steps on
+                the hardcoded ``stochastic_sampling_timesteps`` grid. ``False``
+                delegates each step to ``self.scheduler`` (Flow Euler, UniPC,
+                DPM++), which is what a checkpoint distilled with
+                ``--flow_euler_rollout`` was trained on.
             num_frame_per_block: Granularity of the clean-context commit - the
                 same uniform block size the inherited rollout uses, doing the
                 same job. ``None`` commits the whole clip in one bidirectional
@@ -875,8 +975,28 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
                             * (latents.shape[4] // patch_size[2]))
         seq_len = frame_seq_length * num_latent_frames
 
-        timesteps = stochastic_sampling_timesteps(
-            num_inference_steps, shift, device)
+        # The same schedule dispatch `__call__` makes: the stochastic grid is
+        # the hardcoded few-step one, the deterministic grid comes from the
+        # scheduler so that whichever sampler is configured supplies its own
+        # multi-step formula. `num_inference_steps` may be rewritten by
+        # `retrieve_timesteps` and `step_offset` below is derived from it, so
+        # this has to run first.
+        if stochastic_sampling:
+            timesteps = stochastic_sampling_timesteps(
+                num_inference_steps, shift, device)
+        elif isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+            timesteps, num_inference_steps = retrieve_timesteps(
+                self.scheduler, num_inference_steps, device, None)
+        elif isinstance(self.scheduler, FlowUniPCMultistepScheduler):
+            self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
+            timesteps = self.scheduler.timesteps
+        elif isinstance(self.scheduler, FlowDPMSolverMultistepScheduler):
+            sampling_sigmas = get_sampling_sigmas(num_inference_steps, shift)
+            timesteps, _ = retrieve_timesteps(
+                self.scheduler, device=device, sigmas=sampling_sigmas)
+        else:
+            timesteps, num_inference_steps = retrieve_timesteps(
+                self.scheduler, num_inference_steps, device, None)
         self._num_timesteps = len(timesteps)
 
         self._initialize_kv_cache(
@@ -957,18 +1077,27 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
         # 3. Denoise the span at the refinement levels only, with the whole clip
         #    visible so it conditions on past *and* future clean tokens - the
         #    any-order half of 4.2, which is what `flex_state` opens here and
-        #    step 1 deliberately left closed. This call is always stochastic, so
-        #    the schedule is the trailing `edit_steps` of `timesteps[:-1]` and
-        #    the step is the schedule's own re-noising. `self.scheduler` takes no
-        #    part in this method at all - `timesteps` above comes from
-        #    `stochastic_sampling_timesteps`, and neither loop calls
-        #    `scheduler.step()` - so there is no scheduler state to reset here.
+        #    step 1 deliberately left closed. Either rollout runs the trailing
+        #    `edit_steps` of the schedule; they differ in whether the grid
+        #    carries a tail. The stochastic one does - the re-noising below reads
+        #    it as `timesteps[global_idx + 1]` - so it is sliced off here. The
+        #    deterministic one needs no tail of its own, because its transfer
+        #    reads `scheduler.sigmas`, which `set_timesteps` already terminates
+        #    with a 0.
         flex_state = {"attn_window": (0, seq_len)}
-        schedule = timesteps[:-1][step_offset:]
+        denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps
+        schedule = denoise_timesteps[step_offset:]
         span_frames = edit_end - edit_start
         current_start = edit_start * frame_seq_length
         noisy_input = noisy_span
         edited = noisy_span
+        # Only the deterministic branch touches `self.scheduler`, but that one
+        # keeps its own position there and, for the multi-step samplers, a
+        # history of model outputs. Start both fresh, the way `__call__` does at
+        # the head of every span.
+        self.scheduler._step_index = None
+        if hasattr(self.scheduler, "model_outputs"):
+            self.scheduler.model_outputs = []
         with self.progress_bar(total=len(schedule)) as progress_bar:
             for local_idx, t in enumerate(schedule):
                 timestep = torch.ones([batch_size, span_frames], device=device,
@@ -1042,13 +1171,24 @@ class WanFlexForcingPipeline(WanSelfForcingPipeline):
                             f"Unexpected flow_pred dim: {flow_pred.dim()}, "
                             f"shape: {flow_pred.shape}")
 
-                global_idx = step_offset + local_idx
-                t_i = (timesteps[global_idx] / 1000).to(weight_dtype)
-                t_i_1 = (timesteps[global_idx + 1] / 1000).to(weight_dtype)
-                edited = noisy_input - flow_pred * t_i
-                noisy_input = (1 - t_i_1) * edited + t_i_1 * torch.randn(
-                    edited.shape, dtype=edited.dtype, device=device,
-                    generator=generator)
+                if stochastic_sampling:
+                    global_idx = step_offset + local_idx
+                    t_i = (timesteps[global_idx] / 1000).to(weight_dtype)
+                    t_i_1 = (timesteps[global_idx + 1] / 1000).to(weight_dtype)
+                    edited = noisy_input - flow_pred * t_i
+                    noisy_input = (1 - t_i_1) * edited + t_i_1 * torch.randn(
+                        edited.shape, dtype=edited.dtype, device=device,
+                        generator=generator)
+                else:
+                    # Delegate to the scheduler's own step(), exactly as
+                    # `__call__` does. This refinement is a single linear span
+                    # rather than a walk, so the position `step()` re-derives
+                    # from the `t` it is handed is the one the schedule is at,
+                    # and the trailing sigma of 0 lands the last step on x0 -
+                    # which is what `output` below writes back into the clip.
+                    noisy_input = self.scheduler.step(
+                        flow_pred, t, noisy_input, return_dict=False)[0]
+                    edited = noisy_input
                 progress_bar.update()
 
         output = latents.clone()

videox_fun/utils/flex_chunking.py

@@ -26,7 +26,7 @@ Conventions used across this module:
 * ``num_frames`` always denotes **latent** frames (post-VAE), matching
   ``target_shape[1]`` in the trainers and ``num_latent_frames`` in the pipeline.
 """
-from typing import List, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import torch
 
@@ -45,6 +45,8 @@ __all__ = [
     "chunk_ends_tensor",
     "broadcast_chunk_sizes",
     "chunk_sizes_to_block_kwargs",
+    "FLEX_RANDOM_PROB",
+    "sample_flex_level0",
 ]
 
 
@@ -89,7 +91,8 @@ def uniform_chunks(num_frames: int,
 
 def normalize_chunk_spec(spec: Union[None, int, str, Sequence[int]],
                          num_frames: int,
-                         chunk_size: Optional[int] = None) -> ChunkSizes:
+                         chunk_size: Optional[int] = None,
+                         independent_first_frame: bool = False) -> ChunkSizes:
     """Parse a user-facing chunk specification into a validated partition.
 
     Accepted forms (``num_frames`` = latent frames the partition must cover):
@@ -102,18 +105,26 @@ def normalize_chunk_spec(spec: Union[None, int, str, Sequence[int]],
     * ``"ar"``                 -> fully autoregressive ``[1] * num_frames``
     * ``"bidir"`` / ``"full"`` -> fully bidirectional ``[num_frames]``
 
-    ``independent_first_frame`` needs no separate flag: it is exactly the
-    partition whose first chunk has size 1 (read it back with
-    ``sizes[0] == 1``).
+    ``independent_first_frame`` is applied to the *uniform* forms only (``None``,
+    ``int``, ``"uniform[:N]"``), where it inserts the leading size-1 chunk that
+    makes the layout ``[1, N, N, ...]``. Without it those forms produce a
+    partition the ``independent_first_frame`` trainer never trained on. The
+    keyword forms leave it alone on purpose: ``"ar"`` is already one frame per
+    chunk, and prefixing ``"bidir"`` would break the single bidirectional chunk.
+    An explicit size list also leaves it alone - a leading 1 there is written by
+    the caller, and is read back with ``sizes[0] == 1``.
     """
     if spec is None:
-        sizes = uniform_chunks(num_frames, chunk_size if chunk_size is not None else 1)
+        sizes = uniform_chunks(num_frames, chunk_size if chunk_size is not None else 1,
+                               independent_first_frame=independent_first_frame)
     elif isinstance(spec, bool):
         raise TypeError(f"chunk spec must not be a bool, got {spec!r}")
     elif isinstance(spec, int):
-        sizes = uniform_chunks(num_frames, spec)
+        sizes = uniform_chunks(num_frames, spec,
+                               independent_first_frame=independent_first_frame)
     elif isinstance(spec, str):
-        sizes = _parse_chunk_string(spec, num_frames)
+        sizes = _parse_chunk_string(spec, num_frames,
+                                    independent_first_frame=independent_first_frame)
     elif isinstance(spec, (list, tuple)):
         sizes = [int(s) for s in spec]
     else:
@@ -124,7 +135,8 @@ def normalize_chunk_spec(spec: Union[None, int, str, Sequence[int]],
     return _validate_sizes(sizes, num_frames, spec)
 
 
-def _parse_chunk_string(spec: str, num_frames: int) -> ChunkSizes:
+def _parse_chunk_string(spec: str, num_frames: int,
+                        independent_first_frame: bool = False) -> ChunkSizes:
     text = spec.strip()
     lowered = text.lower()
     if lowered in ("ar", "autoregressive", "causal"):
@@ -132,9 +144,11 @@ def _parse_chunk_string(spec: str, num_frames: int) -> ChunkSizes:
     if lowered in ("bidir", "bidirectional", "full"):
         return [int(num_frames)]
     if lowered.startswith("uniform:"):
-        return uniform_chunks(num_frames, int(lowered.split(":", 1)[1]))
+        return uniform_chunks(num_frames, int(lowered.split(":", 1)[1]),
+                              independent_first_frame=independent_first_frame)
     if lowered.startswith("uniform"):
-        return uniform_chunks(num_frames, 1)
+        return uniform_chunks(num_frames, 1,
+                              independent_first_frame=independent_first_frame)
     body = text.strip("[]() ")
     sep = "," if "," in body else ("-" if "-" in body else None)
     if sep is None:
@@ -159,18 +173,41 @@ def _validate_sizes(sizes: ChunkSizes, num_frames: int, spec) -> ChunkSizes:
     return [int(s) for s in sizes]
 
 
-# Share of iterations that pin the launcher's own uniform ``num_frame_per_block``
-# layout instead of drawing a random partition. ``sample_flexible_chunks`` can
-# reach every partition whose chunks are individually legal, but it reaches any
-# *particular* one only by chance: the uniform ``[3] * 7`` at 21 latent frames
-# has probability 1/183708, i.e. 0.11 expected hits in 20000 iterations. Without
-# this band ``--num_frame_per_block`` would name a layout the model never trains
-# on, and the validation pass - which pins exactly that layout - would render
-# out of distribution. Raising the sampler's randomness makes this *worse*, not
-# better: more reachable partitions means each one is drawn less often.
-# Deliberately not a CLI flag, and shared by both trainers so the two stages
-# cannot drift apart.
-UNIFORM_BLOCK_PROB = 0.1
+# Share of iterations the level-0 mixture's *random* arm takes; the other arms
+# split the remainder equally, so with a pyramid this is 30% coarse / 30%
+# uniform / 30% full_then_blocks / 10% random. ``sample_flex_level0`` below is
+# the single source of truth for the coarsest layout that **both** trainers
+# draw: stage 2 (``train_distill.py``) nests it into the §3.2 ladder, stage 1
+# (``train_ode.py``) regresses against it directly because it predicts x0 in one
+# forward per chunk and has no ladder to build. Sharing the draw is what keeps
+# the two stages on the same layout distribution - stage 1's checkpoint is stage
+# 2's initialisation (``--ode_transformer_path``), so a layout stage 2 spends
+# most of its iterations on must not be one stage 1 never regressed.
+#
+# The layouts ``--num_frame_per_block`` names are arms, not something the random
+# draw is expected to stumble onto: ``sample_flexible_chunks`` reaches every
+# individually legal partition but any *particular* one only by chance. At 21
+# latent frames the uniform ``[3] * 7`` has probability 1/183708, and the fully
+# causal ``[1] * 21`` is 1/21! - unreachable in practice at any ``max_chunk``,
+# which is why no ``--flex_chunk_min/max`` setting can substitute for these arms.
+# Pinning them is what makes ``--num_frame_per_block`` name a layout the model
+# genuinely trains on, and keeps ``log_validation`` - which pins exactly that
+# layout - in distribution. Deliberately not a CLI flag.
+#
+# Deliberate deviation from the paper, and the reason this constant exists at
+# all: arXiv 2607.03509 §5 says "we randomly pick the chunk size in training
+# from 2 to 10" and pins no arm of its own, so its level 0 is drawn entirely
+# from that random family. That family spans thousands of distinct layouts, so
+# the random arm buys exactly what §3.3 asks of it - coverage of "attention
+# configurations that range from strictly causal to fully non-causal" - but it
+# cannot train any *particular* layout, and the layouts inference actually
+# selects (``PAPER_CHUNK_CONFIGS``, plus the one ``log_validation`` pins) are
+# the deterministic ones. Holding it to a minority share is what leaves those
+# arms enough iterations to be trained at all. Setting this to 1.0 reproduces
+# the paper's mixture exactly: every threshold below collapses to 0, so the
+# three deterministic arms become unreachable and the uniform draw always lands
+# in the random one.
+FLEX_RANDOM_PROB = 0.10
 
 
 def sample_flexible_chunks(num_frames: int,
@@ -205,21 +242,13 @@ def sample_flexible_chunks(num_frames: int,
         # Two legal moves, offered everywhere along the clip including the tail:
         # take the whole remainder, or leave at least min_chunk frames behind.
         # Anything in between would strand a sub-min_chunk sliver mid-video.
-        #
-        # The tail used to be closed unconditionally (`if remaining <=
-        # max_chunk: append(remaining); break`), which silently removed most of
-        # the sampler's randomness: at 21 frames with 2..10 a run of 3s stopped
-        # at `[3, 3, 3, 3, 9]` after four draws, so the uniform `[3] * 7` was
-        # unreachable outright - exhaustive enumeration found 679 partitions and
-        # none of them was it. Uniform layouts also bias short: the old rule
-        # topped out at 7 chunks with 65% of draws at just 3.
         upper = min(max_chunk, remaining - min_chunk)
         choices = list(range(min_chunk, upper + 1)) if upper >= min_chunk else []
         if remaining <= max_chunk:
             choices.append(remaining)
         if not choices:
             # Degenerate only: min_chunk == max_chunk and the remainder is
-            # strictly between them, so nothing legal fits. Same as before.
+            # strictly between them, so nothing legal fits.
             sizes.append(remaining)
             break
         step = choices[int(torch.randint(
@@ -414,3 +443,172 @@ def chunk_sizes_to_block_kwargs(chunk_sizes: Sequence[int]) -> dict:
         "num_frame_per_block": max(sizes) if sizes else 1,
         "independent_first_frame": bool(sizes and sizes[0] == 1),
     }
+
+
+def sample_flex_level0(num_frames: int,
+                       num_frame_per_block: int,
+                       chunk_min: int,
+                       chunk_max: int,
+                       pyramid: bool,
+                       generator: Optional[torch.Generator] = None,
+                       device: Union[str, torch.device] = "cpu",
+                       independent_first_frame: bool = False) -> Dict[str, Any]:
+    """Draw one iteration's coarsest Flex-Forcing layout (the four-arm mixture).
+
+    Single source of truth for the level-0 partition **both** trainers draw:
+    ``train_distill.py`` nests the result into the §3.2 ladder, ``train_ode.py``
+    regresses against it directly. Keeping the draw here is what stops the two
+    stages from drifting apart - stage 1's checkpoint is stage 2's
+    initialisation, so every layout stage 2 spends iterations on has to be one
+    stage 1 actually regressed.
+
+    Four arms, decided by a single uniform draw. ``"random"`` takes
+    :data:`FLEX_RANDOM_PROB` (10%) and the other three split what is left
+    equally, i.e. 30% each:
+
+    * ``"coarse"`` - one chunk over the whole clip (inference's planning step),
+      i.e. fully bidirectional. Only drawn when ``pyramid`` is on: without a
+      ladder a ``[F]`` level 0 would be a degenerate single-chunk rollout, so
+      the two surviving arms widen to 45% each and ``"random"`` stays at 10%.
+    * ``"uniform"`` - a pure Self-Forcing grid (no binary refinement) in one of
+      three sub-layouts drawn from ``num_frame_per_block`` and fully causal:
+      ``[1]*F`` alone, ``[nfpb]*k`` alone, or ``[nfpb]*k`` planning followed by
+      ``[1]*F`` for every remaining step - the coarse-block -> per-frame descent
+      the pyramid never samples, because it halves through ``[2, 1x14]`` first.
+      The three are equiprobable *within* the arm, so 10% of iterations each.
+    * ``"full_then_blocks"`` - ``denoise_mode="full_then_blocks"``: a whole-clip
+      planning chunk, then block-major refinement at a width drawn from
+      ``{1, num_frame_per_block}``.
+    * ``"random"`` - the §3.3 / §5 spectrum via :func:`sample_flexible_chunks`.
+
+    Distributed contract: the arm selector is drawn once and broadcast from
+    rank 0, then **every** arm issues exactly one :func:`broadcast_chunk_sizes`.
+    All ranks therefore end up with a byte-identical layout *and* the same
+    collective count. A per-rank arm draw would leave ranks in different arms
+    with different walk forward counts -> FSDP all-gather desync (NCCL hang).
+
+    Args:
+        pyramid: whether the caller ladders at all, i.e.
+            ``args.flex_pyramid_levels > 1``. Stage 1 has no ladder but must
+            still pass stage 2's setting, because this flag decides whether the
+            coarse arm exists and therefore what the other three arms weigh.
+
+    Returns:
+        ``"sizes"`` - the level-0 partition, already rank-reconciled.
+        ``"arm"`` - ``"coarse"`` / ``"uniform_causal"`` / ``"uniform_blocks"`` /
+        ``"uniform_plan_then_causal"`` / ``"full_then_blocks"`` / ``"random"``.
+        The uniform arm is resolved to its sub-layout here so the caller can
+        rebuild the ladder depth without re-deriving the band arithmetic.
+        ``"arm_text"`` - human-readable description for the layout log.
+        ``"causal"`` - ``[1]*F``, present on ``"uniform_plan_then_causal"``.
+        ``"block_level"`` - the block-major level, present on
+        ``"full_then_blocks"``.
+    """
+    num_frames = int(num_frames)
+    nfpb = max(1, int(num_frame_per_block))
+    # The arm selector must be IDENTICAL on every rank, not just its broadcast
+    # count: the arms build their ladder from different tensors, so a per-rank
+    # `u` would desync the group. Broadcast rank 0's draw; the arm's own single
+    # broadcast then reconciles its within-arm randomness (the random base, the
+    # block width) to rank 0.
+    u = torch.rand((), generator=generator, device=device)
+    try:
+        import torch.distributed as dist
+    except ImportError:  # pragma: no cover - torch always ships distributed
+        dist = None
+    if dist is not None and dist.is_available() and dist.is_initialized():
+        dist.broadcast(u.reshape(1), src=0)
+    u = u.item()
+
+    # `random` holds a fixed slice and the deterministic arms divide the rest,
+    # so switching the coarse arm off widens the two survivors instead of
+    # silently fattening `random`. `ftb_hi` is both the end of the last
+    # deterministic arm and the start of the random one, and it is the final
+    # threshold tested, so no rounding remainder can fall through unhandled.
+    det_arms = 3 if pyramid else 2
+    det_w = (1.0 - FLEX_RANDOM_PROB) / det_arms
+    coarse_hi = det_w if pyramid else 0.0
+    uniform_hi = coarse_hi + det_w
+    ftb_hi = uniform_hi + det_w
+
+    if u < coarse_hi:
+        # Coarse end of §3.1: reuse the ladder builder's own level-0 rule so the
+        # `independent_first_frame` handling cannot drift from inference's.
+        sizes = build_pyramid_partitions(
+            num_frames, num_levels=1, base_chunks=None,
+            independent_first_frame=independent_first_frame)[0]
+        return {
+            "sizes": broadcast_chunk_sizes(sizes, device=device),
+            "arm": "coarse",
+            "arm_text": "whole clip, the 3.2 coarse planning layout",
+        }
+
+    if u < uniform_hi:
+        # The three sub-layouts change the ladder DEPTH (hence stage 2's walk
+        # forward count), so the choice must be rank-consistent. It is free here:
+        # `u` was already broadcast, so slicing its position within the uniform
+        # band gives every rank the same sub-arm with no extra collective. Like
+        # every arm this issues EXACTLY ONE broadcast. `uniform_chunks` is called
+        # directly rather than through `normalize_chunk_spec`, which only applies
+        # `independent_first_frame` to its uniform forms via a keyword - the
+        # direct call keeps this arm visibly identical to the bands around it.
+        band = (u - coarse_hi) / det_w                 # [0, 1) in the uniform arm
+        sub = min(int(band * 3), 2)                     # 0 / 1 / 2, equiprobable
+        causal = uniform_chunks(num_frames, 1,
+                                independent_first_frame=independent_first_frame)
+        if sub == 0:
+            return {
+                "sizes": broadcast_chunk_sizes(causal, device=device),
+                "arm": "uniform_causal",
+                "arm_text": "uniform 1-frame blocks, single level (fully causal)",
+            }
+        sizes = broadcast_chunk_sizes(
+            uniform_chunks(num_frames, nfpb,
+                           independent_first_frame=independent_first_frame),
+            device=device)
+        if sub == 1:
+            return {"sizes": sizes, "arm": "uniform_blocks",
+                    "arm_text": f"uniform {nfpb}-frame blocks, single level"}
+        return {
+            "sizes": sizes, "arm": "uniform_plan_then_causal",
+            "causal": list(causal),
+            "arm_text": (f"uniform {nfpb}-frame plan -> fully causal [1x"
+                         f"{num_frames}] for the remaining steps"),
+        }
+
+    if u < ftb_hi:
+        # "First step full, every later step block-major": a fixed two-level
+        # ladder that does NOT go through the binary pyramid refinement, so the
+        # block-major level is reached at high noise (the 2nd step) instead of as
+        # a deep refinement - the one trajectory the pyramid arms never sample.
+        # The block width is drawn between fully causal (1 frame / block) and the
+        # launcher's `num_frame_per_block`, so both the tight-AR and the
+        # coarse-block refinement of the whole-clip plan get trained. Level 0 is
+        # just `[F]`, identical on every rank and needing no sync, so the ONE
+        # broadcast this arm issues covers the block-major level; rank 0's draw
+        # wins and every rank rebuilds the same two-level ladder locally.
+        block_choices = sorted({1, nfpb})
+        pick = int(torch.rand((), generator=generator, device=device).item()
+                   * len(block_choices))
+        block = block_choices[min(pick, len(block_choices) - 1)]
+        block = min(block, num_frames)
+        ftb = build_full_then_blocks_partitions(
+            num_frames, block, independent_first_frame=independent_first_frame)
+        ftb[1] = broadcast_chunk_sizes(ftb[1], device=device)
+        return {
+            "sizes": list(ftb[0]), "arm": "full_then_blocks",
+            "block_level": list(ftb[1]),
+            "arm_text": (f"full clip -> uniform {max(ftb[1])}-frame blocks, "
+                         "the full_then_blocks ladder"),
+        }
+
+    sizes = sample_flexible_chunks(
+        num_frames, min_chunk=chunk_min, max_chunk=chunk_max,
+        generator=generator, device=device,
+        independent_first_frame=independent_first_frame)
+    return {
+        "sizes": broadcast_chunk_sizes(sizes, device=device),
+        "arm": "random",
+        "arm_text": (f"random {chunk_min}..{chunk_max}-frame blocks, "
+                     "the 3.3 stochastic-chunk spectrum"),
+    }