Update multigpus in vsa, sla, spark

examples/fastvideo/predict_t2v_minimax_h3.py

@@ -40,9 +40,10 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
 # GPU at full load, so keep an offloading mode here.
 GPU_memory_mode     = "model_group_offload"
 # Multi GPUs config
-# NOTE: the VSA sparse kernel has no sequence-parallel path, so a FastVideo-H3 student runs
-# single-GPU only; keep ulysses_degree = ring_degree = 1 (enable_multi_gpus_inference() raises
-# otherwise, rather than silently falling back to the dense multi-GPU processor).
+# The FastVideo-H3 VSA student supports sequence-parallel inference through ulysses head-parallel
+# only (the VSA kernel has no ring-attention path): put every GPU into ulysses_degree, keep
+# ring_degree = 1, and ulysses_degree must divide the 56 attention heads (2 / 4 / 7 / 8 / ...).
+# The defaults stay single-GPU because the ~19B stack still needs an offloading mode below.
 ulysses_degree      = 1
 ring_degree         = 1
 # Use FSDP to save more GPU memory in multi gpus. Unused at ulysses_degree = ring_degree = 1.
@@ -247,8 +248,14 @@ if use_qfloat8:
 wants_offload = GPU_memory_mode in ("model_cpu_offload", "model_group_offload", "sequential_cpu_offload")
 
 if ulysses_degree > 1 or ring_degree > 1:
+    # The VSA sparse kernel only supports ulysses head-parallel (no ring-attention path); a VSA
+    # student under ring_degree>1 would silently fold ring into a larger ulysses degree and give
+    # none of ring's per-rank memory saving, so reject it instead of mis-attending.
+    if ring_degree > 1:
+        raise ValueError(
+            "FastVideo-H3 VSA multi-GPU requires ring_degree=1; put every GPU into ulysses_degree "
+            "instead (ulysses_degree must divide the 56 attention heads).")
     from functools import partial
-    # NOTE: raises for the VSA student (no sequence-parallel path); keep ulysses_degree=ring_degree=1.
     transformer.enable_multi_gpus_inference()
     if fsdp_dit:
         fp32_modules = [m for m in transformer.modules()

examples/fastvideo/predict_t2v_wan2.1.py

@@ -239,6 +239,12 @@ pipeline = WanPipeline(
     scheduler=scheduler,
 )
 if ulysses_degree > 1 or ring_degree > 1:
+    # The VSA kernel has no ring-attention path; a sparse student under ring_degree>1 would
+    # silently mis-attend. student_topk >= 1.0 is the dense fallback, which ring does support.
+    if ring_degree > 1 and student_topk < 1.0:
+        raise ValueError(
+            "FastVideo VSA multi-GPU requires ring_degree=1; put every GPU into ulysses_degree "
+            "instead (num_heads must divide ulysses_degree).")
     from functools import partial
     transformer.enable_multi_gpus_inference()
     if fsdp_dit:

examples/sparkdiffusion/predict_t2v_wan2.1.py

@@ -295,6 +295,7 @@ def save_results():
     else:
         video_path = os.path.join(save_path, prefix + ".mp4")
         save_videos_grid(sample, video_path, fps=fps)
+        print(f"Saved video to: {video_path}")
 
 if ulysses_degree * ring_degree > 1:
     import torch.distributed as dist

examples/turbodiffusion/predict_i2v_wan2.2.py

@@ -213,6 +213,12 @@ pipeline = Wan2_2I2VPipeline(
     scheduler=scheduler,
 )
 if ulysses_degree > 1 or ring_degree > 1:
+    # The SLA sparse kernel only supports ulysses head-parallel (no ring-attention path); a
+    # sparse expert under ring_degree>1 would silently mis-attend.
+    if ring_degree > 1:
+        raise ValueError(
+            "TurboDiffusion SLA multi-GPU requires ring_degree=1; put every GPU into "
+            "ulysses_degree instead (num_heads must divide ulysses_degree).")
     from functools import partial
     transformer.enable_multi_gpus_inference()
     transformer_2.enable_multi_gpus_inference()
@@ -289,7 +295,6 @@ def save_results():
         image = (image * 255).numpy().astype(np.uint8)
         image = Image.fromarray(image)
         image.save(video_path)
-        print(f"Saved image to: {video_path}")
     else:
         video_path = os.path.join(save_path, prefix + ".mp4")
         save_videos_grid(sample, video_path, fps=fps)

examples/turbodiffusion/predict_t2v_wan2.1.py

@@ -67,7 +67,7 @@ sampler_name        = "Flow"
 # Used when the sampler is in "Flow_Unipc", "Flow_DPM++".
 # If you want to generate a 480p video, it is recommended to set the shift value to 3.0.
 # If you want to generate a 720p video, it is recommended to set the shift value to 5.0.
-shift               = 3 
+shift               = 5
 
 # Load pretrained model if need
 # Download from https://huggingface.co/TurboDiffusion/TurboWan2.1-T2V-1.3B-480P/
@@ -172,6 +172,12 @@ pipeline = WanPipeline(
     scheduler=scheduler,
 )
 if ulysses_degree > 1 or ring_degree > 1:
+    # The SLA sparse kernel only supports ulysses head-parallel (no ring-attention path); a
+    # sparse student under ring_degree>1 would silently mis-attend.
+    if ring_degree > 1:
+        raise ValueError(
+            "TurboDiffusion SLA multi-GPU requires ring_degree=1; put every GPU into "
+            "ulysses_degree instead (num_heads must divide ulysses_degree).")
     from functools import partial
     transformer.enable_multi_gpus_inference()
     if fsdp_dit:
@@ -234,7 +240,6 @@ def save_results():
         image = (image * 255).numpy().astype(np.uint8)
         image = Image.fromarray(image)
         image.save(video_path)
-        print(f"Saved image to: {video_path}")
     else:
         video_path = os.path.join(save_path, prefix + ".mp4")
         save_videos_grid(sample, video_path, fps=fps)

examples/wan2.1_one_forcing/predict_t2v.py

@@ -0,0 +1,303 @@
+import os
+import sys
+
+import numpy as np
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from omegaconf import OmegaConf
+from PIL import Image
+
+current_file_path = os.path.abspath(__file__)
+project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
+for project_root in project_roots:
+    sys.path.insert(0, project_root) if project_root not in sys.path else None
+
+from videox_fun.dist import set_multi_gpus_devices, shard_model
+from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
+                               WanT5EncoderModel,
+                               WanTransformer3DModel_SelfForcing)
+from videox_fun.pipeline import WanSelfForcingPipeline
+from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
+                              FlowUniPCMultistepScheduler,
+                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
+                              save_videos_grid, unmerge_lora)
+
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# model_full_load means that the entire model will be moved to the GPU.
+#
+# model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
+# and the transformer model has been quantized to float8, which can save more GPU memory.
+#
+# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
+#
+# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use,
+# and the transformer model has been quantized to float8, which can save more GPU memory.
+#
+# model_group_offload transfers internal layer groups between CPU/CUDA,
+# balancing memory efficiency and speed between full-module and leaf-level offloading methods.
+#
+# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use,
+# resulting in slower speeds but saving a large amount of GPU memory.
+GPU_memory_mode     = "model_full_load"
+# Multi GPUs config
+# Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used.
+# For example, if you are using 8 GPUs, you can set ulysses_degree = 2 and ring_degree = 4.
+# If you are using 1 GPU, you can set ulysses_degree = 1 and ring_degree = 1.
+ulysses_degree      = 1
+ring_degree         = 1
+# Use FSDP to save more GPU memory in multi gpus.
+fsdp_dit            = False
+fsdp_text_encoder   = True
+# Compile will give a speedup in fixed resolution and need a little GPU memory.
+# The compile_dit is not compatible with the fsdp_dit and sequential_cpu_offload.
+compile_dit         = False
+
+# Config and model path
+config_path         = "config/wan2.1/wan_civitai.yaml"
+# One-Forcing distills a Wan2.1-T2V-1.3B student (the 14B model is only the DMD teacher at training time).
+model_name          = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+
+# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++".
+sampler_name        = "Flow"
+# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics.
+# Used when the sampler is in "Flow_Unipc", "Flow_DPM++".
+shift               = 5.0
+# One-Forcing is a DMD-distilled student, so it uses stochastic sampling.
+stochastic_sampling = True
+
+# One-Forcing checkpoint to overlay on top of the Wan2.1 base model. Point at the trainer
+# output checkpoint dir (its live/EMA weights are auto-resolved below) or a plain `.pt`/`.safetensors` file.
+# Default: the official One-Forcing checkpoint, downloaded via
+#   hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+# The GAN discriminator head is a training-only submodule of fake_score, so it never appears in
+# the generator checkpoint and inference needs no discriminator wiring.
+transformer_path    = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
+# Load the EMA generator (transformer_ema/) when True, otherwise the live generator weights.
+# Only applies to trainer-output checkpoint dirs; a plain `.pt`/`.safetensors` file is loaded as-is.
+use_ema             = False
+vae_path            = None
+lora_path           = None
+
+# Other params
+sample_size         = [480, 832]
+video_length        = 81
+fps                 = 16
+
+# One-Forcing causal inference config
+# `num_frame_per_block`: 1 = frame-wise (One-Forcing default), 3 = chunk-wise.
+num_frame_per_block     = 1
+# Local attention window size (-1 for global attention)
+local_attn_size         = -1
+# Others
+independent_first_frame = False
+context_noise           = 0.0
+
+# Use torch.float16 if GPU does not support torch.bfloat16
+# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+weight_dtype        = torch.bfloat16
+prompt              = "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about."
+negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+# -------- One-Forcing step selector --------
+# One-Forcing bakes CFG into the student weights, so inference MUST use guidance_scale=1.0.
+# The official recipe generates each frame in a single step (denoising_step_list=[1000]):
+#   1-step (One-Forcing default): num_inference_steps = 1
+#   4-step (more stable, slower): num_inference_steps = 4 -> [1000, 750, 500, 250]
+# With stochastic_sampling=True, num_inference_steps=1 yields the warped single timestep [1000],
+# matching the official framewise 1-step generation.
+guidance_scale      = 1.0
+num_inference_steps = 1
+seed                = 43
+lora_weight         = 0.55
+save_path           = "samples/wan-videos-one-forcing"
+
+device = set_multi_gpus_devices(ulysses_degree, ring_degree)
+config = OmegaConf.load(config_path)
+
+# Load transformer with causal inference support if enabled
+transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs'])
+transformer_additional_kwargs['local_attn_size'] = local_attn_size
+
+transformer = WanTransformer3DModel_SelfForcing.from_pretrained(
+    os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+    transformer_additional_kwargs=transformer_additional_kwargs,
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+
+if transformer_path is not None:
+    def _resolve_transformer_path(raw_path: str, prefer_ema: bool) -> str:
+        """Resolve a checkpoint path to the actual weights file to load.
+
+        File paths are returned unchanged so external `.pt` ckpts keep working. For trainer-output
+        dirs, prefer the EMA generator when asked and available, otherwise fall back to the live
+        generator. The FSDP trainer writes the EMA under `transformer_ema/`; `ema_transformer/` is
+        kept as a compatibility alias.
+        """
+        if os.path.isfile(raw_path):
+            return raw_path
+        if os.path.isdir(raw_path):
+            candidates = []
+            if prefer_ema:
+                candidates.append(os.path.join(raw_path, "transformer_ema", "diffusion_pytorch_model.safetensors"))
+                candidates.append(os.path.join(raw_path, "ema_transformer", "diffusion_pytorch_model.safetensors"))
+            candidates.append(os.path.join(raw_path, "transformer", "diffusion_pytorch_model.safetensors"))
+            candidates.append(os.path.join(raw_path, "diffusion_pytorch_model.safetensors"))
+            for c in candidates:
+                if os.path.isfile(c):
+                    return c
+        raise FileNotFoundError(
+            f"transformer_path={raw_path!r} is neither a file nor a checkpoint dir "
+            f"with a known safetensors layout (transformer_ema/, transformer/ or root)."
+        )
+
+    _raw_transformer_path = transformer_path
+    transformer_path = _resolve_transformer_path(transformer_path, prefer_ema=use_ema)
+    if transformer_path != _raw_transformer_path:
+        print(f"use_ema={use_ema}: resolved {_raw_transformer_path} -> {transformer_path}")
+    print(f"From checkpoint: {transformer_path}")
+    if transformer_path.endswith("safetensors"):
+        from safetensors.torch import load_file
+        state_dict = load_file(transformer_path)
+    else:
+        state_dict = torch.load(transformer_path, map_location="cpu")
+
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+    state_dict = state_dict["generator_ema"] if "generator_ema" in state_dict else state_dict
+    state_dict = state_dict["generator"] if "generator" in state_dict else state_dict
+    # FSDP-saved ckpts keep the `model._fsdp_wrapped_module.` prefix; strip it before the generic
+    # `model.` strip so both FSDP and plain ckpts land at bare parameter names.
+    if any("._fsdp_wrapped_module." in k for k in state_dict.keys()):
+        state_dict = {k.replace("model._fsdp_wrapped_module.", "model.", 1) if k.startswith("model._fsdp_wrapped_module.") else k: v for k, v in state_dict.items()}
+    if any(k.startswith("model.") for k in state_dict.keys()):
+        state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()}
+
+    m, u = transformer.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# Get Vae
+vae = AutoencoderKLWan.from_pretrained(
+    os.path.join(model_name, config['vae_kwargs'].get('vae_subpath', 'vae')),
+    additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
+).to(weight_dtype)
+
+if vae_path is not None:
+    print(f"From checkpoint: {vae_path}")
+    if vae_path.endswith("safetensors"):
+        from safetensors.torch import load_file
+        state_dict = load_file(vae_path)
+    else:
+        state_dict = torch.load(vae_path, map_location="cpu")
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+
+    m, u = vae.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# Get Tokenizer
+tokenizer = AutoTokenizer.from_pretrained(
+    os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
+)
+
+# Get Text encoder
+text_encoder = WanT5EncoderModel.from_pretrained(
+    os.path.join(model_name, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
+    additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+
+# Get Scheduler
+Chosen_Scheduler = {
+    "Flow": FlowMatchEulerDiscreteScheduler,
+    "Flow_Unipc": FlowUniPCMultistepScheduler,
+    "Flow_DPM++": FlowDPMSolverMultistepScheduler,
+}[sampler_name]
+if sampler_name == "Flow_Unipc" or sampler_name == "Flow_DPM++":
+    config['scheduler_kwargs']['shift'] = 1
+scheduler = Chosen_Scheduler(
+    **filter_kwargs(Chosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
+)
+
+# Get Pipeline
+pipeline = WanSelfForcingPipeline(
+    transformer=transformer,
+    vae=vae,
+    tokenizer=tokenizer,
+    text_encoder=text_encoder,
+    scheduler=scheduler,
+)
+
+if ulysses_degree > 1 or ring_degree > 1:
+    from functools import partial
+    transformer.enable_multi_gpus_inference()
+    if fsdp_dit:
+        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
+        pipeline.transformer = shard_fn(pipeline.transformer)
+        print("Add FSDP DIT")
+    if fsdp_text_encoder:
+        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
+        pipeline.text_encoder = shard_fn(pipeline.text_encoder)
+        print("Add FSDP TEXT ENCODER")
+
+if compile_dit:
+    for i in range(len(pipeline.transformer.blocks)):
+        pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
+    print("Add Compile")
+
+# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
+# The order lives inside the helper: quantization has to happen before the offload hooks
+# are installed.
+apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+
+generator = torch.Generator(device=device).manual_seed(seed)
+
+if lora_path is not None:
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
+
+with torch.no_grad():
+    video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
+    latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1
+
+    sample = pipeline(
+        prompt,
+        num_frames = video_length,
+        negative_prompt = negative_prompt,
+        height      = sample_size[0],
+        width       = sample_size[1],
+        generator   = generator,
+        guidance_scale          = guidance_scale,
+        num_inference_steps     = num_inference_steps,
+        shift                   = shift,
+        num_frame_per_block     = num_frame_per_block,
+        independent_first_frame = independent_first_frame,
+        context_noise           = context_noise,
+        stochastic_sampling     = stochastic_sampling,
+    ).videos
+
+if lora_path is not None:
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
+
+def save_results():
+    if not os.path.exists(save_path):
+        os.makedirs(save_path, exist_ok=True)
+
+    index = len([path for path in os.listdir(save_path)]) + 1
+    prefix = str(index).zfill(8)
+    if video_length == 1:
+        video_path = os.path.join(save_path, prefix + ".png")
+
+        image = sample[0, :, 0]
+        image = image.transpose(0, 1).transpose(1, 2)
+        image = (image * 255).numpy().astype(np.uint8)
+        image = Image.fromarray(image)
+        image.save(video_path)
+        print(f"Saved image to: {video_path}")
+    else:
+        video_path = os.path.join(save_path, prefix + ".mp4")
+        save_videos_grid(sample, video_path, fps=fps)
+
+if ulysses_degree * ring_degree > 1:
+    import torch.distributed as dist
+    if dist.get_rank() == 0:
+        save_results()
+else:
+    save_results()

scripts/taehv/README_TRAIN.md

@@ -180,25 +180,32 @@ NCCL_DEBUG=INFO
 
 accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
   --config_path="config/wan2.2/wan_civitai_5b.yaml" \
-  --pretrained_model_name_or_path=$MODEL_NAME \
-  --tae_path=$TAE_PATH \
-  --train_data_dir=$DATASET_NAME \
-  --train_data_meta=$DATASET_META_NAME \
-  --video_sample_size=512 \
-  --video_sample_stride=1 \
-  --video_sample_n_frames=33 \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --tae_path="$TAE_PATH" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --fix_sample_size 720 1280 \
+  --video_sample_strides 1 2 3 \
+  --video_sample_n_frames=25 \
+  --validation_n_frames=49 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=1 \
-  --dataloader_num_workers=4 \
-  --num_train_epochs=100 \
-  --checkpointing_steps=500 \
+  --dataloader_num_workers=8 \
+  --max_train_steps=10000 \
+  --checkpointing_steps=1000 \
+  --validation_steps=1000 \
+  --validation_paths "asset/inpaint_video.mp4" \
   --learning_rate=1e-04 \
   --lr_scheduler="constant_with_warmup" \
-  --lr_warmup_steps=100 \
+  --lr_warmup_steps=200 \
   --latent_loss_weight=1.0 \
+  --use_ema \
+  --ema_decay=0.999 \
   --seed=42 \
   --output_dir="output_dir_taehv_w2.2" \
   --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
   --adam_weight_decay=1e-2 \
   --adam_epsilon=1e-08 \
   --vae_mini_batch=1 \
@@ -208,6 +215,8 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
   --resume_from_checkpoint=latest
 ```
 
+These settings are also provided in `scripts/taehv/train_taehv.sh`. Edit its training arguments and run `bash scripts/taehv/train_taehv.sh`.
+
 **Wan2.1 / Wan2.2 14B (taew2_1) Example**:
 
 Same as above, with these changes:
@@ -246,6 +255,7 @@ export TAE_PATH="taew2_1.safetensors"
 | `--gradient_accumulation_steps` | Gradient accumulation steps | 1 |
 | `--dataloader_num_workers` | Number of DataLoader workers | 4 |
 | `--num_train_epochs` | Number of training epochs | 100 |
+| `--max_train_steps` | Total training steps; overrides `--num_train_epochs` when set | 10000 |
 | `--checkpointing_steps` | Save a checkpoint every N steps | 500 |
 | `--checkpoints_total_limit` | Max number of checkpoints to store | `None` |
 | `--learning_rate` | Initial learning rate | 1e-4 |
@@ -265,7 +275,7 @@ export TAE_PATH="taew2_1.safetensors"
 | `--trainable_modules_low_learning_rate` | Trainable modules with lr/2 | `[]` |
 | `--resume_from_checkpoint` | Resume training from checkpoint, use `"latest"` to auto-select | `latest` |
 | `--validation_steps` / `--validation_epochs` | Run validation every N steps / epochs | 2000 / 5 |
-| `--validation_paths` | Video paths for validation reconstruction comparison | `"asset/1.mp4"` |
+| `--validation_paths` | Video paths for validation reconstruction comparison | `"asset/inpaint_video.mp4"` |
 | `--validation_n_frames` | Validation clip length (4k+1); defaults to `--video_sample_n_frames` | `None` |
 
 **Sample Size Configuration Guide**:
@@ -282,11 +292,11 @@ You can configure validation parameters to periodically reconstruct test videos
 |-----------|-------------|-------------------|
 | `--validation_steps` | Run validation every N steps | 2000 |
 | `--validation_epochs` | Run validation every N epochs | 5 |
-| `--validation_paths` | Validation video paths | `"asset/1.mp4"` |
+| `--validation_paths` | Validation video paths | `"asset/inpaint_video.mp4"` |
 | `--validation_n_frames` | Validation clip length, **4k+1**; omit to reuse `--video_sample_n_frames` | 81 |
 
 ```bash
-  --validation_paths "asset/1.mp4" \
+  --validation_paths "asset/inpaint_video.mp4" \
   --validation_steps=2000 \
   --validation_epochs=5 \
   --validation_n_frames=81
@@ -304,7 +314,7 @@ You can configure validation parameters to periodically reconstruct test videos
 - **Decoder-only distillation**: Add `--freeze_tae_encoder` to only improve decoding quality.
 - **Memory**: The teacher full VAE is the main memory consumer. Use `--low_vram` to keep it on CPU between encoding steps, or reduce `--video_sample_n_frames` / `--video_sample_size`.
 - **Sequential TAE**: `--use_taehv_sequential` trades speed for O(1) activation memory w.r.t. video length.
-- **EMA**: `--use_ema` is recommended for the final deliverable; the final saved `taehv` directory contains the EMA weights. Pass `--ema_decay` (e.g. `0.999`, as `train_taehv_super.sh` does) when short runs need the average to catch up faster than the 0.9999 default.
+- **EMA**: `--use_ema` is recommended for the final deliverable; the final saved `taehv` directory contains the EMA weights. Pass `--ema_decay` (e.g. `0.999`, as both `train_taehv.sh` and `train_taehv_super.sh` do) when short runs need the average to catch up faster than the 0.9999 default.
 
 ### 3.6 Multi-Node Distributed Training
 

scripts/taehv/README_TRAIN_SUPER.md

@@ -186,14 +186,16 @@ LPIPS and DINOv2 weights are loaded separately by the loss implementation.
 
 **Wan2.2 TI2V-5B / Fun-2.2VAE (taew2_2_super) Example**:
 
-The launcher uses a 720x1280 canvas, 25 training frames, 81 validation frames,
-10,000 updates, and EMA. Set the GPU/process count with `accelerate config`;
-four GPUs with batch size 1 and accumulation 1 give an effective batch of 4.
+The launcher uses a 720x1280 canvas, 25 training frames, and EMA, with no
+`--max_train_steps`: training runs for `--num_train_epochs` (100 by default)
+epochs until you kill the processes. Set the GPU/process count with
+`accelerate config`; four GPUs with batch size 1 and accumulation 1 give an
+effective batch of 4.
 
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
-export DATASET_NAME="${DATASET_NAME:-datasets/X-Fun-Videos-Demo}"
-export DATASET_META_NAME="${DATASET_META_NAME:-$DATASET_NAME/metadata_add_width_height.json}"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth"
 NCCL_DEBUG=INFO
 
@@ -206,18 +208,23 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --fix_sample_size 720 1280 \
   --video_sample_strides 1 2 3 \
   --video_sample_n_frames=25 \
-  --validation_n_frames=81 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
-  --max_train_steps=10000 \
   --checkpointing_steps=1000 \
-  --validation_steps=1000 \
-  --validation_max_samples=0 \
   --learning_rate=5e-06 \
   --lr_scheduler="cosine" \
   --lr_warmup_steps=200 \
   --lr_min_ratio=0.2 \
+  --seed=42 \
+  --output_dir="output_dir_taehv_w2.2_super" \
+  --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
+  --adam_weight_decay=1e-2 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=1.0 \
   --l1_loss_weight=0.1 \
   --lpips_loss_weight=0.1 \
   --dino_loss_weight=0.02 \
@@ -227,16 +234,7 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --use_ema \
   --ema_decay=0.999 \
   --selection_lpips_ratio=1.02 \
-  --selection_grid_ratio=1.05 \
-  --seed=42 \
-  --output_dir="output_dir_taehv_w2.2_super" \
-  --mixed_precision="bf16" \
-  --adam_beta1=0.9 \
-  --adam_beta2=0.99 \
-  --adam_weight_decay=1e-2 \
-  --adam_epsilon=1e-08 \
-  --vae_mini_batch=1 \
-  --max_grad_norm=1.0
+  --selection_grid_ratio=1.05
 ```
 
 These settings are also provided in `scripts/taehv/train_taehv_super.sh`. Edit its
@@ -253,10 +251,9 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --fix_sample_size 720 1280 \
   --video_sample_n_frames=25 \
   --video_sample_strides 1 2 3 \
-  --max_train_steps=10000 \
   --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
   --use_ema --ema_decay=0.999 \
-  --dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \
+  --dataloader_num_workers=8 --checkpointing_steps=1000 \
   --output_dir="output_dir_taehv_w2.2_super_scratch"
 ```
 
@@ -271,19 +268,20 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --train_data_dir="$DATASET_NAME" \
   --train_data_meta="$DATASET_META_NAME" \
   --fix_sample_size 720 1280 --video_sample_n_frames=25 \
-  --video_sample_strides 1 2 3 --max_train_steps=10000 \
+  --video_sample_strides 1 2 3 \
   --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
   --use_ema --ema_decay=0.999 \
-  --dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \
+  --dataloader_num_workers=8 --checkpointing_steps=1000 \
   --output_dir="output_dir_taehv_w2.2_super" \
   --resume_from_checkpoint=latest
 ```
 
-Restore with the same settings used to create the checkpoint, including validation
-metadata when supplied. `--max_train_steps` is the total update budget. Keep the
-original schedule horizon. The checkpoint restores optimizer, scheduler, EMA, RNG,
-step and batch cursor; changing workers or sampling settings can change the resumed
-sample sequence. Use a new output directory for a new run.
+Restore with the same settings used to create the checkpoint. Without
+`--max_train_steps`, the schedule horizon derives from `--num_train_epochs`;
+keep the original value when resuming. The checkpoint restores optimizer,
+scheduler, EMA, RNG, step and batch cursor; changing workers or sampling
+settings can change the resumed sample sequence. Use a new output directory
+for a new run.
 
 ### 3.3 Training Parameter Reference
 
@@ -302,7 +300,7 @@ Values below describe the launcher configuration; they are not all Python CLI de
 | `--train_batch_size` / `--gradient_accumulation_steps` | Batch per device / gradient accumulation | 1 / 1 |
 | `--vae_mini_batch` | Official VAE encoding mini-batch size | 1 |
 | `--dataloader_num_workers` | DataLoader workers | 8 |
-| `--max_train_steps` | Total optimizer updates; overrides `--num_train_epochs` | 10000 |
+| `--num_train_epochs` | Epoch budget used when `--max_train_steps` is not set | 100 |
 | `--checkpointing_steps` / `--checkpoints_total_limit` | Save interval / checkpoint count limit | 1000 / `None` |
 | `--learning_rate` / `--lr_scheduler` | Peak learning rate / scheduler | `5e-6` / `cosine` |
 | `--lr_warmup_steps` / `--lr_min_ratio` | Warmup updates / final-to-peak ratio for cosine decay | 200 / 0.2 |
@@ -317,8 +315,8 @@ Values below describe the launcher configuration; they are not all Python CLI de
 **Sample Size Configuration Guide**:
 
 - `--fix_sample_size 720 1280` resizes to cover the canvas and center crops while preserving aspect ratio. Without it, `--video_sample_size` selects a square canvas.
-- A 25-frame clip at stride 3 needs at least 73 source frames; 81-frame validation needs at least 81. Short inputs shorten the batch to the shortest 4k+1 clip length, recorded as `frames` in the training log.
-- Training uses random temporal starts. Validation starts at frame 0 with stride 1. Data loading retains the original random replacement/retry behavior on decode failures.
+- A 25-frame clip at stride 3 needs at least 73 source frames. Short inputs shorten the batch to the shortest 4k+1 clip length, recorded as `frames` in the training log.
+- Training uses random temporal starts. Data loading retains the original random replacement/retry behavior on decode failures.
 
 **Loss Configuration Guide**:
 
@@ -328,9 +326,10 @@ Values below describe the launcher configuration; they are not all Python CLI de
 
 ### 3.4 Training Validation
 
-Provide a separate JSON array through `--validation_data_meta` in the training
-command, or add that argument to the launcher. Relative video paths resolve under
-`--train_data_dir`, so use absolute video paths when validation data has another root.
+Validation is optional and disabled in the default launcher. To enable it, add
+`--validation_data_meta` (plus the related validation arguments) to the training
+command. Relative video paths resolve under `--train_data_dir`, so use absolute
+video paths when validation data has another root.
 
 | Parameter | Description | Recommended Value |
 |-----------|-------------|-------------------|
@@ -359,7 +358,7 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --train_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
   --validation_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
   --fix_sample_size 64 96 --video_sample_n_frames=5 --validation_n_frames=5 \
-  --video_sample_strides 1 --max_train_steps=2 --lr_warmup_steps=0 \
+  --video_sample_strides 1 --num_train_epochs=1 --lr_warmup_steps=0 \
   --checkpointing_steps=1 --validation_steps=1 --validation_max_samples=1 \
   --dataloader_num_workers=0 --use_ema --ema_decay=0.999 \
   --output_dir="output_dir_taehv_w2.2_super_smoke"

scripts/taehv/README_TRAIN_SUPER_zh-CN.md

@@ -179,13 +179,14 @@ LPIPS 与 DINOv2 权重由损失实现单独加载。
 
 **Wan2.2 TI2V-5B / Fun-2.2VAE（taew2_2_super）训练示例**：
 
-启动器采用 720x1280 画布、训练 25 帧、验证 81 帧、10,000 步和 EMA。
+启动器采用 720x1280 画布、训练 25 帧和 EMA，不设置 `--max_train_steps`：
+训练按 `--num_train_epochs`（默认 100）个 epoch 运行，直到手动终止进程。
 使用 `accelerate config` 配置 GPU/进程数；四卡、每卡 batch=1、累积=1 时有效 batch 为 4。
 
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
-export DATASET_NAME="${DATASET_NAME:-datasets/X-Fun-Videos-Demo}"
-export DATASET_META_NAME="${DATASET_META_NAME:-$DATASET_NAME/metadata_add_width_height.json}"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth"
 NCCL_DEBUG=INFO
 
@@ -198,18 +199,23 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --fix_sample_size 720 1280 \
   --video_sample_strides 1 2 3 \
   --video_sample_n_frames=25 \
-  --validation_n_frames=81 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
-  --max_train_steps=10000 \
   --checkpointing_steps=1000 \
-  --validation_steps=1000 \
-  --validation_max_samples=0 \
   --learning_rate=5e-06 \
   --lr_scheduler="cosine" \
   --lr_warmup_steps=200 \
   --lr_min_ratio=0.2 \
+  --seed=42 \
+  --output_dir="output_dir_taehv_w2.2_super" \
+  --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
+  --adam_weight_decay=1e-2 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=1.0 \
   --l1_loss_weight=0.1 \
   --lpips_loss_weight=0.1 \
   --dino_loss_weight=0.02 \
@@ -219,16 +225,7 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --use_ema \
   --ema_decay=0.999 \
   --selection_lpips_ratio=1.02 \
-  --selection_grid_ratio=1.05 \
-  --seed=42 \
-  --output_dir="output_dir_taehv_w2.2_super" \
-  --mixed_precision="bf16" \
-  --adam_beta1=0.9 \
-  --adam_beta2=0.99 \
-  --adam_weight_decay=1e-2 \
-  --adam_epsilon=1e-08 \
-  --vae_mini_batch=1 \
-  --max_grad_norm=1.0
+  --selection_grid_ratio=1.05
 ```
 
 相同参数已写入 `scripts/taehv/train_taehv_super.sh`，也可修改其中的训练参数后，
@@ -245,10 +242,9 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --fix_sample_size 720 1280 \
   --video_sample_n_frames=25 \
   --video_sample_strides 1 2 3 \
-  --max_train_steps=10000 \
   --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
   --use_ema --ema_decay=0.999 \
-  --dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \
+  --dataloader_num_workers=8 --checkpointing_steps=1000 \
   --output_dir="output_dir_taehv_w2.2_super_scratch"
 ```
 
@@ -263,16 +259,16 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --train_data_dir="$DATASET_NAME" \
   --train_data_meta="$DATASET_META_NAME" \
   --fix_sample_size 720 1280 --video_sample_n_frames=25 \
-  --video_sample_strides 1 2 3 --max_train_steps=10000 \
+  --video_sample_strides 1 2 3 \
   --learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
   --use_ema --ema_decay=0.999 \
-  --dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \
+  --dataloader_num_workers=8 --checkpointing_steps=1000 \
   --output_dir="output_dir_taehv_w2.2_super" \
   --resume_from_checkpoint=latest
 ```
 
-恢复时保持创建 checkpoint 时的配置，使用过的验证索引也需继续提供。
-`--max_train_steps` 表示总更新预算，应保留原 schedule 预算。Checkpoint 恢复
+恢复时保持创建 checkpoint 时的配置。未设置 `--max_train_steps` 时，schedule
+预算由 `--num_train_epochs` 决定，恢复时应保留原值。Checkpoint 恢复
 optimizer、scheduler、EMA、RNG、step 和 batch 游标；改变 workers 或采样设置可能改变
 恢复后的样本顺序。新训练应使用新的输出目录。
 
@@ -293,7 +289,7 @@ optimizer、scheduler、EMA、RNG、step 和 batch 游标；改变 workers 或
 | `--train_batch_size` / `--gradient_accumulation_steps` | 每卡 batch / 梯度累积 | 1 / 1 |
 | `--vae_mini_batch` | 官方 VAE 编码 mini-batch | 1 |
 | `--dataloader_num_workers` | DataLoader workers | 8 |
-| `--max_train_steps` | 总优化器更新次数，覆盖 `--num_train_epochs` | 10000 |
+| `--num_train_epochs` | 未设置 `--max_train_steps` 时的 epoch 预算 | 100 |
 | `--checkpointing_steps` / `--checkpoints_total_limit` | 保存间隔 / checkpoint 数量上限 | 1000 / `None` |
 | `--learning_rate` / `--lr_scheduler` | 峰值学习率 / 调度器 | `5e-6` / `cosine` |
 | `--lr_warmup_steps` / `--lr_min_ratio` | 预热步数 / cosine 终点与峰值之比 | 200 / 0.2 |
@@ -308,8 +304,8 @@ optimizer、scheduler、EMA、RNG、step 和 batch 游标；改变 workers 或
 **Sample Size 配置指南**：
 
 - `--fix_sample_size 720 1280` 等比例缩放覆盖画布后中心裁剪；未指定时，`--video_sample_size` 控制正方形画布。
-- 25 帧、stride=3 至少需要 73 个源视频帧；81 帧验证至少需要 81 帧。短视频会使 batch 按最短视频截为 4k+1 帧，实际长度记录在训练日志的 `frames` 字段。
-- 训练随机选择时间起点，验证固定从第 0 帧开始、stride=1。解码失败时沿用原有随机换视频重试逻辑。
+- 25 帧、stride=3 至少需要 73 个源视频帧。短视频会使 batch 按最短视频截为 4k+1 帧，实际长度记录在训练日志的 `frames` 字段。
+- 训练随机选择时间起点。解码失败时沿用原有随机换视频重试逻辑。
 
 **损失配置指南**：
 
@@ -319,8 +315,9 @@ optimizer、scheduler、EMA、RNG、step 和 batch 游标；改变 workers 或
 
 ### 3.4 训练验证
 
-在训练命令中通过 `--validation_data_meta` 提供独立 JSON 数组，或将该参数加入启动器。
-相对视频路径基于 `--train_data_dir` 解析；验证数据位于其他根目录时，可使用绝对视频路径。
+验证为可选功能，默认启动器中未开启。需要时在训练命令中通过 `--validation_data_meta`
+（及相关验证参数）启用。相对视频路径基于 `--train_data_dir` 解析；验证数据位于其他
+根目录时，可使用绝对视频路径。
 
 | 参数 | 说明 | 推荐值 |
 |------|------|--------|
@@ -347,7 +344,7 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --train_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
   --validation_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
   --fix_sample_size 64 96 --video_sample_n_frames=5 --validation_n_frames=5 \
-  --video_sample_strides 1 --max_train_steps=2 --lr_warmup_steps=0 \
+  --video_sample_strides 1 --num_train_epochs=1 --lr_warmup_steps=0 \
   --checkpointing_steps=1 --validation_steps=1 --validation_max_samples=1 \
   --dataloader_num_workers=0 --use_ema --ema_decay=0.999 \
   --output_dir="output_dir_taehv_w2.2_super_smoke"

scripts/taehv/README_TRAIN_zh-CN.md

@@ -180,25 +180,32 @@ NCCL_DEBUG=INFO
 
 accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
   --config_path="config/wan2.2/wan_civitai_5b.yaml" \
-  --pretrained_model_name_or_path=$MODEL_NAME \
-  --tae_path=$TAE_PATH \
-  --train_data_dir=$DATASET_NAME \
-  --train_data_meta=$DATASET_META_NAME \
-  --video_sample_size=512 \
-  --video_sample_stride=1 \
-  --video_sample_n_frames=33 \
+  --pretrained_model_name_or_path="$MODEL_NAME" \
+  --tae_path="$TAE_PATH" \
+  --train_data_dir="$DATASET_NAME" \
+  --train_data_meta="$DATASET_META_NAME" \
+  --fix_sample_size 720 1280 \
+  --video_sample_strides 1 2 3 \
+  --video_sample_n_frames=25 \
+  --validation_n_frames=49 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=1 \
-  --dataloader_num_workers=4 \
-  --num_train_epochs=100 \
-  --checkpointing_steps=500 \
+  --dataloader_num_workers=8 \
+  --max_train_steps=10000 \
+  --checkpointing_steps=1000 \
+  --validation_steps=1000 \
+  --validation_paths "asset/inpaint_video.mp4" \
   --learning_rate=1e-04 \
   --lr_scheduler="constant_with_warmup" \
-  --lr_warmup_steps=100 \
+  --lr_warmup_steps=200 \
   --latent_loss_weight=1.0 \
+  --use_ema \
+  --ema_decay=0.999 \
   --seed=42 \
   --output_dir="output_dir_taehv_w2.2" \
   --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
   --adam_weight_decay=1e-2 \
   --adam_epsilon=1e-08 \
   --vae_mini_batch=1 \
@@ -208,6 +215,8 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv.py \
   --resume_from_checkpoint=latest
 ```
 
+相同参数已写入 `scripts/taehv/train_taehv.sh`，也可修改其中的训练参数后，运行 `bash scripts/taehv/train_taehv.sh`。
+
 **Wan2.1 / Wan2.2 14B（taew2_1）训练示例**：
 
 与上面相同，仅做如下修改：
@@ -246,6 +255,7 @@ export TAE_PATH="taew2_1.safetensors"
 | `--gradient_accumulation_steps` | 梯度累积步数 | 1 |
 | `--dataloader_num_workers` | DataLoader 子进程数 | 4 |
 | `--num_train_epochs` | 训练 epoch 数 | 100 |
+| `--max_train_steps` | 总训练步数；设置后覆盖 `--num_train_epochs` | 10000 |
 | `--checkpointing_steps` | 每 N 步保存 checkpoint | 500 |
 | `--checkpoints_total_limit` | 最多保留的 checkpoint 数量 | `None` |
 | `--learning_rate` | 初始学习率 | 1e-4 |
@@ -265,7 +275,7 @@ export TAE_PATH="taew2_1.safetensors"
 | `--trainable_modules_low_learning_rate` | 以 lr/2 训练的模块 | `[]` |
 | `--resume_from_checkpoint` | 恢复训练路径，使用 `"latest"` 自动选择最新 checkpoint | `latest` |
 | `--validation_steps` / `--validation_epochs` | 每 N 步 / 每 N 个 epoch 执行一次验证 | 2000 / 5 |
-| `--validation_paths` | 用于重建对比验证的视频路径 | `"asset/1.mp4"` |
+| `--validation_paths` | 用于重建对比验证的视频路径 | `"asset/inpaint_video.mp4"` |
 | `--validation_n_frames` | 验证片段帧数（4k+1），不传则复用 `--video_sample_n_frames` | `None` |
 
 **Sample Size 配置指南**：
@@ -282,11 +292,11 @@ export TAE_PATH="taew2_1.safetensors"
 |------|------|--------|
 | `--validation_steps` | 每 N 步执行一次验证 | 2000 |
 | `--validation_epochs` | 每 N 个 epoch 执行一次验证 | 5 |
-| `--validation_paths` | 验证视频路径 | `"asset/1.mp4"` |
+| `--validation_paths` | 验证视频路径 | `"asset/inpaint_video.mp4"` |
 | `--validation_n_frames` | 验证片段帧数，**必须为 4k+1**；不传则复用 `--video_sample_n_frames` | 81 |
 
 ```bash
-  --validation_paths "asset/1.mp4" \
+  --validation_paths "asset/inpaint_video.mp4" \
   --validation_steps=2000 \
   --validation_epochs=5 \
   --validation_n_frames=81
@@ -304,7 +314,7 @@ export TAE_PATH="taew2_1.safetensors"
 - **Decoder-only 蒸馏**：加上 `--freeze_tae_encoder` 只提升解码质量。
 - **显存**：teacher 完整 VAE 是主要的显存消耗者。使用 `--low_vram` 可以将其放在 CPU、仅在编码时搬上 GPU；也可以降低 `--video_sample_n_frames` / `--video_sample_size`。
 - **串行 TAE**：`--use_taehv_sequential` 以速度换显存，激活显存与视频长度无关（O(1)）。
-- **EMA**：推荐开启 `--use_ema` 作为最终交付物；最终保存的 `taehv` 目录里是 EMA 权重。短训练希望均值更快跟上时，可传 `--ema_decay`（如 `train_taehv_super.sh` 用的 `0.999`）替代默认的 0.9999。
+- **EMA**：推荐开启 `--use_ema` 作为最终交付物；最终保存的 `taehv` 目录里是 EMA 权重。短训练希望均值更快跟上时，可传 `--ema_decay`（如 `train_taehv.sh` 与 `train_taehv_super.sh` 用的 `0.999`）替代默认的 0.9999。
 
 ### 3.6 多机分布式训练
 

scripts/taehv/train_taehv_super.py

@@ -614,12 +614,8 @@ def main():
     if args.allow_tf32:
         torch.backends.cuda.matmul.allow_tf32 = True
     if accelerator.is_main_process:
-        os.makedirs(args.output_dir, exist_ok=True)
-        if not args.resume_from_checkpoint and (
-            os.path.exists(os.path.join(args.output_dir, "train.jsonl")) or
-            any(name.startswith("checkpoint-") for name in os.listdir(args.output_dir))
-        ):
-            raise FileExistsError("Use a new output_dir or --resume_from_checkpoint")
+        if args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
 
     # For mixed precision training, keep trainable parameters and Adam moments FP32.
     weight_dtype = torch.float32

scripts/taehv/train_taehv_super.sh

@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
-export DATASET_NAME="${DATASET_NAME:-datasets/X-Fun-Videos-Demo}"
-export DATASET_META_NAME="${DATASET_META_NAME:-$DATASET_NAME/metadata_add_width_height.json}"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
 # export NCCL_IB_DISABLE=1
@@ -17,18 +17,23 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --fix_sample_size 720 1280 \
   --video_sample_strides 1 2 3 \
   --video_sample_n_frames=25 \
-  --validation_n_frames=81 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
-  --max_train_steps=10000 \
   --checkpointing_steps=1000 \
-  --validation_steps=1000 \
-  --validation_max_samples=0 \
   --learning_rate=5e-06 \
   --lr_scheduler="cosine" \
   --lr_warmup_steps=200 \
   --lr_min_ratio=0.2 \
+  --seed=42 \
+  --output_dir="output_dir_taehv_w2.2_super" \
+  --mixed_precision="bf16" \
+  --adam_beta1=0.9 \
+  --adam_beta2=0.99 \
+  --adam_weight_decay=1e-2 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=1.0 \
   --l1_loss_weight=0.1 \
   --lpips_loss_weight=0.1 \
   --dino_loss_weight=0.02 \
@@ -38,13 +43,4 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
   --use_ema \
   --ema_decay=0.999 \
   --selection_lpips_ratio=1.02 \
-  --selection_grid_ratio=1.05 \
-  --seed=42 \
-  --output_dir="output_dir_taehv_w2.2_super" \
-  --mixed_precision="bf16" \
-  --adam_beta1=0.9 \
-  --adam_beta2=0.99 \
-  --adam_weight_decay=1e-2 \
-  --adam_epsilon=1e-08 \
-  --vae_mini_batch=1 \
-  --max_grad_norm=1.0
+  --selection_grid_ratio=1.05

scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md

@@ -173,7 +173,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --num_train_epochs=100 \
   --checkpointing_steps=200 \
   --learning_rate=2.0e-06 \
-  --learning_rate_critic=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
@@ -192,8 +192,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --enable_bucket \
   --num_frame_per_block=3 \
   --use_kv_cache_training \
-  --denoising_step_indices_list 1000 667 334 1 \
-  --real_guidance_scale=6.0 \
+  --use_ema \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --randomize_step_indices \
   --fake_guidance_scale=0.0 \
   --gen_update_interval=5 \
@@ -222,7 +224,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--num_train_epochs` | Number of training epochs | 100 |
 | `--checkpointing_steps` | Save checkpoint every N steps | 200 |
 | `--learning_rate` | Generator learning rate | 2e-06 |
-| `--learning_rate_critic` | Critic learning rate | 2e-06 |
+| `--learning_rate_critic` | Critic learning rate | 4.0e-07 |
 | `--lr_scheduler` | LR scheduler type | `constant_with_warmup` |
 | `--lr_warmup_steps` | LR warmup steps | 100 |
 | `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_dmd` |
@@ -235,6 +237,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--max_grad_norm` | Gradient clipping threshold | 10.0 |
 | `--trainable_modules` | Trainable modules (`"."` = all) | `"."` |
 | `--low_vram` | Enable low VRAM mode (offload VAE/text encoder) | - |
+| `--use_ema` | Maintain an EMA copy of the generator weights | - |
 
 **Video Sampling Parameters**:
 
@@ -255,8 +258,8 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 
 | Parameter | Description | Example Value |
 |-----------|-------------|---------------|
-| `--denoising_step_indices_list` | Denoising step indices (DMD core param). Example runs 4-step; parser default `[1000, 500]` = 2-step | `1000 667 334 1` |
-| `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 6.0 |
+| `--denoising_step_indices_list` | Denoising step indices (DMD core param). Example runs 4-step; parser default `[1000, 500]` = 2-step | `1000 750 500 250` |
+| `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 4.0 |
 | `--randomize_step_indices` | Randomize the denoising step indices during training | - |
 | `--fake_guidance_scale` | CFG scale for the fake-score (generator). 0.0 = no CFG | 0.0 |
 | `--gen_update_interval` | Generator update interval (generator updates every N critic steps) | 5 |
@@ -273,7 +276,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 
 ### 5.4 Optional One-Forcing GAN and Precision Controls
 
-These options reuse the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms in `scripts/wan2.1_self_forcing/train_distill.py`. All GAN/R1/R2 weights default to **0**, and both precision flags default to **off** in argparse. The AR diffusion, CCD, and DMD launchers enable `--fsdp_reduce_fp32`; GAN/R1/R2 and `--score_timestep_clamp` remain off. AR diffusion and CCD do not gain GAN losses.
+These options reuse the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms in `scripts/wan2.1_self_forcing/train_distill.py`. All GAN/R1/R2 weights default to **0**, and both precision flags default to **off** in argparse. The AR diffusion, CCD, and DMD launchers enable `--fsdp_reduce_fp32`; GAN/R1/R2 remain off, and the DMD launcher additionally enables `--score_timestep_clamp`. AR diffusion and CCD do not gain GAN losses.
 
 The discriminator shares the **fake-score backbone**, with one learnable register token per selected feature layer. Its head is optimized by the critic optimizer, not the generator optimizer. Generator loss is `DMD + weighted GAN`; critic loss is `denoising + weighted GAN + R1 + R2`.
 
@@ -290,7 +293,7 @@ The discriminator shares the **fake-score backbone**, with one learnable registe
 | `--r2_weight` | Normalized finite-difference penalty on noisy generated latents | `0.0` |
 | `--r2_sigma` | R2 perturbation standard deviation | `0.01` |
 | `--fsdp_reduce_fp32` | FSDP FP32 gradient reduction; parameter gathering stays low precision | off (enabled in the launcher) |
-| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off |
+| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off (enabled in the DMD launcher) |
 
 For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` for the generator and `softplus(-D_real) + softplus(D_fake)` for the discriminator. Relativistic losses instead use `softplus(-(D_fake - D_real))` and `softplus(-(D_real - D_fake))`.
 

scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md

@@ -173,7 +173,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --num_train_epochs=100 \
   --checkpointing_steps=200 \
   --learning_rate=2.0e-06 \
-  --learning_rate_critic=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
@@ -192,8 +192,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --enable_bucket \
   --num_frame_per_block=3 \
   --use_kv_cache_training \
-  --denoising_step_indices_list 1000 667 334 1 \
-  --real_guidance_scale=6.0 \
+  --use_ema \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --randomize_step_indices \
   --fake_guidance_scale=0.0 \
   --gen_update_interval=5 \
@@ -222,7 +224,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--num_train_epochs` | 训练 epoch 数 | 100 |
 | `--checkpointing_steps` | 每 N 步保存 checkpoint | 200 |
 | `--learning_rate` | 生成器学习率 | 2e-06 |
-| `--learning_rate_critic` | 判别器学习率 | 2e-06 |
+| `--learning_rate_critic` | 判别器学习率 | 4.0e-07 |
 | `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
 | `--lr_warmup_steps` | 学习率预热步数 | 100 |
 | `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_dmd` |
@@ -235,6 +237,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--max_grad_norm` | 梯度裁剪阈值 | 10.0 |
 | `--trainable_modules` | 可训练模块（`"."` = 全部） | `"."` |
 | `--low_vram` | 低显存模式（卸载 VAE/文本编码器） | - |
+| `--use_ema` | 维护生成器权重的 EMA 副本 | - |
 
 **视频采样参数**：
 
@@ -255,8 +258,8 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 
 | 参数 | 说明 | 示例值 |
 |------|------|--------|
-| `--denoising_step_indices_list` | 去噪步骤索引（DMD 核心参数）。示例为 4 步；解析器默认 `[1000, 500]` = 2 步 | `1000 667 334 1` |
-| `--real_guidance_scale` | real-score（14B teacher）的 CFG scale | 6.0 |
+| `--denoising_step_indices_list` | 去噪步骤索引（DMD 核心参数）。示例为 4 步；解析器默认 `[1000, 500]` = 2 步 | `1000 750 500 250` |
+| `--real_guidance_scale` | real-score（14B teacher）的 CFG scale | 4.0 |
 | `--randomize_step_indices` | 训练时随机化去噪步骤索引 | - |
 | `--fake_guidance_scale` | fake-score（生成器）的 CFG scale。0.0 = 无 CFG | 0.0 |
 | `--gen_update_interval` | 生成器更新间隔（每 N 步判别器更新后更新 1 次生成器） | 5 |
@@ -272,7 +275,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 
 ### 5.4 可选 One-Forcing GAN 与精度控制
 
-这些选项复用 `scripts/wan2.1_self_forcing/train_distill.py` 中的 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 机制。GAN/R1/R2 权重默认均为 **0**，两个精度开关在 argparse 中默认均**关闭**。AR diffusion、CCD 和 DMD 的 launcher 均启用 `--fsdp_reduce_fp32`；GAN/R1/R2 和 `--score_timestep_clamp` 仍关闭。AR diffusion 和 CCD 不添加 GAN 损失。
+这些选项复用 `scripts/wan2.1_self_forcing/train_distill.py` 中的 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 机制。GAN/R1/R2 权重默认均为 **0**，两个精度开关在 argparse 中默认均**关闭**。AR diffusion、CCD 和 DMD 的 launcher 均启用 `--fsdp_reduce_fp32`；GAN/R1/R2 仍关闭，DMD 的 launcher 额外启用 `--score_timestep_clamp`。AR diffusion 和 CCD 不添加 GAN 损失。
 
 判别器共享 **fake-score backbone**，每个选定特征层对应一个可学习 register token。判别器头由 critic optimizer 更新，不加入 generator optimizer。生成器损失为 `DMD + 加权 GAN`，critic 损失为 `denoising + 加权 GAN + R1 + R2`。
 
@@ -289,7 +292,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--r2_weight` | 噪声生成 latent 上的归一化有限差分正则权重 | `0.0` |
 | `--r2_sigma` | R2 扰动标准差 | `0.01` |
 | `--fsdp_reduce_fp32` | FSDP 使用 FP32 梯度归约，参数仍以低精度聚合 | 关闭（launcher 已启用） |
-| `--score_timestep_clamp` | 将 DMD、生成器 GAN 和 critic 时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关闭 |
+| `--score_timestep_clamp` | 将 DMD、生成器 GAN 和 critic 时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关闭（DMD launcher 已启用） |
 
 设 logits 为 `D_fake` 和 `D_real`，普通 GAN 的生成器损失为 `softplus(-D_fake)`，判别器损失为 `softplus(-D_real) + softplus(D_fake)`。相对论模式分别使用 `softplus(-(D_fake - D_real))` 和 `softplus(-(D_real - D_fake))`。
 

scripts/wan2.1_causal_forcing/train_ar_diffusion.sh

@@ -37,6 +37,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \

scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh

@@ -41,6 +41,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --output_dir="output_dir_wan2.1_causal_forcing_ccd" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=0.0 \
   --adam_beta1=0.0 \
   --adam_beta2=0.999 \

scripts/wan2.1_causal_forcing/train_causal_dmd.sh

@@ -38,13 +38,14 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --num_train_epochs=100 \
   --checkpointing_steps=200 \
   --learning_rate=2.0e-06 \
-  --learning_rate_critic=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_causal_forcing_dmd" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=0.0 \
   --adam_beta1=0.0 \
   --adam_beta2=0.999 \
@@ -56,8 +57,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --enable_bucket \
   --num_frame_per_block=3 \
   --use_kv_cache_training \
-  --denoising_step_indices_list 1000 667 334 1 \
-  --real_guidance_scale=6.0 \
+  --use_ema \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --randomize_step_indices \
   --fake_guidance_scale=0.0 \
   --gen_update_interval=5 \

scripts/wan2.1_flex_forcing/README_TRAIN_FLEX_DMD_DISTILL.md

@@ -34,6 +34,7 @@ This document provides the complete workflow for **Flex-Forcing Stage 2 — Dist
   - [3.5 Training with FSDP](#35-training-with-fsdp)
   - [3.6 Other Backends](#36-other-backends)
   - [3.7 Multi-Node Distributed Training](#37-multi-node-distributed-training)
+  - [3.8 Optional One-Forcing GAN, EMA, and Precision Controls](#38-optional-one-forcing-gan-ema-and-precision-controls)
 - [4. Inference Testing](#4-inference-testing)
   - [4.1 Inference Parameters](#41-inference-parameters)
   - [4.2 Text-to-Video (T2V) Inference](#42-text-to-video-t2v-inference)
@@ -107,6 +108,8 @@ DMD uses a **TextDataset** (`--train_mode="normal"`) — the generator creates i
 
 `train_distill.sh` defaults to `datasets/internal_datasets/metadata.json`. For a real run, swap in the **VidProM extended** prompts the paper uses.
 
+Enabling any GAN/R1/R2 weight (see **3.8**) switches the loader to **ImageVideoDataset**: captioned real videos are then required, cropped to the critic's fake input frame count.
+
 ---
 
 ## 3. Distillation Training
@@ -174,15 +177,16 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -204,6 +208,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -220,6 +226,7 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
 |-----------|-------------|---------------|
 | `--config_path` | Component subpaths / `transformer_additional_kwargs` | `config/wan2.1/wan_civitai.yaml` |
 | `--pretrained_model_name_or_path` | Base model directory (1.3B) | `$MODEL_NAME` |
+| `--real_score_pretrained_model_name_or_path` | 14B real-score teacher directory; defaults to `--pretrained_model_name_or_path` | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
 | `--train_data_meta` | Prompt JSON (only the `"text"` field is read) | `$DATASET_META_NAME` |
 | `--ode_transformer_path` | Stage 1 ODE checkpoint for generator/critic init | `$STAGE1_CKPT` |
 | `--fix_sample_size` | Fixed `H W` of the training clip | `480 832` |
@@ -230,9 +237,9 @@ Output: `output_dir_wan2.1_flex_forcing_distill/`.
 | `--dataloader_num_workers` | DataLoader workers | 8 |
 | `--num_train_epochs` | Number of training epochs (overridden by `--max_train_steps`) | 100 |
 | `--max_train_steps` | Hard step cap; the paper uses 600 | 600 |
-| `--checkpointing_steps` | Save checkpoint every N steps | 50 |
-| `--learning_rate` | Generator learning rate | `2e-06` |
-| `--learning_rate_critic` | Critic learning rate | `2e-06` |
+| `--checkpointing_steps` | Save checkpoint every N steps | 200 |
+| `--learning_rate` | Generator learning rate | `2.0e-06` |
+| `--learning_rate_critic` | Critic learning rate | `4.0e-07` |
 | `--lr_scheduler` | LR scheduler type | `constant_with_warmup` |
 | `--lr_warmup_steps` | LR warmup steps | 100 |
 | `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 4.5 |
@@ -322,15 +329,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -352,6 +360,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -400,15 +410,16 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -430,10 +441,13 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
 
+`--use_ema` is dropped here: `train_distill.py` raises `NotImplementedError` for EMA under DeepSpeed Zero-3.
+
 #### 3.6.2 Training without DeepSpeed and FSDP
 
 **This approach is not recommended because there is no memory-saving backend, which can easily cause out-of-memory errors**. We only provide the training shell for reference.
@@ -464,15 +478,16 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -494,6 +509,8 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -538,15 +555,16 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -568,6 +586,8 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -602,6 +622,44 @@ NCCL_DEBUG=INFO
 
 - **Data Synchronization**: All machines must be able to access the same data paths (NFS/shared storage)
 
+### 3.8 Optional One-Forcing GAN, EMA, and Precision Controls
+
+These options port the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms already used by the causal-forcing DMD trainer. All GAN/R1/R2 weights default to **0** and stay off in the launcher. `--use_ema`, `--fsdp_reduce_fp32` and `--score_timestep_clamp` default to **off** in argparse, but `train_distill.sh` enables all three; the DeepSpeed Zero-3 command in **3.6.1** drops `--use_ema`, which Zero-3 rejects.
+
+The discriminator shares the **fake-score backbone**, with one learnable register token per selected feature layer. Its head is optimized by the critic optimizer, not the generator optimizer. Generator loss is `DMD + weighted GAN`; critic loss is `denoising + weighted GAN + R1 + R2`.
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--use_ema` | EMA over the generator; checkpoints save a `transformer_ema` subfolder and validation samples from the EMA weights | off (enabled in `train_distill.sh`) |
+| `--gan_g_weight` | Non-saturating generator GAN loss weight | `0.0` |
+| `--gan_d_weight` | Discriminator GAN loss weight | `0.0` |
+| `--gan_feature_layers` | Fake-score block indices used by the head | `21 29` |
+| `--gan_block_ffn_dim` | Head cross-attention FFN dimension | `2048` |
+| `--gan_block_num_heads` | Head count; must divide the fake-score hidden dimension | `12` |
+| `--relativistic_discriminator` | Use paired real/fake logit differences | off |
+| `--r1_weight` | APT-style finite-difference penalty on noisy real latents | `0.0` |
+| `--r1_sigma` | R1 perturbation standard deviation | `0.01` |
+| `--r2_weight` | Normalized finite-difference penalty on noisy generated latents | `0.0` |
+| `--r2_sigma` | R2 perturbation standard deviation | `0.01` |
+| `--fsdp_reduce_fp32` | FSDP FP32 gradient reduction; parameter gathering stays low precision | off (enabled in `train_distill.sh`) |
+| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off (enabled in `train_distill.sh`) |
+
+For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` for the generator and `softplus(-D_real) + softplus(D_fake)` for the discriminator. Relativistic losses instead use `softplus(-(D_fake - D_real))` and `softplus(-(D_real - D_fake))`.
+
+R1 follows the source's APT form: `r1_weight * mean((D_real_perturbed - D_real)^2)`, **without dividing by sigma**. R2 follows `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`. These are finite-difference penalties, not input-gradient double-backward regularizers. Real/fake samples share timestep and noise; perturbations use independent noise. The non-relativistic critic runs fake and real as two sequential 1x-batch forwards with a backward after each (the memory-optimized split), the relativistic critic runs one packed 2x-batch forward, and R1/R2 each re-forward their own 2x batch after the main backward(s).
+
+To enable GAN, pass positive generator/discriminator weights and, optionally, R1/R2 weights to `train_distill.py`. Requirements:
+- Captioned real videos, not a prompt-only JSON. The real latents are cropped to the critic's fake input frame count, so the clip must cover at least that many latent frames.
+- `--gradient_accumulation_steps=1`; DeepSpeed is not supported by the GAN branch.
+- Valid feature-layer indices for the selected fake-score model. The default `21 29` layers and `12` heads fit Wan2.1-T2V-1.3B; other backbones may need different settings.
+- Enabling only `gan_g_weight` does not supply a discriminator classification-learning objective. Choose the discriminator weight deliberately; nonzero values are not a guaranteed quality improvement.
+
+Trainable fake-score and head master weights stay **FP32**, while forward uses the configured autocast/FSDP precision. Logits, softplus losses, and finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient communication precision, not the forward dtype or model architecture. The clamp changes score-training timesteps, not the generator rollout or validation schedule.
+
+Resume GAN checkpoints with the same GAN/head configuration. The head is saved with fake-score, including FP32 head tensors, and is recreated before loading. It is not part of the exported generator, so inference needs no discriminator or new GAN flags.
+
+Under `--use_ema`, the shadow stays in FP32 and mirrors the live model's mixed-precision policy; FSDP EMA requires FSDP1 and DeepSpeed Zero-3 is rejected.
+
 ---
 
 ## 4. Inference Testing
@@ -752,6 +810,7 @@ Any-order editing needs the full KV cache: both entries set `local_attn_size = -
 ## 5. Additional Resources
 
 - **Flex-Forcing paper**: https://arxiv.org/abs/2607.03509
+- **One-Forcing (optional GAN mechanism)**: https://github.com/Aurora-edu/One-Forcing
 - **Stage 1**: [README_TRAIN_ODE.md](./README_TRAIN_ODE.md)
 - **Shared chunking helpers**: `videox_fun/utils/flex_chunking.py` — `normalize_chunk_spec`, `sample_flexible_chunks`, `sample_flex_level0` (the level-0 four-arm mixture, drawn by **both** trainers), `refine_partition`, `build_pyramid_partitions`, `validate_nested_partitions`, `broadcast_chunk_sizes`
 - **Model / pipeline**: `videox_fun/models/wan_transformer3d_flex_forcing.py`, `videox_fun/pipeline/pipeline_wan_flex_forcing.py`

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

@@ -34,6 +34,7 @@
   - [3.5 使用 FSDP 训练](#35-使用-fsdp-训练)
   - [3.6 其他后端](#36-其他后端)
   - [3.7 多机分布式训练](#37-多机分布式训练)
+  - [3.8 可选的 One-Forcing GAN、EMA 与精度开关](#38-可选的-one-forcing-ganema-与精度开关)
 - [四、推理测试](#四推理测试)
   - [4.1 推理参数解析](#41-推理参数解析)
   - [4.2 文生视频（T2V）推理](#42-文生视频t2v推理)
@@ -107,6 +108,8 @@ DMD 使用 **TextDataset**（`--train_mode="normal"`）——生成器通过 rol
 
 `train_distill.sh` 默认指向 `datasets/internal_datasets/metadata.json`。真实训练请换成论文使用的 **VidProM extended** prompt 集。
 
+开启任何 GAN/R1/R2 权重（见 **3.8**）会把数据加载切到 **ImageVideoDataset**：此时需要带标注的真实视频，并会裁剪到 critic fake 输入的帧数。
+
 ---
 
 ## 三、蒸馏训练
@@ -174,15 +177,16 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -204,6 +208,8 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -220,6 +226,7 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
 |-----------|-------------|---------------|
 | `--config_path` | 组件子路径 / `transformer_additional_kwargs` | `config/wan2.1/wan_civitai.yaml` |
 | `--pretrained_model_name_or_path` | 底座模型目录（1.3B） | `$MODEL_NAME` |
+| `--real_score_pretrained_model_name_or_path` | 14B 真实分数教师目录；默认回落 `--pretrained_model_name_or_path` | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
 | `--train_data_meta` | prompt JSON（只读 `"text"` 字段） | `$DATASET_META_NAME` |
 | `--ode_transformer_path` | 用于生成器 / critic 初始化的阶段 1 ODE checkpoint | `$STAGE1_CKPT` |
 | `--fix_sample_size` | 固定训练片段的 `H W` | `480 832` |
@@ -230,9 +237,9 @@ launcher 钉死的是：81 个像素帧 = **21 个 latent 帧**（一段 5 秒
 | `--dataloader_num_workers` | DataLoader 工作进程数 | 8 |
 | `--num_train_epochs` | 训练 epoch 数（被 `--max_train_steps` 覆盖） | 100 |
 | `--max_train_steps` | 硬性步数上限；论文用 600 | 600 |
-| `--checkpointing_steps` | 每 N 步保存一次 checkpoint | 50 |
-| `--learning_rate` | 生成器学习率 | `2e-06` |
-| `--learning_rate_critic` | critic 学习率 | `2e-06` |
+| `--checkpointing_steps` | 每 N 步保存一次 checkpoint | 200 |
+| `--learning_rate` | 生成器学习率 | `2.0e-06` |
+| `--learning_rate_critic` | critic 学习率 | `4.0e-07` |
 | `--lr_scheduler` | 学习率调度器类型 | `constant_with_warmup` |
 | `--lr_warmup_steps` | 学习率预热步数 | 100 |
 | `--real_guidance_scale` | 真实分数（14B 教师）的 CFG scale | 4.5 |
@@ -333,15 +340,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -363,6 +371,8 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -405,15 +415,16 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -435,10 +446,13 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
 
+此处去掉 `--use_ema`：DeepSpeed Zero-3 下 `train_distill.py` 会对 EMA 抛出 `NotImplementedError`。
+
 DeepSpeed-Zero-3 训练完成之后，得到 `checkpoint-xxx` 状态的保存权重，我们需要将权重文件合并为单个，可以使用提供的`scripts/zero_to_bf16.py`脚本。
 
 ```shell
@@ -473,15 +487,16 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -503,6 +518,8 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_flex_forcing/train_dis
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -548,15 +565,16 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -578,6 +596,8 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"
 ```
@@ -604,6 +624,44 @@ export RANK=1
 > - Use `NCCL_DEBUG=INFO` to troubleshoot issues in multi-machine training.
 > - If there is no RDMA environment, set `export NCCL_IB_DISABLE=1` and `export NCCL_P2P_DISABLE=1` to avoid communication issues.
 
+### 3.8 可选的 One-Forcing GAN、EMA 与精度开关
+
+这些开关把 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 的机制（因果 forcing 的 DMD trainer 已在用）移植过来。所有 GAN/R1/R2 权重默认是 **0**，launcher 中也保持关闭；`--use_ema`、`--fsdp_reduce_fp32` 与 `--score_timestep_clamp` 在 argparse 中默认**关闭**，但 `train_distill.sh` 把三者都打开了；**3.6.1** 的 DeepSpeed Zero-3 命令则去掉了 Zero-3 不支持的 `--use_ema`。
+
+判别器与 **fake-score 骨干共享**，每个被选中的特征层有一个可学习的 register token。判别器头由 critic 优化器优化，而不是生成器优化器。生成器损失为 `DMD + 加权 GAN`；critic 损失为 `去噪 + 加权 GAN + R1 + R2`。
+
+| 参数 | 说明 | 默认 |
+|-----------|-------------|---------|
+| `--use_ema` | 对生成器做 EMA；checkpoint 额外保存 `transformer_ema` 子目录，验证样片采样自 EMA 权重 | 关（launcher 已启用） |
+| `--gan_g_weight` | 非饱和生成器 GAN 损失权重 | `0.0` |
+| `--gan_d_weight` | 判别器 GAN 损失权重 | `0.0` |
+| `--gan_feature_layers` | 判别器头使用的 fake-score block 索引 | `21 29` |
+| `--gan_block_ffn_dim` | 判别器头 cross-attention FFN 维度 | `2048` |
+| `--gan_block_num_heads` | 判别器头 head 数；必须整除 fake-score 隐维度 | `12` |
+| `--relativistic_discriminator` | 使用成对的 real/fake logit 差值 | 关 |
+| `--r1_weight` | 对含噪真实 latent 的 APT 式有限差分惩罚 | `0.0` |
+| `--r1_sigma` | R1 扰动标准差 | `0.01` |
+| `--r2_weight` | 对含噪生成 latent 的归一化有限差分惩罚 | `0.0` |
+| `--r2_sigma` | R2 扰动标准差 | `0.01` |
+| `--fsdp_reduce_fp32` | FSDP 梯度归约用 FP32；参数收集保持低精度 | 关（launcher 已启用） |
+| `--score_timestep_clamp` | 把 DMD、生成器 GAN 与 critic 的时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关（launcher 已启用） |
+
+对 logit `D_fake` 与 `D_real`，普通 GAN 损失是生成器 `softplus(-D_fake)`、判别器 `softplus(-D_real) + softplus(D_fake)`；相对论版本改用 `softplus(-(D_fake - D_real))` 与 `softplus(-(D_real - D_fake))`。
+
+R1 沿用来源的 APT 形式：`r1_weight * mean((D_real_perturbed - D_real)^2)`，**不除以 sigma**。R2 是 `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`。两者都是有限差分惩罚，不是输入梯度二重反向的正则。real/fake 样本共享时间步与噪声；扰动使用独立噪声。非相对论 critic 把 fake 与 real 拆成两次 1x batch 前向、各自一次 backward（省显存的拆分路径），相对论 critic 跑一次打包的 2x batch 前向，R1/R2 则在主 backward 之后各自再前向自己的 2x batch。
+
+启用 GAN 需要给 `train_distill.py` 传入正的生成器/判别器权重，可选再加 R1/R2 权重。要求：
+- 带标注的真实视频，不能是只有 prompt 的 JSON。真实 latent 会裁剪到 critic fake 输入的帧数，因此片段至少要覆盖那么多 latent 帧。
+- `--gradient_accumulation_steps=1`；GAN 分支不支持 DeepSpeed。
+- 特征层索引必须对所选 fake-score 模型有效。默认的 `21 29` 层与 `12` head 适配 Wan2.1-T2V-1.3B；换骨干可能需要不同配置。
+- 只开 `gan_g_weight` 不会给判别器提供分类学习目标。请有意识地选择判别器权重；非零值不保证质量提升。
+
+可训练的 fake-score 与判别器头 master 权重保持 **FP32**，前向使用配置的 autocast/FSDP 精度。logit、softplus 损失与有限差分归约都在 FP32 下计算。`--fsdp_reduce_fp32` 只改梯度通信精度，不改前向 dtype 或模型结构。时间步 clamp 只改分数训练的时间步，不改生成器 rollout 或验证调度。
+
+恢复 GAN checkpoint 时请使用相同的 GAN/判别器头配置。判别器头随 fake-score 一起保存（含 FP32 的 head 张量），加载前重建。它不属于导出的生成器，因此推理不需要判别器或新的 GAN 旗标。
+
+开启 `--use_ema` 时，shadow 保持 FP32 并与 live 模型的混合精度策略对齐；FSDP EMA 仅支持 FSDP1，DeepSpeed Zero-3 会被拒绝。
+
 ---
 
 ## 四、推理测试
@@ -741,4 +799,5 @@ num_frame_per_block = 7          # 干净上下文的提交粒度；None = 整
 ## 五、更多资源
 
 - **Flex-Forcing 论文**：https://arxiv.org/abs/2607.03509
+- **One-Forcing（可选 GAN 机制）**：https://github.com/Aurora-edu/One-Forcing
 - **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun

scripts/wan2.1_flex_forcing/train_distill.py

@@ -17,6 +17,7 @@
 
 import argparse
 import contextlib
+import dataclasses
 import gc
 import json
 import logging
@@ -76,7 +77,8 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512,
 from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
                                WanTransformer3DModel,
                                WanTransformer3DModel_FlexForcing,
-                               WanTransformer3DModel_SelfForcing)
+                               WanTransformer3DModel_SelfForcing,
+                               add_wan_gan_discriminator)
 from videox_fun.pipeline import (WanFlexForcingPipeline, WanI2VPipeline,
                                  WanPipeline, WanSelfForcingPipeline)
 from videox_fun.utils.discrete_sampler import DiscreteSampling
@@ -84,6 +86,7 @@ from videox_fun.utils.flex_chunking import (build_pyramid_partitions,
                                             chunk_boundaries,
                                             sample_flex_level0,
                                             validate_nested_partitions)
+from videox_fun.utils.fsdp_ema import FSDPEMA
 from videox_fun.utils.tqdm_bar import PauseAwareTqdm
 from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
                                     get_image_to_video_latent,
@@ -171,6 +174,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi
     return pred_for_score, score_num_frames, True
 
 
+def repeat_batch(value, repeats):
+    """Repeat a batched tensor or per-sample list along the batch dimension."""
+    if value is None or repeats == 1:
+        return value
+    if torch.is_tensor(value):
+        return torch.cat([value] * repeats)
+    return list(value) * repeats
+
+
 def filter_kwargs(cls, kwargs):
     import inspect
     sig = inspect.signature(cls.__init__)
@@ -743,6 +755,15 @@ def parse_args():
             " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
         ),
     )
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.")
+    parser.add_argument(
+        "--fsdp_reduce_fp32", action="store_true",
+        help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
+    )
+    parser.add_argument(
+        "--score_timestep_clamp", action="store_true",
+        help="Clamp DMD, GAN and critic timesteps to [0.02, 0.98] * num_train_timesteps.",
+    )
     parser.add_argument(
         "--non_ema_revision",
         type=str,
@@ -965,6 +986,16 @@ def parse_args():
         default=None,
         help=("If you want to load the ode-trained weight into generator transformer3d, input its path."),
     )
+    parser.add_argument(
+        "--real_score_pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        help=(
+            "Path to a non-causal pretrained model used as the DMD real_score teacher. "
+            "For CF Stage 3 this should point at Wan2.1-T2V-14B. If unset, falls back to "
+            "`--pretrained_model_name_or_path` (1.3B) — useful only for sanity testing."
+        ),
+    )
     parser.add_argument(
         "--vae_path",
         type=str,
@@ -1170,6 +1201,47 @@ def parse_args():
         help="Probability of applying teacher forcing per step (1.0 = always)"
     )
 
+    parser.add_argument(
+        "--gan_g_weight", type=float, default=0.0,
+        help="One-Forcing: non-saturating generator GAN loss weight (0 disables). Positive GAN/R1/R2 weights require real videos.",
+    )
+    parser.add_argument(
+        "--gan_d_weight", type=float, default=0.0,
+        help="One-Forcing: discriminator GAN loss weight added to critic denoising loss (0 disables).",
+    )
+    parser.add_argument(
+        "--gan_feature_layers", nargs="+", type=int, default=[21, 29],
+        help="fake_score block indices feeding the discriminator head, one register token per layer.",
+    )
+    parser.add_argument(
+        "--gan_block_ffn_dim", type=int, default=2048,
+        help="FFN dimension of the discriminator cross-attention blocks.",
+    )
+    parser.add_argument(
+        "--gan_block_num_heads", type=int, default=12,
+        help="Number of heads of the discriminator cross-attention blocks.",
+    )
+    parser.add_argument(
+        "--relativistic_discriminator", action="store_true",
+        help="Use relativistic GAN losses based on real/fake logit differences.",
+    )
+    parser.add_argument(
+        "--r1_weight", type=float, default=0.0,
+        help="Finite-difference R1 penalty weight on noisy real latents (APT form, not divided by sigma).",
+    )
+    parser.add_argument(
+        "--r1_sigma", type=float, default=0.01,
+        help="Perturbation standard deviation for the R1 penalty.",
+    )
+    parser.add_argument(
+        "--r2_weight", type=float, default=0.0,
+        help="Finite-difference R2 penalty weight on noisy generated latents (0.5 * squared logit difference / sigma squared)."
+    )
+    parser.add_argument(
+        "--r2_sigma", type=float, default=0.01,
+        help="Perturbation standard deviation for the R2 penalty.",
+    )
+
     args = parser.parse_args()
     env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
     if env_local_rank != -1 and env_local_rank != args.local_rank:
@@ -1184,6 +1256,14 @@ def parse_args():
 
 def main():
     args = parse_args()
+    use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight))
+    if any(not math.isfinite(w) or w < 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight)):
+        raise ValueError("GAN and regularization weights must be finite and non-negative.")
+    if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1):
+        raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.")
+    if use_gan and (args.r1_weight > 0 and (not math.isfinite(args.r1_sigma) or args.r1_sigma <= 0)
+                    or args.r2_weight > 0 and (not math.isfinite(args.r2_sigma) or args.r2_sigma <= 0)):
+        raise ValueError("Enabled R1/R2 penalties require finite, positive perturbation standard deviations.")
 
     if args.report_to == "wandb" and args.hub_token is not None:
         raise ValueError(
@@ -1247,6 +1327,9 @@ def main():
         else:
             fsdp_stage = 0
         print(f"Using FSDP stage: {fsdp_stage}")
+        if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
+            fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
+            print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")
 
         args.use_fsdp = True
         if fsdp_stage == 3:
@@ -1374,8 +1457,11 @@ def main():
         transformer_additional_kwargs=generator_transformer_kwargs,
         low_cpu_mem_usage=True,
     ).to(weight_dtype)
+    # CF Stage 3 uses a 14B non-causal teacher as real_score; pass an explicit
+    # `--real_score_pretrained_model_name_or_path` to point at Wan2.1-T2V-14B.
+    real_score_base = args.real_score_pretrained_model_name_or_path or args.pretrained_model_name_or_path
     real_score_transformer3d = WanTransformer3DModel.from_pretrained(
-        os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+        os.path.join(real_score_base, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
         transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
         low_cpu_mem_usage=True,
     ).to(weight_dtype)
@@ -1476,6 +1562,41 @@ def main():
                 param.requires_grad = True
                 break
 
+    if use_gan:
+        if (not args.gan_feature_layers or any(i < 0 or i >= len(fake_score_transformer3d.blocks) for i in args.gan_feature_layers)
+                or args.gan_block_ffn_dim <= 0 or args.gan_block_num_heads <= 0
+                or fake_score_transformer3d.dim % args.gan_block_num_heads != 0):
+            raise ValueError("GAN feature layers must exist and the positive head count must divide the fake_score hidden dimension.")
+        # Keep trainable discriminator master weights in fp32; forward uses autocast.
+        add_wan_gan_discriminator(
+            fake_score_transformer3d,
+            args.gan_feature_layers,
+            args.gan_block_ffn_dim,
+            args.gan_block_num_heads,
+            torch.float32,
+            seed=args.seed or 0,
+        )
+
+    # Create EMA for the generator_transformer3d.
+    if args.use_ema:
+        if zero_stage == 3:
+            raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.")
+
+        if fsdp_plugin is not None:
+            if getattr(fsdp_plugin, "fsdp_version", 1) != 1:
+                raise NotImplementedError("FSDPEMA 仅支持 FSDP1。")
+            ema_module = generator_transformer_cls.from_config(generator_transformer3d.config)
+            ema_generator_transformer3d = FSDPEMA(
+                ema_module, source=generator_transformer3d,
+                accelerator=accelerator, fsdp_plugin=fsdp_plugin,
+            )
+        else:
+            ema_generator_transformer3d = EMAModel(
+                generator_transformer3d.parameters(),
+                model_cls=generator_transformer_cls,
+                model_config=generator_transformer3d.config,
+            )
+
     # `accelerate` 0.16.0 will have better support for customized saving
     if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
         # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
@@ -1486,7 +1607,10 @@ def main():
                     from safetensors.torch import save_file
 
                     safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors")
-                    accelerate_state_dict = {k: v.to(dtype=weight_dtype) for k, v in accelerate_state_dict.items()}
+                    accelerate_state_dict = {
+                        k: v.to(dtype=torch.float32 if k.startswith("gan_head.") else weight_dtype)
+                        for k, v in accelerate_state_dict.items()
+                    }
                     save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"})
 
                     with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file:
@@ -1503,7 +1627,14 @@ def main():
             # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
             def save_model_hook(models, weights, output_dir):
                 if accelerator.is_main_process:
-                    models[0].save_pretrained(os.path.join(output_dir, "transformer"))
+                    model = accelerator.unwrap_model(models[0])
+                    save_dir = os.path.join(output_dir, "transformer")
+                    if hasattr(model, "gan_head"):
+                        from safetensors.torch import save_file
+                        model.save_config(save_dir)
+                        save_file(model.state_dict(), os.path.join(save_dir, "diffusion_pytorch_model.safetensors"), metadata={"format": "pt"})
+                    else:
+                        model.save_pretrained(save_dir)
                     if not args.use_deepspeed:
                         weights.pop()
 
@@ -1529,14 +1660,19 @@ def main():
                         else WanTransformer3DModel
                     )
 
-                    # load diffusers style into model
-                    load_model = model_cls.from_pretrained(
-                        input_dir, subfolder="transformer"
-                    )
-                    unwrapped.register_to_config(**load_model.config)
+                    # GAN heads are recreated from the same CLI configuration before resume.
+                    if hasattr(unwrapped, "gan_head"):
+                        from safetensors.torch import load_file
+                        unwrapped.load_state_dict(load_file(os.path.join(input_dir, "transformer", "diffusion_pytorch_model.safetensors")), strict=True)
+                    else:
+                        # load diffusers style into model
+                        load_model = model_cls.from_pretrained(
+                            input_dir, subfolder="transformer"
+                        )
+                        unwrapped.register_to_config(**load_model.config)
 
-                    unwrapped.load_state_dict(load_model.state_dict())
-                    del load_model
+                        unwrapped.load_state_dict(load_model.state_dict())
+                        del load_model
 
                 pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
                 if os.path.exists(pkl_path):
@@ -1550,6 +1686,32 @@ def main():
         accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
         accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)
 
+        if args.use_ema:
+            def save_ema_hook(models, weights, output_dir):
+                if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process:
+                    ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
+
+            def load_ema_hook(models, input_dir):
+                ema_path = os.path.join(input_dir, "transformer_ema")
+                if isinstance(ema_generator_transformer3d, FSDPEMA):
+                    ema_generator_transformer3d.load_pretrained(ema_path)
+                else:
+                    _, ema_kwargs = generator_transformer_cls.load_config(ema_path, return_unused_kwargs=True)
+                    load_model = generator_transformer_cls.from_pretrained(
+                        input_dir, subfolder="transformer_ema",
+                        transformer_additional_kwargs=generator_transformer_kwargs,
+                        low_cpu_mem_usage=True,
+                        torch_dtype=torch.float32,
+                    )
+                    load_model = EMAModel(load_model.parameters(), model_cls=generator_transformer_cls, model_config=load_model.config)
+                    load_model.load_state_dict(ema_kwargs)
+                    ema_generator_transformer3d.load_state_dict(load_model.state_dict())
+                    ema_generator_transformer3d.to(accelerator.device)
+                    del load_model
+
+            accelerator.register_save_state_pre_hook(save_ema_hook)
+            accelerator.register_load_state_pre_hook(load_ema_hook)
+
     if args.gradient_checkpointing:
         generator_transformer3d.enable_gradient_checkpointing()
         fake_score_transformer3d.enable_gradient_checkpointing()
@@ -1646,6 +1808,12 @@ def main():
                     print(f"Set {name} to lr : {args.learning_rate_critic / 2}")
                 break
 
+    if use_gan:
+        fake_trainable_params_optim[0]['params'].extend(
+            param for name, param in fake_score_transformer3d.named_parameters()
+            if name.startswith("gan_head.") and name not in in_already
+        )
+
     if args.use_came:
         optimizer = optimizer_cls(
             trainable_params_optim,
@@ -1687,7 +1855,7 @@ def main():
         args.random_hw_adapt = False
 
     # Get the dataset
-    if args.train_mode != "normal" or args.use_teacher_forcing:
+    if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
         train_dataset = ImageVideoDataset(
             args.train_data_meta, args.train_data_dir,
             video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames, 
@@ -1715,7 +1883,7 @@ def main():
 
         return length_to_frame_num
 
-    if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing:
+    if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan:
         aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
         batch_sampler_generator = torch.Generator().manual_seed(args.seed)
         batch_sampler = AspectRatioBatchImageVideoSampler(
@@ -2058,6 +2226,9 @@ def main():
     if args.train_mode != "normal":
         clip_image_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
 
+    if args.use_ema:
+        ema_generator_transformer3d.to(accelerator.device)
+
     # We need to recalculate our total training steps as the size of the training dataloader may have changed.
     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
     if overrode_max_train_steps:
@@ -2148,6 +2319,8 @@ def main():
         vae_stream_2 = None
 
     idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
+    min_score_timestep = int(0.02 * noise_scheduler.config.num_train_timesteps)
+    max_score_timestep = int(0.98 * noise_scheduler.config.num_train_timesteps)
 
     def randomize_denoising_step_indices(
         denoising_step_indices_list,
@@ -2208,6 +2381,8 @@ def main():
         # must be averaged by contribution count, not by gradient_accumulation_steps.
         train_gen_log_count = 0
         train_denoising_loss = 0.0
+        train_gan_g_loss = 0.0
+        train_gan_d_loss = 0.0
         batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
         for step, batch in enumerate(train_dataloader):
             # Data batch sanity check
@@ -2340,7 +2515,7 @@ def main():
                                 pixel_values_bs = pixel_values_bs.sample()
                                 new_pixel_values.append(pixel_values_bs)
                             return torch.cat(new_pixel_values, dim = 0)
-                        if args.use_teacher_forcing:
+                        if args.use_teacher_forcing or use_gan:
                             clean_latents = _batch_encode_vae(pixel_values)
                         else:
                             clean_latents = None
@@ -2372,7 +2547,7 @@ def main():
                                 clip_context.append(_clip_context)
                             clip_context = torch.cat(clip_context)
 
-                    if args.use_teacher_forcing:
+                    if args.use_teacher_forcing or use_gan:
                         target_shape = clean_latents.size()
                     else:
                         target_shape = mask_latents.size()
@@ -3120,6 +3295,8 @@ def main():
                     # Common code for both KV cache and block mask training
                     indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
                     generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
+                    if args.score_timestep_clamp:
+                        generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep)
                     generator_denoised_input = add_noise(
                         generator_pred,
                         torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
@@ -3225,17 +3402,53 @@ def main():
                     train_dmd_loss += avg_dmd_loss.item()
                     train_gen_log_count += 1
 
+                    # --- One-Forcing Generator GAN Loss ---
+                    generator_loss = dmd_loss
+                    if args.gan_g_weight > 0:
+                        indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
+                        gan_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
+                        if args.score_timestep_clamp:
+                            gan_timestep = gan_timestep.clamp(min_score_timestep, max_score_timestep)
+                        gan_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng)
+                        gan_inputs = [add_noise(generator_pred, gan_noise, gan_timestep)]
+                        if args.relativistic_discriminator:
+                            real_latents = clean_latents[:, :, :generator_pred.size(2)].to(generator_pred.dtype)
+                            gan_inputs.append(add_noise(real_latents, gan_noise, gan_timestep))
+                        num_gan_inputs = len(gan_inputs)
+                        with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                            _, gan_logits = fake_score_transformer3d(
+                                x=torch.cat(gan_inputs).to(dtype=weight_dtype),
+                                context=repeat_batch(prompt_embeds, num_gan_inputs),
+                                t=gan_timestep.repeat(num_gan_inputs),
+                                seq_len=seq_len,
+                                y=repeat_batch(inpaint_latents, num_gan_inputs) if args.train_mode != "normal" else None,
+                                clip_fea=repeat_batch(clip_context, num_gan_inputs) if args.train_mode != "normal" else None,
+                                gan_classify=True,
+                            )
+                        gan_logits = gan_logits.float().chunk(num_gan_inputs)
+                        if args.relativistic_discriminator:
+                            gan_g_loss = F.softplus(-(gan_logits[0] - gan_logits[1])).mean()
+                        else:
+                            gan_g_loss = F.softplus(-gan_logits[0]).mean()
+                        gan_g_loss = gan_g_loss * args.gan_g_weight
+                        train_gan_g_loss += accelerator.gather(gan_g_loss.detach().repeat(args.train_batch_size)).mean().item()
+                        generator_loss = generator_loss + gan_g_loss
+
                     if args.low_vram:
                         real_score_transformer3d = real_score_transformer3d.to("cpu")
-                        fake_score_transformer3d = fake_score_transformer3d.to("cpu")
+                        if args.gan_g_weight == 0:
+                            fake_score_transformer3d = fake_score_transformer3d.to("cpu")
                         torch.cuda.empty_cache()
 
-                    accelerator.backward(dmd_loss)
+                    accelerator.backward(generator_loss)
                     if accelerator.sync_gradients:
                         accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
                     optimizer.step()
                     lr_scheduler.step()
                     optimizer.zero_grad()
+                    if args.gan_g_weight > 0:
+                        # Discard discriminator gradients from the generator GAN backward.
+                        critic_optimizer.zero_grad()
                     
                     if args.low_vram:
                         fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device)
@@ -3626,6 +3839,8 @@ def main():
 
                 indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
                 critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
+                if args.score_timestep_clamp:
+                    critic_timestep = critic_timestep.clamp(min_score_timestep, max_score_timestep)
                 critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng)
 
                 fake_score_denoised_input = add_noise(
@@ -3634,15 +3849,25 @@ def main():
                     critic_timestep
                 )
 
-                with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
-                    fake_score_denoised_output = fake_score_transformer3d(
-                        x=fake_score_denoised_input,
-                        context=prompt_embeds,
-                        t=critic_timestep,
-                        seq_len=seq_len,
-                        y=inpaint_latents if args.train_mode != "normal" else None,
-                        clip_fea=clip_context if args.train_mode != "normal" else None,
-                    )
+                # Generated and real latents share critic_noise / critic_timestep. Prepare the
+                # real input and the R1/R2 perturbations once; they are consumed by the critic
+                # forward below and by the R1/R2 penalties after the main backward(s).
+                real_noisy = None
+                r1_perturbation = None
+                r2_perturbation = None
+                if use_gan:
+                    real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype)
+                    real_noisy = add_noise(real_latents, critic_noise, critic_timestep)
+                    del real_latents
+                    if args.r1_weight > 0:
+                        r1_perturbation = args.r1_sigma * torch.randn(
+                            real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng
+                        )
+                    if args.r2_weight > 0:
+                        r2_perturbation = args.r2_sigma * torch.randn(
+                            fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype,
+                            device=fake_score_denoised_input.device, generator=torch_rng
+                        )
 
                 def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
                     noise_pred = noise_pred.float()
@@ -3656,11 +3881,136 @@ def main():
                     final_loss = masked_loss.mean()
                     return final_loss
 
-                denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
-                avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
-                train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
-                
-                accelerator_fake_score_transformer3d.backward(denoising_loss)
+                # --- One-Forcing Discriminator GAN Loss ---
+                if use_gan and not args.relativistic_discriminator:
+                    # Memory-optimized critic: the non-relativistic gan_d_loss =
+                    # softplus(-real) + softplus(fake) is separable, so run fake and real as two
+                    # sequential 1x-batch forwards and backward the fake pass first. Backwarding
+                    # fake frees its gradient-checkpointing block inputs before the real forward,
+                    # dropping the packed fake+real+R1+R2 forward peak (up to 4x batch) down to
+                    # 1x batch here; R1/R2 re-forward their own 2x batch below. Both backward(s)
+                    # accumulate into fake_score.grad, giving the identical update as the packed
+                    # single-forward critic.
+                    # -- Fake pass (1x batch): denoising output + fake discriminator logit --
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        fake_score_denoised_output, fake_gan_logits = fake_score_transformer3d(
+                            x=fake_score_denoised_input,
+                            context=prompt_embeds,
+                            t=critic_timestep,
+                            seq_len=seq_len,
+                            y=inpaint_latents if args.train_mode != "normal" else None,
+                            clip_fea=clip_context if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                    denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
+                    avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
+                    train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
+
+                    fake_logit = fake_gan_logits.float()
+                    fake_d_loss = F.softplus(fake_logit).mean() * args.gan_d_weight
+                    critic_loss = denoising_loss + fake_d_loss
+                    accelerator_fake_score_transformer3d.backward(critic_loss)
+                    # Fake graph freed: the transformer blocks' stored inputs are released here.
+                    fake_d_detached = fake_d_loss.detach()
+                    del critic_loss, fake_gan_logits, fake_logit, fake_d_loss, fake_score_denoised_output
+
+                    # -- Real pass (1x batch): discriminator logit only --
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        _, real_gan_logits = fake_score_transformer3d(
+                            x=real_noisy,
+                            context=prompt_embeds,
+                            t=critic_timestep,
+                            seq_len=seq_len,
+                            y=inpaint_latents if args.train_mode != "normal" else None,
+                            clip_fea=clip_context if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                    real_logit = real_gan_logits.float()
+                    real_d_loss = F.softplus(-real_logit).mean() * args.gan_d_weight
+                    accelerator_fake_score_transformer3d.backward(real_d_loss)
+                    # Real graph freed. Combine the two detached halves for logging.
+                    gan_d_loss = fake_d_detached + real_d_loss.detach()
+                    train_gan_d_loss += accelerator.gather(gan_d_loss.repeat(args.train_batch_size)).mean().item()
+                    del fake_d_detached, real_gan_logits, real_logit, real_d_loss
+                else:
+                    # Combined path: 2x-batch forward for relativistic GAN (both logits coupled by
+                    # softplus(-(real - fake))), or 1x-batch forward when GAN is disabled.
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        if use_gan:
+                            fake_score_denoised_output, gan_logits = fake_score_transformer3d(
+                                x=torch.cat([fake_score_denoised_input, real_noisy]),
+                                context=repeat_batch(prompt_embeds, 2),
+                                t=critic_timestep.repeat(2),
+                                seq_len=seq_len,
+                                y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
+                                clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
+                                gan_classify=True,
+                            )
+                            fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0]
+                            gan_logits = gan_logits.float().chunk(2)
+                        else:
+                            fake_score_denoised_output = fake_score_transformer3d(
+                                x=fake_score_denoised_input,
+                                context=prompt_embeds,
+                                t=critic_timestep,
+                                seq_len=seq_len,
+                                y=inpaint_latents if args.train_mode != "normal" else None,
+                                clip_fea=clip_context if args.train_mode != "normal" else None,
+                            )
+
+                    denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
+                    avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
+                    train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
+
+                    critic_loss = denoising_loss
+                    if use_gan:
+                        fake_logit, real_logit = gan_logits[0], gan_logits[1]
+                        if args.relativistic_discriminator:
+                            gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
+                        else:
+                            gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
+                        gan_d_loss = gan_d_loss * args.gan_d_weight
+                        train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
+                        critic_loss = critic_loss + gan_d_loss
+
+                    accelerator_fake_score_transformer3d.backward(critic_loss)
+                    del critic_loss
+
+                # R1/R2 gradient penalties: each re-forwards its own 2x batch after the main
+                # backward(s) and finite-differences the perturbed logit against the clean one,
+                # matching the packed-formula value while keeping the peak at 2x batch.
+                if use_gan and args.r1_weight > 0:
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        _, r1_logits = fake_score_transformer3d(
+                            x=torch.cat([real_noisy + r1_perturbation, real_noisy]),
+                            context=repeat_batch(prompt_embeds, 2),
+                            t=critic_timestep.repeat(2),
+                            seq_len=seq_len,
+                            y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
+                            clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                        r1_logits = r1_logits.float().chunk(2)
+                        r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean()
+                    accelerator_fake_score_transformer3d.backward(r1_loss)
+                    del r1_loss
+
+                if use_gan and args.r2_weight > 0:
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        _, r2_logits = fake_score_transformer3d(
+                            x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]),
+                            context=repeat_batch(prompt_embeds, 2),
+                            t=critic_timestep.repeat(2),
+                            seq_len=seq_len,
+                            y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
+                            clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                        r2_logits = r2_logits.float().chunk(2)
+                        r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean()
+                    accelerator_fake_score_transformer3d.backward(r2_loss)
+                    del r2_loss
+
                 if accelerator_fake_score_transformer3d.sync_gradients:
                     accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm)
                 critic_optimizer.step()
@@ -3674,12 +4024,19 @@ def main():
             # Checks if the accelerator has performed an optimization step behind the scenes
             if accelerator.sync_gradients:
 
+                if args.use_ema:
+                    ema_generator_transformer3d.step(generator_transformer3d.parameters())
                 progress_bar.update(1)
                 global_step += 1
-                accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step)
+                logs = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}
+                if use_gan:
+                    logs.update({"train_gan_g_loss": train_gan_g_loss / max(train_gen_log_count, 1), "train_gan_d_loss": train_gan_d_loss})
+                accelerator.log(logs, step=global_step)
                 train_dmd_loss = 0.0
                 train_gen_log_count = 0
                 train_denoising_loss = 0.0
+                train_gan_g_loss = 0.0
+                train_gan_d_loss = 0.0
 
                 if global_step % args.checkpointing_steps == 0:
                     if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
@@ -3722,6 +4079,10 @@ def main():
 
                 if args.validation_prompts is not None and global_step % args.validation_steps == 0:
                     with progress_bar.paused():
+                        if args.use_ema:
+                            # Store the generator parameters temporarily and load the EMA parameters to perform inference.
+                            ema_generator_transformer3d.store(generator_transformer3d.parameters())
+                            ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
                         log_validation(
                             vae,
                             text_encoder,
@@ -3734,6 +4095,9 @@ def main():
                             weight_dtype,
                             global_step,
                         )
+                        if args.use_ema:
+                            # Switch back to the original generator parameters.
+                            ema_generator_transformer3d.restore(generator_transformer3d.parameters())
 
             logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
             progress_bar.set_postfix(**logs)
@@ -3743,6 +4107,10 @@ def main():
 
         if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
             with progress_bar.paused():
+                if args.use_ema:
+                    # Store the generator parameters temporarily and load the EMA parameters to perform inference.
+                    ema_generator_transformer3d.store(generator_transformer3d.parameters())
+                    ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
                 log_validation(
                     vae,
                     text_encoder,
@@ -3755,6 +4123,9 @@ def main():
                     weight_dtype,
                     global_step,
                 )
+                if args.use_ema:
+                    # Switch back to the original generator parameters.
+                    ema_generator_transformer3d.restore(generator_transformer3d.parameters())
 
     # Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever
     # something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto

scripts/wan2.1_flex_forcing/train_distill.sh

@@ -28,15 +28,16 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --gradient_accumulation_steps=1 \
   --dataloader_num_workers=8 \
   --num_train_epochs=100 \
-  --checkpointing_steps=50 \
-  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --checkpointing_steps=200 \
+  --learning_rate=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_flex_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
   --adam_weight_decay=3e-2 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
@@ -58,5 +59,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
   --randomize_step_indices \
   --flow_euler_rollout \
   --use_kv_cache_training \
+  --use_ema \
+  --score_timestep_clamp \
   --num_frame_per_block=3 \
   --ode_transformer_path="output_dir_wan2.1_flex_forcing_ode_regression/checkpoint-1000/diffusion_pytorch_model.safetensors"

scripts/wan2.1_one_forcing/README_ONE_FORCING.md

@@ -0,0 +1,356 @@
+# Wan2.1 One-Forcing: Stable One-Step Autoregressive Video Generation
+
+This document provides the complete workflow for **One-Forcing** on Wan2.1-T2V-1.3B, a recipe ported from [Aurora-edu/One-Forcing](https://github.com/Aurora-edu/One-Forcing).
+
+> **What is One-Forcing?**
+>
+> One-Forcing (*Towards Stable One-Step Autoregressive Video Generation*) is **Causal-Forcing Stage 3 — Distribution Matching Distillation (DMD)** *plus* a **shared noised-latent adversarial critic**. The warm-start generator is produced by the first two Causal-Forcing stages, and One-Forcing then polishes it into a **stable one-step, frame-wise autoregressive** video model:
+>
+> 1. **Stage 1 — AR Diffusion** (`scripts/wan2.1_causal_forcing/train_ar_diffusion.py`): Train the model with **teacher forcing** on clean video latents, learning to denoise chunk-by-chunk (or frame-by-frame) in a causal autoregressive manner. This produces a strong AR backbone. See [README_TRAIN_AR_DIFFUSION.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md).
+> 2. **Stage 2 — Causal Consistency Distillation (CCD)** (`scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py`): Distill the multi-step AR model into a **one-step-per-block** consistency model using an EMA teacher with CFG guidance. Its output checkpoint is the **warm-start generator** for One-Forcing. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md).
+> 3. **Stage 3 — One-Forcing** (`train_one_forcing.py`, this document), which adds two mechanisms on top of DMD:
+>    - **DMD core**: the generator matches the real distribution through a 14B real-score teacher and a 1.3B fake-score critic (`train_causal_dmd.py`).
+>    - **Adversarial critic**: a discriminator head rides on the **fake-score backbone** (one register token per selected feature layer). Both the DMD denoising loss and the adversarial loss operate on the **same noised generated latent**, and the critic uses **one shared noise/timestep** for fake and real samples.
+>    - **One step**: `denoising_step_list = [1000]` — each frame is generated in a single denoising step (the "One" in One-Forcing).
+>
+> Instead of training Stage 1/2 yourself, you can also warm-start from the **official ODE-initialized checkpoint** released by the One-Forcing authors (see Section 3).
+>
+> **The trainer** `scripts/wan2.1_one_forcing/train_one_forcing.py` implements the One-Forcing generator/critic GAN losses, the relativistic option, and finite-difference R1/R2 penalties. For the underlying DMD stage, see [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md).
+
+---
+
+## Table of Contents
+- [1. Prerequisites](#1-prerequisites)
+- [2. Environment Setup](#2-environment-setup)
+- [3. Download Pretrained Models](#3-download-pretrained-models)
+- [4. Prepare Training Data](#4-prepare-training-data)
+  - [4.1 Quick Demo Dataset](#41-quick-demo-dataset)
+  - [4.2 Dataset Structure](#42-dataset-structure)
+  - [4.3 metadata.json Format](#43-metadatajson-format)
+- [5. Training](#5-training)
+  - [5.1 Quick Start](#51-quick-start)
+  - [5.2 Recipe Parameters](#52-recipe-parameters)
+  - [5.3 GAN Prerequisites](#53-gan-prerequisites)
+  - [5.4 Optional Regularizers (Relativistic / R1 / R2)](#54-optional-regularizers-relativistic--r1--r2)
+- [6. Inference](#6-inference)
+  - [6.1 Quick Start](#61-quick-start)
+  - [6.2 1-Step vs 4-Step Selector](#62-1-step-vs-4-step-selector)
+  - [6.3 First-Frame Schedule](#63-first-frame-schedule)
+- [7. Checkpoint Round-Trip](#7-checkpoint-round-trip)
+- [8. Additional Resources](#8-additional-resources)
+
+---
+
+## 1. Prerequisites
+
+One-Forcing is the final polish on top of a distilled generator, so it needs:
+
+1. A **warm-start generator checkpoint** — a distilled/ODE generator that the DMD + adversarial critic then polishes (e.g. a **Causal-Forcing Stage-2 CCD** checkpoint, or the official ODE-initialized checkpoint in Section 3).
+2. A **Wan2.1-T2V-14B** model as the DMD real-score teacher.
+3. **Captioned real videos** — the adversarial critic needs clean real latents, so a prompt-only JSON is **not** sufficient (unlike plain DMD).
+
+---
+
+## 2. Environment Setup
+
+**Method 1: Using requirements.txt**
+
+```bash
+pip install -r requirements.txt
+```
+
+**Method 2: Manual Installation**
+
+```bash
+pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
+pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
+pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
+pip install yunchang xfuser modelscope openpyxl
+pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
+pip install opencv-python-headless
+pip install deepspeed==0.17.0 numpy==1.26.4
+```
+
+**Method 3: Using Docker**
+
+```bash
+# pull image
+docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+
+# enter image
+docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+```
+
+---
+
+## 3. Download Pretrained Models
+
+One-Forcing requires **two** pretrained models:
+
+- **Wan2.1-T2V-1.3B**: base model for the generator/critic.
+- **Wan2.1-T2V-14B**: the non-causal real-score teacher used by DMD to compute the real distribution score.
+
+```bash
+# Create model directory
+mkdir -p models/Diffusion_Transformer
+
+# Download Wan2.1 T2V 1.3B (student base model)
+modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
+
+# Download Wan2.1 T2V 14B (DMD real-score teacher)
+modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B
+```
+
+The official One-Forcing checkpoints are hosted on HuggingFace (`JiaqiFeng/OneForcing`):
+
+```bash
+# Warm-start generator for training (the official ODE-initialized checkpoint),
+# loaded via --ode_transformer_path.
+hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
+
+# Trained One-Forcing generator for inference (see Section 6).
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+If you trained Causal-Forcing Stage 1/2 yourself, you can instead point `--ode_transformer_path` at your own Stage-2 CCD checkpoint.
+
+---
+
+## 4. Prepare Training Data
+
+One-Forcing enables the GAN branch by default (`--gan_g_weight=0.03 --gan_d_weight=0.03`), so the data loader always uses an **ImageVideoDataset**: the adversarial critic needs **captioned real videos** to build its clean latents. A prompt-only JSON is **not** sufficient (unlike plain DMD).
+
+### 4.1 Quick Demo Dataset
+
+```bash
+# Download official demo dataset
+modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
+```
+
+### 4.2 Dataset Structure
+
+Same as Causal-Forcing Stage 1/2. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md#42-dataset-structure) for details.
+
+### 4.3 metadata.json Format
+
+Each entry needs `file_path`, `type="video"`, and `text`; the launcher uses the `metadata_add_width_height.json` variant, which additionally carries `width`/`height` so bucketing does not need to decode the media files:
+
+```json
+[
+  {
+    "file_path": "train/00000000.mp4",
+    "text": "A young woman gently turns her head to the right, tilting it slightly as she offers a soft greeting...",
+    "type": "video",
+    "width": 832,
+    "height": 480
+  }
+]
+```
+
+> **Note**: Keep enough frames per video for `--video_sample_stride=1` and `--video_sample_n_frames=81`: the GAN branch aborts when the source yields fewer latent frames than `score_num_frames=21` (see Section 5.3).
+
+---
+
+## 5. Training
+
+### 5.1 Quick Start
+
+The ready-to-use launcher is [train_one_forcing.sh](./train_one_forcing.sh). It calls the One-Forcing trainer (`train_one_forcing.py`) with the GAN recipe enabled:
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
+
+accelerate launch --mixed_precision="bf16" --use_fsdp \
+    --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+    scripts/wan2.1_one_forcing/train_one_forcing.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --ode_transformer_path=$STAGE2_CKPT \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --fix_sample_size 480 832 \
+  --video_sample_stride=1 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --max_train_steps=200 \
+  --checkpointing_steps=100 \
+  --learning_rate=1.0e-05 \
+  --learning_rate_critic=1.0e-05 \
+  --lr_scheduler="constant" \
+  --seed=42 \
+  --output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.01 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=10.0 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --num_frame_per_block=1 \
+  --use_kv_cache_training \
+  --use_ema \
+  --ema_decay=0.99 \
+  --denoising_step_indices_list 1000 \
+  --real_guidance_scale=6.0 \
+  --score_timestep_clamp \
+  --fake_guidance_scale=0.0 \
+  --gen_update_interval=5 \
+  --gan_g_weight=0.03 \
+  --gan_d_weight=0.03 \
+  --gan_feature_layers 21 29 \
+  --gan_block_ffn_dim=2048 \
+  --gan_block_num_heads=12 \
+  --resume_from_checkpoint="latest" \
+  --trainable_modules "."
+```
+
+Or simply:
+
+```bash
+bash scripts/wan2.1_one_forcing/train_one_forcing.sh
+```
+
+### 5.2 Recipe Parameters
+
+The One-Forcing recipe sets these launcher flags:
+
+| Parameter | Launcher flag | Value |
+|------------------------|---------------|-------|
+| `denoising_step_list: [1000]` | `--denoising_step_indices_list` | `1000` (one step) |
+| `num_frame_per_block: 1` | `--num_frame_per_block` | `1` (frame-wise) |
+| `num_training_frames: 21` / `score_num_frames: 21` | `--video_sample_n_frames=81` @ 480×832 | 21 latent frames |
+| `guidance_scale: 5.0` | `--real_guidance_scale` | `6.0` (see CFG note below) |
+| `lr: 1e-5` / `lr_critic: 1e-5` | `--learning_rate` / `--learning_rate_critic` | `1.0e-05` |
+| `beta1: 0.0`, `beta2: 0.999`, `weight_decay: 0.01` | `--adam_beta1`, `--adam_beta2`, `--adam_weight_decay` | `0.0`, `0.999`, `0.01` |
+| `ema_weight: 0.99` | `--use_ema`, `--ema_decay` | on, `0.99` |
+| `max_steps: 200` | `--max_train_steps` | `200` |
+| `dfake_gen_update_ratio: 5` | `--gen_update_interval` | `5` |
+| `gan_feature_layers: [21, 29]` | `--gan_feature_layers` | `21 29` |
+| `gan_block_ffn_dim: 2048` | `--gan_block_ffn_dim` | `2048` |
+| `gan_block_num_heads: 12` | `--gan_block_num_heads` | `12` |
+| `gan_g_weight: 0.03`, `gan_d_weight: 0.03` | `--gan_g_weight`, `--gan_d_weight` | `0.03`, `0.03` |
+| `relativistic_discriminator: false` | `--relativistic_discriminator` | omitted (off) |
+| `r1_weight: 0.0`, `r2_weight: 0.0` | `--r1_weight`, `--r2_weight` | omitted (off) |
+| `context_noise: 0` | `--context_noise` | `0` (default) |
+
+> **CFG convention**: this repo computes guidance as `cond + scale * (cond - uncond)`, and this recipe uses `--real_guidance_scale=6.0`.
+
+For the general training/video-sampling parameters shared with the DMD stage, see [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md#52-key-parameters) Section 5.2.
+
+### 5.3 GAN Prerequisites
+
+The GAN branch is validated at startup in `train_one_forcing.py`. It requires:
+
+- **Captioned real videos** (`type="video"` + `file_path` + `text`), not a prompt-only JSON. Enabling any positive GAN/R1/R2 weight switches the loader from `TextDataset` to `ImageVideoDataset`.
+- **`--gradient_accumulation_steps=1`** and **no DeepSpeed** — the trainer raises `ValueError` otherwise.
+- **`--video_sample_stride=1`**: with `stride=2`, an 81-frame source yields ~10 latent frames, which is below `score_num_frames=21` and aborts.
+- **Valid feature-layer indices** for the fake-score backbone. The defaults `21 29` (layers) and `12` (heads) fit Wan2.1-T2V-1.3B; the head count must divide the fake-score hidden dimension.
+
+Trainable fake-score and head master weights stay **FP32**; the forward pass uses the configured autocast/FSDP precision, and logits/softplus/finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient-communication precision only.
+
+### 5.4 Optional Regularizers (Relativistic / R1 / R2)
+
+These regularizers are **off** by default. To experiment, append to the launcher:
+
+```bash
+# Relativistic discriminator: use paired real/fake logit differences
+--relativistic_discriminator
+
+# Finite-difference gradient penalties (APT form, not double-backward)
+--r1_weight=1.0 --r1_sigma=0.1     # on noisy real latents:   r1_weight * mean((D_real' - D_real)^2)
+--r2_weight=0.001 --r2_sigma=0.01  # on noisy gen latents: 0.5 * r2_weight * mean(((D_fake' - D_fake)/sigma)^2)
+```
+
+For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` (generator) and `softplus(-D_real) + softplus(D_fake)` (discriminator); relativistic losses use the logit differences instead. Real/fake samples share timestep and noise; perturbations use independent noise.
+
+---
+
+## 6. Inference
+
+### 6.1 Quick Start
+
+The inference entry is [examples/wan2.1_one_forcing/predict_t2v.py](../../examples/wan2.1_one_forcing/predict_t2v.py). It loads the **1.3B student** (the 14B model is only the DMD teacher at training time) and runs frame-wise one-step generation. By default it uses the **official One-Forcing checkpoint**:
+
+```bash
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+```python
+# examples/wan2.1_one_forcing/predict_t2v.py
+model_name          = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+transformer_path    = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
+use_ema             = False   # only applies to trainer-output checkpoint dirs
+stochastic_sampling = True    # DMD-distilled student
+guidance_scale      = 1.0     # CFG is baked into the distilled weights
+num_inference_steps = 1       # One-Forcing one-step generation
+num_frame_per_block = 1       # frame-wise
+sample_size         = [480, 832]
+video_length        = 81
+```
+
+Run it:
+
+```bash
+python examples/wan2.1_one_forcing/predict_t2v.py
+```
+
+`transformer_path` may be the official checkpoint file, a plain `.pt`/`.safetensors` file, or a trainer-output **checkpoint dir** (its live/EMA weights are auto-resolved). The GAN discriminator head is a training-only submodule of fake_score, so inference needs **no discriminator wiring**.
+
+### 6.2 1-Step vs 4-Step Selector
+
+With `stochastic_sampling=True`, `num_inference_steps` selects the warped timestep schedule:
+
+| `num_inference_steps` | Denoise timesteps | Use case |
+|-----------------------|-------------------|----------|
+| `1` (One-Forcing default) | `[1000]` | Fastest; frame-wise one-step generation |
+| `4` | `[1000, 750, 500, 250]` | More stable, slower |
+
+### 6.3 First-Frame Schedule
+
+The shared `WanSelfForcingPipeline` applies **one denoising schedule to all blocks**. A separate first-frame schedule (e.g. `[1000, 750, 500, 250]` for the first frame and `[1000]` for later frames) is therefore not applied; the pipeline is shared by multiple examples and is intentionally left unmodified.
+
+---
+
+## 7. Checkpoint Round-Trip
+
+The FSDP trainer writes each `checkpoint-{step}/` as follows, so training output feeds inference with no key surgery:
+
+| Component | Saved to | Contains `gan_head`? |
+|-----------|----------|----------------------|
+| **Generator** (live) | `checkpoint-{step}/diffusion_pytorch_model.safetensors` (root, via FSDP hook) | No |
+| **Generator** (EMA) | `checkpoint-{step}/transformer_ema/` | No |
+| **Fake-score + discriminator head** | `checkpoint-{step}/fake_score/` | Yes (FP32 head tensors) |
+
+- The **generator never contains `gan_head`** keys (the head lives on fake-score), so `predict_t2v.py` loads it cleanly with `strict=False`.
+- `_resolve_transformer_path` in the inference script resolves, in order: `transformer_ema/`, `ema_transformer/` (compat alias), `transformer/`, then the root safetensors.
+- **Resume** GAN training with the same GAN/head configuration: the head is recreated from CLI flags before loading and restored from `fake_score/`.
+
+---
+
+## 8. Additional Resources
+
+- **One-Forcing (source recipe)**: https://github.com/Aurora-edu/One-Forcing
+- **Causal-Forcing (DMD stage)**: https://github.com/thu-ml/Causal-Forcing
+- **Stage 1 (AR Diffusion) training guide**: [README_TRAIN_AR_DIFFUSION.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md)
+- **Stage 2 (CCD) training guide**: [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md)
+- **Underlying DMD trainer guide**: [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md)
+- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun

scripts/wan2.1_one_forcing/README_ONE_FORCING_zh-CN.md

@@ -0,0 +1,355 @@
+# Wan2.1 One-Forcing：稳定的单步自回归视频生成
+
+本文档提供在 Wan2.1-T2V-1.3B 上进行 **One-Forcing** 的完整工作流，该配方移植自 [Aurora-edu/One-Forcing](https://github.com/Aurora-edu/One-Forcing)。
+
+> **什么是 One-Forcing？**
+>
+> One-Forcing（*Towards Stable One-Step Autoregressive Video Generation*）= **Causal-Forcing Stage 3 — 分布匹配蒸馏 (DMD)** *加上*一个**共享噪声 latent 的对抗判别器（critic）**。热启动生成器由 Causal-Forcing 的前两个阶段产出，One-Forcing 再将其打磨为**稳定的单步、逐帧自回归**视频模型：
+>
+> 1. **Stage 1 — AR Diffusion**（`scripts/wan2.1_causal_forcing/train_ar_diffusion.py`）：在干净视频 latent 上以 **teacher forcing** 训练模型，学习按块（或按帧）因果自回归地去噪，产出强 AR 基座。参见 [README_TRAIN_AR_DIFFUSION_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md)。
+> 2. **Stage 2 — 因果一致性蒸馏 (CCD)**（`scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py`）：使用 EMA teacher + CFG 将多步 AR 模型蒸馏为**每块一步**的一致性模型，其产出 checkpoint 即 One-Forcing 的**热启动生成器**。参见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md)。
+> 3. **Stage 3 — One-Forcing**（`train_one_forcing.py`，即本文档），在 DMD 之上增加两个机制：
+>    - **DMD 核心**：生成器通过 14B real-score teacher 与 1.3B fake-score critic 匹配真实分布（`train_causal_dmd.py`）。
+>    - **对抗 critic**：判别器头挂在 **fake-score backbone** 上（每个选定特征层对应一个 register token）。DMD 去噪损失与对抗损失作用在**同一个噪声生成 latent** 上，critic 对真实/生成样本使用**同一份共享噪声与时间步**。
+>    - **单步**：`denoising_step_list = [1000]` —— 每帧在单个去噪步内生成（即 One-Forcing 中的 “One”）。
+>
+> 你也可以不自己训练 Stage 1/2，而是直接使用 One-Forcing 作者发布的**官方 ODE 初始化 checkpoint** 作为热启动权重（见第三节）。
+>
+> **训练器** `scripts/wan2.1_one_forcing/train_one_forcing.py` 实现了 One-Forcing 的生成器/critic GAN 损失、相对论（relativistic）选项以及有限差分 R1/R2 正则。底层 DMD 阶段请参见 [README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md)。
+
+---
+
+## 目录
+- [一、前置条件](#一前置条件)
+- [二、环境配置](#二环境配置)
+- [三、下载预训练模型](#三下载预训练模型)
+- [四、数据准备](#四数据准备)
+  - [4.1 快速测试数据集](#41-快速测试数据集)
+  - [4.2 数据集结构](#42-数据集结构)
+  - [4.3 metadata.json 格式](#43-metadatajson-格式)
+- [五、训练](#五训练)
+  - [5.1 快速开始](#51-快速开始)
+  - [5.2 配方参数](#52-配方参数)
+  - [5.3 GAN 前置要求](#53-gan-前置要求)
+  - [5.4 可选正则（相对论 / R1 / R2）](#54-可选正则相对论--r1--r2)
+- [六、推理](#六推理)
+  - [6.1 快速开始](#61-快速开始)
+  - [6.2 单步与 4 步选择器](#62-单步与-4-步选择器)
+  - [6.3 首帧调度](#63-首帧调度)
+- [七、Checkpoint 往返（训练→推理）](#七checkpoint-往返训练推理)
+- [八、更多资源](#八更多资源)
+
+---
+
+## 一、前置条件
+
+One-Forcing 是在已蒸馏生成器之上的最后打磨，因此需要：
+
+1. **热启动生成器 checkpoint** —— 一个已蒸馏/ODE 生成器，由 DMD + 对抗 critic 进一步打磨（例如 **Causal-Forcing Stage-2 CCD** checkpoint，或第三节的官方 ODE 初始化 checkpoint）。
+2. **Wan2.1-T2V-14B** 模型，作为 DMD real-score teacher。
+3. **带描述的真实视频** —— 对抗 critic 需要干净的真实 latent，因此仅有提示词的 JSON **不够用**（这一点与普通 DMD 不同）。
+
+---
+
+## 二、环境配置
+
+**方式 1：使用 requirements.txt**
+
+```bash
+pip install -r requirements.txt
+```
+
+**方式 2：手动安装依赖**
+
+```bash
+pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
+pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
+pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
+pip install yunchang xfuser modelscope openpyxl
+pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
+pip install opencv-python-headless
+pip install deepspeed==0.17.0 numpy==1.26.4
+```
+
+**方式 3：使用 Docker**
+
+```bash
+# Pull the image
+docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+
+# Enter the container
+docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+```
+
+---
+
+## 三、下载预训练模型
+
+One-Forcing 需要**两个**预训练模型：
+
+- **Wan2.1-T2V-1.3B**：生成器/判别器的基础模型。
+- **Wan2.1-T2V-14B**：非因果的 real-score teacher，用于计算真实分布得分。
+
+```bash
+# Create the model directory
+mkdir -p models/Diffusion_Transformer
+
+# Download the Wan2.1 T2V 1.3B student base model
+modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
+
+# Download the Wan2.1 T2V 14B real-score teacher
+modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B
+```
+
+One-Forcing 官方 checkpoint 托管在 HuggingFace（`JiaqiFeng/OneForcing`）：
+
+```bash
+# 训练用热启动生成器（官方 ODE 初始化 checkpoint），通过 --ode_transformer_path 加载。
+hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
+
+# 推理用已训练 One-Forcing 生成器（见第六节）。
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+如果你自己训练了 Causal-Forcing Stage 1/2，也可以将 `--ode_transformer_path` 指向你产出的 Stage-2 CCD checkpoint。
+
+---
+
+## 四、数据准备
+
+One-Forcing 默认开启 GAN 分支（`--gan_g_weight=0.03 --gan_d_weight=0.03`），因此数据加载始终使用 **ImageVideoDataset**：对抗 critic 需要**带描述的真实视频**来构建干净 latent，仅有提示词的 JSON **不够用**（这一点与普通 DMD 不同）。
+
+### 4.1 快速测试数据集
+
+```bash
+# Download the official demo dataset
+modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
+```
+
+### 4.2 数据集结构
+
+与 Causal-Forcing Stage 1/2 相同。详见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md#42-数据集结构)。
+
+### 4.3 metadata.json 格式
+
+每条记录需要 `file_path`、`type="video"` 与 `text`；启动脚本使用 `metadata_add_width_height.json` 变体，额外携带 `width`/`height`，分桶训练无需现场解码媒体文件：
+
+```json
+[
+  {
+    "file_path": "train/00000000.mp4",
+    "text": "A young woman gently turns her head to the right, tilting it slightly as she offers a soft greeting...",
+    "type": "video",
+    "width": 832,
+    "height": 480
+  }
+]
+```
+
+> **说明**：每条视频须有足够的帧数以满足 `--video_sample_stride=1` 与 `--video_sample_n_frames=81`：当源视频产出的 latent 帧数少于 `score_num_frames=21` 时，GAN 分支会中止（见 5.3）。
+
+---
+
+## 五、训练
+
+### 5.1 快速开始
+
+可直接使用的启动脚本为 [train_one_forcing.sh](./train_one_forcing.sh)。它调用 One-Forcing 训练器（`train_one_forcing.py`）并开启 GAN 配方：
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
+
+accelerate launch --mixed_precision="bf16" --use_fsdp \
+    --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+    scripts/wan2.1_one_forcing/train_one_forcing.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --ode_transformer_path=$STAGE2_CKPT \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --fix_sample_size 480 832 \
+  --video_sample_stride=1 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --max_train_steps=200 \
+  --checkpointing_steps=100 \
+  --learning_rate=1.0e-05 \
+  --learning_rate_critic=1.0e-05 \
+  --lr_scheduler="constant" \
+  --seed=42 \
+  --output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.01 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=10.0 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --num_frame_per_block=1 \
+  --use_kv_cache_training \
+  --use_ema \
+  --ema_decay=0.99 \
+  --denoising_step_indices_list 1000 \
+  --real_guidance_scale=6.0 \
+  --score_timestep_clamp \
+  --fake_guidance_scale=0.0 \
+  --gen_update_interval=5 \
+  --gan_g_weight=0.03 \
+  --gan_d_weight=0.03 \
+  --gan_feature_layers 21 29 \
+  --gan_block_ffn_dim=2048 \
+  --gan_block_num_heads=12 \
+  --resume_from_checkpoint="latest" \
+  --trainable_modules "."
+```
+
+或直接运行：
+
+```bash
+bash scripts/wan2.1_one_forcing/train_one_forcing.sh
+```
+
+### 5.2 配方参数
+
+One-Forcing 配方设置以下 launcher 参数：
+
+| 参数 | Launcher 参数 | 取值 |
+|--------------------|---------------|------|
+| `denoising_step_list: [1000]` | `--denoising_step_indices_list` | `1000`（单步） |
+| `num_frame_per_block: 1` | `--num_frame_per_block` | `1`（逐帧） |
+| `num_training_frames: 21` / `score_num_frames: 21` | `--video_sample_n_frames=81` @ 480×832 | 21 个 latent 帧 |
+| `guidance_scale: 5.0` | `--real_guidance_scale` | `6.0`（见下方 CFG 说明） |
+| `lr: 1e-5` / `lr_critic: 1e-5` | `--learning_rate` / `--learning_rate_critic` | `1.0e-05` |
+| `beta1: 0.0`、`beta2: 0.999`、`weight_decay: 0.01` | `--adam_beta1`、`--adam_beta2`、`--adam_weight_decay` | `0.0`、`0.999`、`0.01` |
+| `ema_weight: 0.99` | `--use_ema`、`--ema_decay` | 开启、`0.99` |
+| `max_steps: 200` | `--max_train_steps` | `200` |
+| `dfake_gen_update_ratio: 5` | `--gen_update_interval` | `5` |
+| `gan_feature_layers: [21, 29]` | `--gan_feature_layers` | `21 29` |
+| `gan_block_ffn_dim: 2048` | `--gan_block_ffn_dim` | `2048` |
+| `gan_block_num_heads: 12` | `--gan_block_num_heads` | `12` |
+| `gan_g_weight: 0.03`、`gan_d_weight: 0.03` | `--gan_g_weight`、`--gan_d_weight` | `0.03`、`0.03` |
+| `relativistic_discriminator: false` | `--relativistic_discriminator` | 不传（关闭） |
+| `r1_weight: 0.0`、`r2_weight: 0.0` | `--r1_weight`、`--r2_weight` | 不传（关闭） |
+| `context_noise: 0` | `--context_noise` | `0`（默认） |
+
+> **CFG 约定**：本仓库的引导计算为 `cond + scale * (cond - uncond)`，本配方取 `--real_guidance_scale=6.0`。
+
+与 DMD 阶段共用的通用训练/视频采样参数，参见 [README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md#52-关键参数) 5.2 节。
+
+### 5.3 GAN 前置要求
+
+GAN 分支在 `train_one_forcing.py` 启动时会被校验，要求：
+
+- **带描述的真实视频**（`type="video"` + `file_path` + `text`），不能只提供提示词 JSON。任一 GAN/R1/R2 权重为正时，数据加载会从 `TextDataset` 切换为 `ImageVideoDataset`。
+- **`--gradient_accumulation_steps=1`** 且**不使用 DeepSpeed** —— 否则训练器抛出 `ValueError`。
+- **`--video_sample_stride=1`**：当 `stride=2` 时，81 帧源仅得到约 10 个 latent 帧，低于 `score_num_frames=21` 而中止。
+- fake-score backbone 上**有效的特征层索引**。默认 `21 29`（层）与 `12`（头）适用于 Wan2.1-T2V-1.3B；头数须整除 fake-score 隐藏维度。
+
+可训练的 fake-score 与判别器头 master weights 保持 **FP32**；前向按配置使用 autocast/FSDP 混合精度，logits/softplus/有限差分归约在 FP32 下计算。`--fsdp_reduce_fp32` 只改变梯度通信精度。
+
+### 5.4 可选正则（相对论 / R1 / R2）
+
+这些正则默认**关闭**。若要实验，向 launcher 追加：
+
+```bash
+# Relativistic discriminator: use paired real/fake logit differences
+--relativistic_discriminator
+
+# Finite-difference gradient penalties (APT form, not double-backward)
+--r1_weight=1.0 --r1_sigma=0.1     # on noisy real latents:   r1_weight * mean((D_real' - D_real)^2)
+--r2_weight=0.001 --r2_sigma=0.01  # on noisy gen latents: 0.5 * r2_weight * mean(((D_fake' - D_fake)/sigma)^2)
+```
+
+设 logits 为 `D_fake` 与 `D_real`，普通 GAN 损失为 `softplus(-D_fake)`（生成器）与 `softplus(-D_real) + softplus(D_fake)`（判别器）；相对论模式改用 logit 差值。真实/生成样本共用时间步与噪声，扰动噪声独立抽样。
+
+---
+
+## 六、推理
+
+### 6.1 快速开始
+
+推理入口为 [examples/wan2.1_one_forcing/predict_t2v.py](../../examples/wan2.1_one_forcing/predict_t2v.py)。它加载 **1.3B 学生模型**（14B 仅在训练时作为 DMD teacher），运行逐帧单步生成。默认使用 **One-Forcing 官方 checkpoint**：
+
+```bash
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+```python
+# examples/wan2.1_one_forcing/predict_t2v.py
+model_name          = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+transformer_path    = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
+use_ema             = False   # only applies to trainer-output checkpoint dirs
+stochastic_sampling = True    # DMD-distilled student
+guidance_scale      = 1.0     # CFG is baked into the distilled weights
+num_inference_steps = 1       # One-Forcing one-step generation
+num_frame_per_block = 1       # frame-wise
+sample_size         = [480, 832]
+video_length        = 81
+```
+
+运行：
+
+```bash
+python examples/wan2.1_one_forcing/predict_t2v.py
+```
+
+`transformer_path` 既可以是官方 checkpoint 文件、纯 `.pt`/`.safetensors` 文件，也可以是训练输出的 **checkpoint 目录**（其 live/EMA 权重会被自动解析）。GAN 判别器头是 fake_score 的训练专用子模块，因此推理**无需任何判别器接线**。
+
+### 6.2 单步与 4 步选择器
+
+当 `stochastic_sampling=True` 时，`num_inference_steps` 决定 warped 时间步调度：
+
+| `num_inference_steps` | 去噪时间步 | 适用场景 |
+|-----------------------|------------|----------|
+| `1`（One-Forcing 默认） | `[1000]` | 最快；逐帧单步生成 |
+| `4` | `[1000, 750, 500, 250]` | 更稳定、更慢 |
+
+### 6.3 首帧调度
+
+共享的 `WanSelfForcingPipeline` 对**所有块使用同一份去噪调度**，因此此处不单独为首帧设置调度（例如首帧 `[1000, 750, 500, 250]`、后续帧 `[1000]`）；该 pipeline 被多个 example 共用，有意保持不变。
+
+---
+
+## 七、Checkpoint 往返（训练→推理）
+
+FSDP 训练器将每个 `checkpoint-{step}/` 按如下布局写出，因此训练产物可直接喂给推理，无需改动 key：
+
+| 组件 | 保存位置 | 是否含 `gan_head`？ |
+|------|----------|---------------------|
+| **生成器**（live） | `checkpoint-{step}/diffusion_pytorch_model.safetensors`（根目录，经 FSDP hook） | 否 |
+| **生成器**（EMA） | `checkpoint-{step}/transformer_ema/` | 否 |
+| **Fake-score + 判别器头** | `checkpoint-{step}/fake_score/` | 是（FP32 判别器头张量） |
+
+- **生成器从不包含 `gan_head`** 键（判别器头挂在 fake-score 上），因此 `predict_t2v.py` 可用 `strict=False` 干净加载。
+- 推理脚本中的 `_resolve_transformer_path` 依次解析：`transformer_ema/`、`ema_transformer/`（兼容别名）、`transformer/`，最后到根目录 safetensors。
+- **恢复** GAN 训练时须使用相同的 GAN/判别器头配置：加载前会根据 CLI 参数重建判别器头，并从 `fake_score/` 还原。
+
+---
+
+## 八、更多资源
+
+- **One-Forcing（源配方）**：https://github.com/Aurora-edu/One-Forcing
+- **Causal-Forcing（DMD 阶段）**：https://github.com/thu-ml/Causal-Forcing
+- **Stage 1（AR Diffusion）训练指南**：[README_TRAIN_AR_DIFFUSION_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md)
+- **Stage 2（CCD）训练指南**：[README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md)
+- **底层 DMD 训练器指南**：[README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md)
+- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun

scripts/wan2.1_one_forcing/train_one_forcing.sh

@@ -0,0 +1,85 @@
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
+# One-Forcing needs REAL videos to build the adversarial critic's clean_latents, so the dataset
+# must be a video meta (type=video + file_path), NOT a prompt-only json. The official demo
+# dataset X-Fun-Videos-Demo ships captioned videos and is normalized by --enable_bucket
+# --random_hw_adapt --training_with_video_token_length.
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+# Warm-start generator: a distilled/ODE checkpoint that the DMD + adversarial critic then polish.
+# Default: the official One-Forcing ODE-initialized checkpoint, downloaded via
+#   hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
+# You can also point it at your own Causal-Forcing Stage-2 CCD checkpoint.
+export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+# One-Forcing = Causal-Forcing Stage-3 DMD + a shared noised-latent adversarial critic on fake_score.
+# - 1-step frame-wise generation: --denoising_step_indices_list 1000, --num_frame_per_block=1;
+#   21 latent frames from --video_sample_n_frames=81 @ 480x832.
+# - GAN critic: gan_g_weight/gan_d_weight=0.03, feature layers [21, 29], ffn 2048, heads 12,
+#   non-relativistic, R1/R2 off. The head rides on fake_score and is training-only.
+# - GAN requires real videos, --gradient_accumulation_steps=1, no DeepSpeed, and --video_sample_stride=1
+#   (stride=2 on 81-frame sources yields ~10 latent frames < score_num_frames=21 and aborts).
+# - guidance: local CFG is cond + scale*(cond - uncond), so --real_guidance_scale=6.0.
+# - optimizer/EMA/steps: lr=lr_critic=1e-5, betas (0.0, 0.999), weight_decay=0.01,
+#   --use_ema --ema_decay=0.99, gen_update_interval=5, max_steps=200.
+accelerate launch --mixed_precision="bf16" --use_fsdp \
+    --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+    scripts/wan2.1_one_forcing/train_one_forcing.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --ode_transformer_path=$STAGE2_CKPT \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --fix_sample_size 480 832 \
+  --video_sample_stride=1 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --max_train_steps=200 \
+  --checkpointing_steps=100 \
+  --learning_rate=1.0e-05 \
+  --learning_rate_critic=1.0e-05 \
+  --lr_scheduler="constant" \
+  --seed=42 \
+  --output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.01 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=10.0 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --num_frame_per_block=1 \
+  --use_kv_cache_training \
+  --use_ema \
+  --ema_decay=0.99 \
+  --denoising_step_indices_list 1000 \
+  --real_guidance_scale=6.0 \
+  --score_timestep_clamp \
+  --fake_guidance_scale=0.0 \
+  --gen_update_interval=5 \
+  --gan_g_weight=0.03 \
+  --gan_d_weight=0.03 \
+  --gan_feature_layers 21 29 \
+  --gan_block_ffn_dim=2048 \
+  --gan_block_num_heads=12 \
+  --resume_from_checkpoint="latest" \
+  --trainable_modules "."

scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md

@@ -179,6 +179,7 @@ The difference between DeepSpeed-Zero-2 and FSDP lies in whether to shard model
 
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -189,6 +190,7 @@ NCCL_DEBUG=INFO
 accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -204,21 +206,26 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \
@@ -247,17 +254,19 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | `--num_train_epochs` | Number of training epochs | 100 |
 | `--checkpointing_steps` | Save checkpoint every N steps | 50 |
 | `--learning_rate` | Initial learning rate (generator) | 2e-06 |
-| `--learning_rate_critic` | Initial learning rate (critic) | 2e-06 |
+| `--learning_rate_critic` | Initial learning rate (critic) | 4.0e-07 |
 | `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` |
 | `--lr_warmup_steps` | Learning rate warmup steps | 100 |
 | `--seed` | Random seed | 42 |
 | `--output_dir` | Output directory | `output_dir_wan2.1_self_forcing_distill` |
 | `--gradient_checkpointing` | Enable gradient checkpointing | - |
 | `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
-| `--adam_weight_decay` | AdamW weight decay | 3e-2 |
+| `--adam_weight_decay` | AdamW weight decay | 0.0 |
+| `--adam_beta1` | AdamW beta1 | 0.0 |
+| `--adam_beta2` | AdamW beta2 | 0.999 |
 | `--adam_epsilon` | AdamW epsilon | 1e-10 |
 | `--vae_mini_batch` | VAE encoding mini-batch size | 1 |
-| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
+| `--max_grad_norm` | Gradient clipping threshold | 10.0 |
 | `--fsdp_reduce_fp32` | Optional: reduce gradients in fp32 under FSDP mixed precision (default reduces in bf16) | - |
 | `--enable_bucket` | Enable bucket training, no cropping, group by resolution | - |
 | `--random_hw_adapt` | Auto-scale images/videos to random sizes in `[min_size, max_size]` range | - |
@@ -277,7 +286,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 |-----|------|-------|
 | `--denoising_step_indices_list` | Denoising step indices list (core distillation parameter) | `1000 750 500 250` |
 | `--randomize_step_indices` | Randomize the denoising timestep indices during training | - |
-| `--real_guidance_scale` | Real guidance scale for scoring | 6.0 |
+| `--real_guidance_scale` | Real guidance scale for scoring | 4.0 |
 | `--score_timestep_clamp` | Optional: clamp DMD/critic timesteps to `[0.02, 0.98] * num_train_timesteps` | - |
 | `--fake_guidance_scale` | Fake guidance scale for scoring | 0.0 |
 | `--gen_update_interval` | Generator update interval | 5 |
@@ -365,6 +374,7 @@ You can configure validation parameters to periodically generate test videos dur
 
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -375,6 +385,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -390,21 +401,26 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \
@@ -427,6 +443,7 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/
 Training shell command is as follows:
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -437,6 +454,7 @@ NCCL_DEBUG=INFO
 accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -452,21 +470,26 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \
@@ -481,6 +504,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
 
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -491,6 +515,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -506,21 +531,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \
@@ -540,6 +570,7 @@ Assuming 2 machines, each with 8 GPUs:
 **Machine 0 (Master)**:
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # Master machine IP
@@ -555,6 +586,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -570,21 +602,26 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \
@@ -596,6 +633,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
 **Machine 1 (Worker)**:
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # Same as Master

scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md

@@ -180,6 +180,7 @@ hf download gdhe17/Self-Forcing --local-dir models/Diffusion_Transformer/Self-Fo
 
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -190,6 +191,7 @@ NCCL_DEBUG=INFO
 accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -205,21 +207,26 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \
@@ -248,17 +255,19 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | `--num_train_epochs` | 训练 epoch 数 | 100 |
 | `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 |
 | `--learning_rate` | 初始学习率（生成器） | 2e-06 |
-| `--learning_rate_critic` | 初始学习率（判别器） | 2e-06 |
+| `--learning_rate_critic` | 初始学习率（判别器） | 4.0e-07 |
 | `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
 | `--lr_warmup_steps` | 学习率预热步数 | 100 |
 | `--seed` | 随机种子 | 42 |
 | `--output_dir` | 输出目录 | `output_dir_wan2.1_self_forcing_distill` |
 | `--gradient_checkpointing` | 激活重计算 | - |
 | `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
-| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
+| `--adam_weight_decay` | AdamW 权重衰减 | 0.0 |
+| `--adam_beta1` | AdamW beta1 | 0.0 |
+| `--adam_beta2` | AdamW beta2 | 0.999 |
 | `--adam_epsilon` | AdamW epsilon 值 | 1e-10 |
 | `--vae_mini_batch` | VAE 编码时的迷你批次大小 | 1 |
-| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
+| `--max_grad_norm` | 梯度裁剪阈值 | 10.0 |
 | `--fsdp_reduce_fp32` | 可选：FSDP 混精下梯度以 fp32 做跨卡 reduce（默认 bf16） | - |
 | `--enable_bucket` | 启用分桶训练，不裁剪图片/视频，按分辨率分组训练 | - |
 | `--random_hw_adapt` | 自动缩放图片/视频到 `[min_size, max_size]` 范围内的随机尺寸 | - |
@@ -278,7 +287,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 |-----|------|-------|
 | `--denoising_step_indices_list` | 去噪步骤列表（蒸馏核心参数） | `1000 750 500 250` |
 | `--randomize_step_indices` | 训练时随机化去噪 timestep 索引 | - |
-| `--real_guidance_scale` | 用于评分的真实 guidance scale | 6.0 |
+| `--real_guidance_scale` | 用于评分的真实 guidance scale | 4.0 |
 | `--score_timestep_clamp` | 可选：将 DMD/critic 的 timestep 截断到 `[0.02, 0.98] * num_train_timesteps` | - |
 | `--fake_guidance_scale` | 用于评分的虚拟 guidance scale | 0.0 |
 | `--gen_update_interval` | 生成器更新间隔 | 5 |
@@ -366,6 +375,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -376,6 +386,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -391,21 +402,26 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \
@@ -428,6 +444,7 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/
 训练 shell 命令如下：
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -438,6 +455,7 @@ NCCL_DEBUG=INFO
 accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -453,21 +471,26 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \
@@ -482,6 +505,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
 
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -492,6 +516,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -507,21 +532,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \
@@ -541,6 +571,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
 **机器 0（Master）**：
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # Master 机器 IP
@@ -556,6 +587,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -571,21 +603,26 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \
@@ -597,6 +634,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
 **机器 1（Worker）**：
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # 与 Master 相同

scripts/wan2.1_self_forcing/train_distill.sh

@@ -1,4 +1,5 @@
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/internal_datasets/"
 export DATASET_META_NAME="datasets/internal_datasets/metadata.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -9,6 +10,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
   --config_path="config/wan2.1/wan_civitai.yaml" \
   --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
   --train_data_dir=$DATASET_NAME \
   --train_data_meta=$DATASET_META_NAME \
   --image_sample_size=640 \
@@ -24,21 +26,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
   --num_train_epochs=100 \
   --checkpointing_steps=50 \
   --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
   --lr_scheduler="constant_with_warmup" \
   --lr_warmup_steps=100 \
   --seed=42 \
   --output_dir="output_dir_wan2.1_self_forcing_distill" \
   --gradient_checkpointing \
   --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
   --adam_epsilon=1e-10 \
   --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
   --random_hw_adapt \
   --training_with_video_token_length \
   --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
   --use_kv_cache_training \
   --num_frame_per_block=3 \
   --train_mode="normal" \

videox_fun/dist/__init__.py

@@ -21,10 +21,13 @@ from .longcatvideo_xfuser import (usp_attn_longcatvideo_avatar_forward,
                                   usp_rope_longcatvideo_forward)
 from .ltx2_xfuser import (LTX2MultiGPUsAttnProcessor,
                           LTX2PerturbedMultiGPUsAttnProcessor)
-from .minimax_h3_xfuser import MiniMaxH3MultiGPUsAttnProcessor
+from .minimax_h3_xfuser import (MiniMaxH3MultiGPUsAttnProcessor,
+                                MiniMaxH3MultiGPUsVSAAttnProcessor)
 from .qwen_xfuser import QwenImageMultiGPUsAttnProcessor2_0
 from .qwenimage21_xfuser import QwenImage21MultiGPUsAttnProcessor
-from .wan_xfuser import (usp_attn_forward, usp_attn_s2v_forward,
+from .wan_xfuser import (usp_attn_forward, usp_attn_fastvideo_forward,
+                         usp_attn_s2v_forward,
                          usp_attn_self_forcing_forward,
-                         usp_attn_sparkwan_forward)
+                         usp_attn_sparkwan_forward,
+                         usp_attn_turbowan_forward)
 from .z_image_xfuser import ZMultiGPUsSingleStreamAttnProcessor

videox_fun/dist/fuser.py

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

videox_fun/dist/minimax_h3_xfuser.py

@@ -1,5 +1,7 @@
 import torch
+import torch.nn.functional as F
 
+from .fuser import get_sp_group, ulysses_all_to_all
 from .fuser import xFuserLongContextAttention
 
 
@@ -80,3 +82,99 @@ class MiniMaxH3MultiGPUsAttnProcessor:
         hidden_states = attn.to_out[0](hidden_states)
         hidden_states = attn.to_out[1](hidden_states)
         return hidden_states
+
+
+class MiniMaxH3MultiGPUsVSAAttnProcessor:
+    """Sequence-parallel attention processor for the FastVideo VSA student on MiniMax-H3.
+
+    Mirrors ``MiniMaxH3MultiGPUsAttnProcessor``'s row-split -> head-split Ulysses scheme, but swaps
+    the dense ``xFuserLongContextAttention`` core for the VSA block-sparse op (``attn.attn_op``) so
+    the trained coarse-to-fine selection and the ``to_gate_compress`` compensation branch are kept
+    instead of being silently dropped by a dense fallback.
+
+    Two VSA-specific points under head-parallelism:
+      * ``to_gate_compress`` is a ``Linear(dim, heads*head_dim)`` that MIXES heads, so it cannot run
+        in the head-shard layout. It is strictly per-token and SP shards rows while keeping the full
+        hidden dim, so the gate is evaluated here on this rank's ROW shard (yielding all heads), then
+        all-to-all'd alongside q/k/v -- element-wise identical to computing it on the full sequence
+        and slicing heads, and it lands on the same contiguous head block the all-to-all hands back.
+      * the parent forward pads the packed sequence up to a multiple of the group size, so the
+        gathered sequence carries ``<= sp_world_size - 1`` zero rows at its tail. VSA builds its tile
+        geometry from ``vsa_segments``, whose lengths sum to the PADLESS ``valid_length``; the padding
+        tail is therefore sliced off before the op and the output padded back afterwards. Unlike the
+        dense processor (which lets the pad rows dilute the softmax by ~1e-5), this keeps VSA
+        bitwise-equal to the single-GPU run -- the pad rows never enter the kernel.
+
+    Requires a padless packed sequence (``attention_mask is None``) and a pure-Ulysses split
+    (``num_heads`` divisible by the group size, checked by the model's ``enable_multi_gpus_inference``).
+    A dense student (``topk >= 1.0``) carries no ``attn_op`` and must keep the dense processor.
+    """
+
+    def __call__(
+        self,
+        attn,
+        hidden_states: torch.Tensor,
+        rotary_emb=None,
+        attention_mask=None,
+        valid_length=None,
+        vsa_segments=None,
+    ) -> torch.Tensor:
+        from ..models.minimax_h3_transformer3d import \
+            apply_minimax_h3_rotary_emb
+
+        if attention_mask is not None:
+            raise ValueError(
+                "MiniMaxH3MultiGPUsVSAAttnProcessor does not support a masked (padded) packed sequence.")
+        if not hasattr(attn, "attn_op"):
+            raise RuntimeError(
+                "The VSA sequence-parallel processor needs attn.attn_op; a dense student (topk>=1.0) "
+                "must bind MiniMaxH3MultiGPUsAttnProcessor instead.")
+
+        query = attn.to_q(hidden_states).unflatten(-1, (attn.heads, -1))
+        key = attn.to_k(hidden_states).unflatten(-1, (attn.heads, -1))
+        value = attn.to_v(hidden_states).unflatten(-1, (attn.heads, -1))
+
+        query = attn.norm_q(query)
+        key = attn.norm_k(key)
+
+        # The rotary coordinates are this rank's rows, so apply them before the all-to-all moves rows.
+        if rotary_emb is not None:
+            query = apply_minimax_h3_rotary_emb(query, *rotary_emb)
+            key = apply_minimax_h3_rotary_emb(key, *rotary_emb)
+
+        # Cross-head compress gate on this rank's ROW shard (full hidden dim present): [B, chunk, H, D].
+        gate = attn.attn_op.to_gate_compress(hidden_states).unflatten(-1, (attn.heads, -1))
+
+        half_dtypes = (torch.float16, torch.bfloat16)
+
+        def half(x):
+            return x if x.dtype in half_dtypes else x.to(torch.bfloat16)
+
+        # row-split -> head-split over the whole (padded) packed sequence, gate alongside q/k/v.
+        query = ulysses_all_to_all(half(query), scatter_dim=2, gather_dim=1)
+        key = ulysses_all_to_all(half(key), scatter_dim=2, gather_dim=1)
+        value = ulysses_all_to_all(half(value), scatter_dim=2, gather_dim=1)
+        gate = ulysses_all_to_all(half(gate), scatter_dim=2, gather_dim=1)
+
+        # Drop the SP padding tail so VSA sees exactly the padless length its segments describe.
+        padded_length = query.shape[1]
+        trim = valid_length is not None and valid_length < padded_length
+        if trim:
+            query = query[:, :valid_length]
+            key = key[:, :valid_length]
+            value = value[:, :valid_length]
+            gate = gate[:, :valid_length]
+
+        out = attn.attn_op(query, key, value, None, segments=vsa_segments, gate=gate)
+
+        # Pad the output back to the gathered length before the inverse all-to-all; the zero rows are
+        # dropped by the model's final all_gather and never touch a real row's attention.
+        if trim:
+            out = F.pad(out, (0, 0, 0, 0, 0, padded_length - valid_length))
+
+        # head-split -> row-split, back to this rank's rows.
+        out = ulysses_all_to_all(out.to(hidden_states.dtype), scatter_dim=1, gather_dim=2)
+        out = out.flatten(2, 3).type_as(query)
+        out = attn.to_out[0](out)
+        out = attn.to_out[1](out)
+        return out

videox_fun/dist/wan_xfuser.py

@@ -224,7 +224,10 @@ def usp_attn_sparkwan_forward(self,
         q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
         k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
         v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
-        x = self.attn_op(q, k, v)
+        # heads_total/head_start keep the SLA linear branch's fp32 normalizer sum bitwise-equal to
+        # the single-GPU result despite this rank holding only the H/P head block.
+        h_per = n // sp_size
+        x = self.attn_op(q, k, v, heads_total=n, head_start=sp_rank * h_per)
         x = ulysses_all_to_all(x, scatter_dim=1, gather_dim=2)
         x = x.flatten(2)
         x = self.o(x)
@@ -267,6 +270,140 @@ def usp_attn_sparkwan_forward(self,
     out = self.o(out)
     return out
 
+def usp_attn_turbowan_forward(self,
+                     x,
+                     seq_lens,
+                     grid_sizes,
+                     freqs,
+                     dtype=torch.bfloat16,
+                     t=0):
+    """Sequence-parallel self-attention for the TurboDiffusion SLA student.
+
+    ``TurboWanTransformer3DModel`` swaps ``WanSelfAttention`` for a local variant whose kernel
+    is ``self.attn_op`` (``FakeLocalAttention`` -> ``SparseLinearAttention``, block-sparse exact
+    softmax + a ``proj_l`` feature-map linear branch). The stock ``usp_attn_forward`` binds
+    ``xFuserLongContextAttention``'s dense full-softmax and only touches the base-class
+    ``q/k/v/norm_q/norm_k/o`` attributes, so patching it onto this model would silently drop the
+    SLA kernel -- every needed attribute is inherited, hence no AttributeError, no shape error,
+    just dense numbers. This variant keeps ``self.attn_op`` live under Ulysses head-parallelism.
+
+    Equivalence: ``SparseLinearAttention.forward`` takes and returns ``[B, L, H, D]``, transposes
+    to BHLD internally and does every reduction over L/D only (``get_block_map`` pooling and
+    top-k, the sparse kernel, the ``k^T v`` linear branch), and ``proj_l`` is a per-head
+    ``head_dim -> head_dim`` projection -- no cross-head mixing and no head-indexed parameter,
+    so an all-to-all from the sequence shard ``[B, S/P, H, D]`` to the head shard
+    ``[B, S, H/P, D]`` reproduces the single-GPU result exactly (top-k is a per-query key-block
+    ratio). RoPE is applied BEFORE the all-to-all because ``rope_apply`` slices the frequency
+    table by sequence-parallel rank.
+
+    Requires a pure-Ulysses split (``ring_degree`` must be 1: the SLA kernel exposes no partial
+    output / logsumexp interface, so ring's online rescale-concat cannot be wired to it) and
+    ``num_heads`` divisible by the SP degree.
+    """
+    # No sparse kernel on this instance: keep the stock dense SP path.
+    if not hasattr(self, "attn_op"):
+        return usp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=dtype, t=t)
+
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    half_dtypes = (torch.float16, torch.bfloat16)
+
+    def half(t_):
+        return t_ if t_.dtype in half_dtypes else t_.to(dtype)
+
+    # query, key, value function
+    def qkv_fn(x):
+        q = self.norm_q(self.q(x.to(dtype))).view(b, s, n, d)
+        k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d)
+        v = self.v(x.to(dtype)).view(b, s, n, d)
+        return q, k, v
+
+    q, k, v = qkv_fn(x)
+    q, k = rope_apply_qk(q, k, grid_sizes, freqs)
+
+    # seq-shard -> head-shard, SLA over the FULL sequence on a subset of heads, then back.
+    q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
+    k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
+    v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
+    # heads_total/head_start keep the SLA linear branch's fp32 normalizer sum bitwise-equal to the
+    # single-GPU result despite this rank holding only the H/P head block.
+    sp_size = get_sequence_parallel_world_size()
+    sp_rank = get_sequence_parallel_rank()
+    h_per = n // sp_size
+    x = self.attn_op(q, k, v, heads_total=n, head_start=sp_rank * h_per)
+    x = ulysses_all_to_all(x.to(dtype), scatter_dim=1, gather_dim=2)
+
+    # output
+    x = x.flatten(2)
+    x = self.o(x)
+    return x
+
+def usp_attn_fastvideo_forward(self,
+                     x,
+                     seq_lens,
+                     grid_sizes,
+                     freqs,
+                     dtype=torch.bfloat16,
+                     t=0):
+    """Sequence-parallel self-attention for the FastVideo VSA student.
+
+    Same Ulysses head-parallel scheme as ``usp_attn_sparkwan_forward``, with one VSA-specific
+    wrinkle: ``VSASparseAttention.to_gate_compress`` is a ``Linear(dim, num_heads*head_dim)``
+    that MIXES heads, so it cannot run inside the head-shard layout the way the rest of the VSA
+    kernel can (``sparse`` / ``compressed`` only reduce over L/D and are head-independent).
+    It is, however, strictly per-token -- and sequence parallelism shards the sequence while
+    leaving the hidden dim intact, so every rank still holds the full ``C`` for its own tokens.
+    The gate is therefore evaluated here on the SEQUENCE shard (yielding all ``H`` heads), then
+    all-to-all'd alongside q/k/v, which is element-wise identical to computing it on the full
+    sequence and slicing heads -- and lands on exactly the contiguous head block
+    ``[sp_rank*H/P, (sp_rank+1)*H/P)`` the all-to-all hands this rank.
+
+    The gathered sequence is ``s * sp_size``, which the parent forward pads to a multiple of the
+    SP degree, so it can exceed ``grid_sizes.prod()``; ``_vsa_dispatch`` routes that case to its
+    per-sample branch (attend over the valid grid tokens, pad the output back), keeping the tile
+    geometry built from the true grid.
+
+    Requires a pure-Ulysses split (``ring_degree`` must be 1) and ``num_heads`` divisible by the
+    SP degree. ``topk >= 1.0`` (the dense fallback) defers to ``usp_attn_forward``.
+    """
+    # Dense student (topk >= 1.0) or a plain Wan self-attention: keep the stock SP path.
+    if getattr(self, "topk", 1.0) >= 1.0 or not hasattr(self, "attn_op"):
+        return usp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=dtype, t=t)
+
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    half_dtypes = (torch.float16, torch.bfloat16)
+
+    def half(t_):
+        return t_ if t_.dtype in half_dtypes else t_.to(dtype)
+
+    # query, key, value function
+    def qkv_fn(x):
+        q = self.norm_q(self.q(x.to(dtype))).view(b, s, n, d)
+        k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d)
+        v = self.v(x.to(dtype)).view(b, s, n, d)
+        return q, k, v
+
+    q, k, v = qkv_fn(x)
+    q, k = rope_apply_qk(q, k, grid_sizes, freqs)
+
+    # Cross-head compress gate, evaluated on this rank's sequence shard where the full hidden
+    # dim is available: [B, S/P, C] -> [B, S/P, H, D] (all heads).
+    gate = self.attn_op.to_gate_compress(x.to(dtype)).view(b, s, n, d)
+
+    # seq-shard -> head-shard for q/k/v and the gate alike.
+    q = ulysses_all_to_all(half(q), scatter_dim=2, gather_dim=1)
+    k = ulysses_all_to_all(half(k), scatter_dim=2, gather_dim=1)
+    v = ulysses_all_to_all(half(v), scatter_dim=2, gather_dim=1)
+    gate = ulysses_all_to_all(half(gate), scatter_dim=2, gather_dim=1)
+
+    # VSA over the FULL sequence on a head shard; x is unused once gate is supplied.
+    out = self._vsa_dispatch(q, k, v, None, seq_lens, grid_sizes, gate=gate)
+
+    # head-shard -> seq-shard, back to this rank's sequence chunk.
+    out = ulysses_all_to_all(out.to(dtype), scatter_dim=1, gather_dim=2)
+    out = out.flatten(2)
+    out = self.o(out)
+    return out
+
 @amp.autocast(enabled=False)
 @torch.compiler.disable()
 def s2v_rope_apply(x, grid_sizes, freqs):

videox_fun/models/attention_kernel.py

@@ -51,17 +51,6 @@ import triton.language as tl
 # loads, a fused tl.dot(p, v, acc) accumulate, and a per-query-tile variable kv-block count (vs a
 # uniform-MAXK loop over a dummy-padded tile).
 #
-# This is deliberately NOT the fastest forward available. That same sm90 CUDA kernel (ThunderKittens)
-# IS buildable from source on this box -- only its prebuilt wheels are pinned -- and measures
-# 4.66 ms/call on the workload above, i.e. 1.35x over this Triton kernel (step-level VSA-vs-dense
-# 1.019x -> 1.115x at 480P, 1.233x -> 1.381x at 720P). It stays out of the production path because
-# it adds a build dependency and fails SILENTLY three ways: on a non-sm90a device pass it compiles
-# to an empty stub returning uninitialised memory; it reinterpret_casts `variable_block_sizes` to
-# int32 with no dtype check, so an int64 tensor yields all-NaN; and its lse is [B,H,L,1] in the
-# natural-log domain scaled by -1/qk_scale rather than our [B,H,L] base-2. The build recipe, the
-# runtime gate it would need, and the measurement set are archived in
-# `_fastvideo_sm90_kernel/README.md`.
-#
 # Deviations from upstream, listed so a diff against the sdist still audits this block:
 #   1. its `import torch / triton / triton.language`, and its SPARSE-ADDITION `import math`, are
 #      hoisted to this module's top import block -- same names, no behavioural change, and the
@@ -495,12 +484,10 @@ class _sparse_linear_attention_varlen(torch.autograd.Function):
     @staticmethod
     def forward(ctx, q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes, kb_id, qk_scale):
         assert q_tile.is_contiguous() and k_tile.is_contiguous() and v_tile.is_contiguous()
-        B, H, L_BUF, D = q_tile.shape
-        # Forward runs the vendored official Triton block-sparse kernel, which consumes the
-        # official (q2k_index, q2k_num) pair map_to_index produced: the first q2k_num[row] entries
-        # of each row are that query tile's kv-tile ids and the -1 pads are never read. Its M is
-        # [B,H,L_BUF] base-2 lse (m_i + log2 l_i), byte-compatible with what our backward expects,
-        # so the save_for_backward contract is unchanged.
+        # The forward consumes the official (q2k_index, q2k_num) pair map_to_index produced:
+        # the first q2k_num[row] entries of each row are that query tile's kv-tile ids and the
+        # -1 pads are never read. It produces M as [B,H,L_BUF] base-2 lse (m_i + log2 l_i),
+        # which is exactly what the varlen backward kernels below consume.
         out, lse = triton_block_sparse_attn_forward(
             q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes)
         ctx.save_for_backward(q_tile, k_tile, v_tile, q2k_idx, q2k_num,
@@ -625,7 +612,13 @@ def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
     pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
 
     K = pooled_score.shape[-1]
-    topk = min(K, int(topk_ratio * K))
+    # Clamp to at least one key block. A truncating int() selects ZERO blocks whenever
+    # topk_ratio * K < 1 (short sequence and/or aggressive ratio -- e.g. RoLa's default 0.03
+    # below 2176 tokens, SLA's 0.1 below 640), and the sparse kernel then attends to nothing and
+    # emits all-NaN. Keeps the truncating count everywhere else, so no existing configuration
+    # changes its selection; matches the lower bound the VSA paths already apply
+    # (VSASparseAttention.sparse / _h3_compute_topk).
+    topk = max(1, min(K, int(topk_ratio * K)))
     lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices
 
     sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)

videox_fun/models/attention_utils.py

@@ -5,6 +5,7 @@ import warnings
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 try:
     import flash_attn_interface
@@ -344,6 +345,32 @@ def attention(
     return out
 
 
+def _head_shard_seq_sum(x, heads_total, head_start):
+    """``torch.sum(x, dim=-2, keepdim=True)`` that is bitwise-invariant to Ulysses head-sharding.
+
+    ``x`` is ``[B, h_local, L, D]``. The pipeline runs under autocast and ``torch.sum`` is in
+    PyTorch's float32 promotion list, so this normalizer accumulates in fp32 over ``L`` and
+    returns fp32; the CUDA ReduceConfig picks its accumulation tree from the output count
+    ``B*H*D``. When head-parallelism shrinks ``H`` (e.g. 12 -> 6 -> 3) the tree changes and the
+    fp32 sum moves in its last bits (~1e-7), which the linear branch's denominator passes through
+    undamped -- a bf16 output would absorb it, fp32 does not. Padding the head dim back to
+    ``heads_total`` makes every rank reduce over the full ``B*heads_total*D`` outputs, i.e. the
+    single-GPU tree, then this shard's own head block is sliced back out, so multi-GPU reproduces
+    the single-GPU result bitwise.
+
+    Uninitialized padding is safe: each output element reads only its own head's rows, so the
+    garbage heads (sliced away) never touch the real ones. ``heads_total=None`` or an already-full
+    head count falls back to the plain reduction, leaving the single-GPU path unchanged.
+    """
+    h_local = x.shape[1]
+    if heads_total is None or h_local == heads_total:
+        return torch.sum(x, dim=-2, keepdim=True)
+    buf = torch.empty(x.shape[0], heads_total, x.shape[2], x.shape[3],
+                      device=x.device, dtype=x.dtype)
+    buf[:, head_start:head_start + h_local].copy_(x)
+    return torch.sum(buf, dim=-2, keepdim=True)[:, head_start:head_start + h_local]
+
+
 class SparseLinearAttention(nn.Module):
     # Modified from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/kernel.py
     def __init__(self, head_dim, topk, feature_map='softmax', BLKQ=64, BLKK=64, use_bf16=True, tie_feature_map_qk=True):
@@ -390,13 +417,18 @@ class SparseLinearAttention(nn.Module):
             nn.init.zeros_(self.proj_l.weight)
             nn.init.zeros_(self.proj_l.bias)
 
-    def forward(self, q, k, v, return_sparsity=False):
+    def forward(self, q, k, v, return_sparsity=False, heads_total=None, head_start=0):
         R'''
         Args:
             q: queries  [B, L, H, D]
             k: keys     [B, L, H, D].
             v: values   [B, L, H, D].
             return_sparsity: whether to return the actual sparsity.
+            heads_total: full head count before Ulysses head-sharding, or None when unsharded.
+                Kept so the linear branch's fp32 normalizer sum stays bitwise-equal to the
+                single-GPU result under sequence parallelism (see _head_shard_seq_sum).
+            head_start: this shard's offset into the full head range; used only when heads_total
+                is not None.
         '''
         dtype = q.dtype
         
@@ -416,7 +448,7 @@ class SparseLinearAttention(nn.Module):
 
         def calc_linear(q, k, v):
             kvsum = k.transpose(-1, -2) @ v
-            ksum = torch.sum(k, dim=-2, keepdim=True)
+            ksum = _head_shard_seq_sum(k, heads_total, head_start)
             return (q @ kvsum) / (1e-5 + (q * ksum).sum(dim=-1, keepdim=True))
         o_l = calc_linear(q, k, v)
 
@@ -942,7 +974,7 @@ class VSASparseAttention(nn.Module):
         oc = oc_blocks.repeat_interleave(self.BLKQ, dim=2)[:, :, :L]
         return oc
 
-    def forward_h3(self, q, k, v, x, segments, return_sparsity=False):
+    def forward_h3(self, q, k, v, x, segments, return_sparsity=False, gate=None):
         R'''Segment-aware H3 VSA -- the official FastVideo segment/tile mask semantics.
 
         Args:
@@ -950,6 +982,12 @@ class VSASparseAttention(nn.Module):
             x:       [B, L, C] attention-input hidden states (drives the compress gate).
             segments: packed layout -- dense row counts (``int``) and video region
                 ``(t, h, w)`` token grid, in packed order. Built by the model forward.
+            gate: [B, L, H, D] pre-computed ``to_gate_compress`` output; when given, ``x`` is
+                ignored. ``to_gate_compress`` is a ``Linear(dim, num_heads*head_dim)`` and thus
+                mixes heads, so under Ulysses head-parallelism the caller must evaluate it on
+                its own sequence shard (where the FULL hidden dim is still present) and
+                all-to-all the result -- passing the head shard in here keeps this module
+                head-local. ``H`` is then the local head count.
         Returns:
             [B, L, H, D] attention output (before the output projection).
 
@@ -984,7 +1022,8 @@ class VSASparseAttention(nn.Module):
         qb = _tile(q, self.dtype, 'q')
         kb = _tile(k, self.dtype, 'k')
         vb = _tile(v, self.dtype, 'v')
-        gate = self.to_gate_compress(x).view(B, L, H, D)
+        gate = self.to_gate_compress(x) if gate is None else gate
+        gate = gate.view(B, L, H, D)
         gb = _tile(gate, gate.dtype, 'g')
 
         pq = _h3_pool_tiles(qb, vbs, tile_elems, n_tiles)     # [B,H,n_tiles,D] fp32
@@ -1016,12 +1055,15 @@ class VSASparseAttention(nn.Module):
             return out, self.topk
         return out
 
-    def forward(self, q, k, v, x, return_sparsity=False, segments=None):
+    def forward(self, q, k, v, x, return_sparsity=False, segments=None, gate=None):
         R'''
         Args:
             q, k, v: [B, L, H, D] POST-RoPE queries/keys/values.
             x:       [B, L, C] attention-input hidden states (drives the compress gate).
             return_sparsity: also return the sparse-branch keep ratio.
+            gate: [B, L, H, D] pre-computed ``to_gate_compress`` output; see ``forward_h3``
+                (required under Ulysses head-parallelism, where ``x`` is unavailable in the
+                head-shard layout because the gate Linear mixes heads).
         Returns:
             [B, L, H, D] attention output (before the output projection).
         '''
@@ -1030,7 +1072,8 @@ class VSASparseAttention(nn.Module):
             # kernel (forward + backward), so training and inference share one selection
             # semantics. Falls back to the flat 1-D block map only when the caller cannot
             # supply a clean segment layout (segments is None).
-            return self.forward_h3(q, k, v, x, segments, return_sparsity=return_sparsity)
+            return self.forward_h3(q, k, v, x, segments, return_sparsity=return_sparsity,
+                                   gate=gate)
         B, L, H, D = q.shape
         qb = q.transpose(1, 2).contiguous()                   # [B,H,L,D]
         kb = k.transpose(1, 2).contiguous()
@@ -1041,7 +1084,8 @@ class VSASparseAttention(nn.Module):
             # Shorter than one block: exact dense attention (sparsity gives no benefit).
             o_s = F.scaled_dot_product_attention(
                 qb.float(), kb.float(), vb.float()).to(qb.dtype)
-        gate = self.to_gate_compress(x).view(B, L, H, D).transpose(1, 2)  # [B,H,L,D]
+        gate = self.to_gate_compress(x) if gate is None else gate
+        gate = gate.view(B, L, H, D).transpose(1, 2)          # [B,H,L,D]
         o_c = self.compressed(qb, kb, vb)
         out = o_s + gate.to(o_s.dtype) * o_c.to(o_s.dtype)
         out = out.transpose(1, 2)                             # [B,L,H,D]

videox_fun/models/fastvideo_minimax_h3_transformer3d.py

@@ -240,16 +240,23 @@ class FastVideoMiniMaxH3Transformer3DModel(MiniMaxH3Transformer3DModel):
         return tuple(segments)
 
     def enable_multi_gpus_inference(self,):
-        """VSA has no sequence-parallel path in this repo.
+        """Enable sequence-parallel inference for the VSA sparse student.
 
-        The in-repo VSA kernel only implements the single-GPU coarse-to-fine block-sparse forward;
-        there is no ulysses head-parallel or ring-attention binding for it (unlike the base
-        ``MiniMaxH3MultiGPUsAttnProcessor``). Fail loudly rather than silently falling back to a
-        DENSE multi-GPU processor (which would bypass VSA and drop the gate params).
+        Reuses the parent's packed-sequence SP assembly -- it sizes ``sp_world_size`` /
+        ``all_gather``, pads the sequence to a group multiple, hands every rank a contiguous row
+        slice and checks ``num_heads`` is divisible by the group -- then re-binds the VSA
+        sequence-parallel processor onto every block that carries an ``attn_op``. The parent
+        binding (``MiniMaxH3MultiGPUsAttnProcessor``) is a DENSE core; leaving it on a sparse
+        block would silently bypass VSA and drop the ``to_gate_compress`` params, so the sparse
+        blocks are swapped to ``MiniMaxH3MultiGPUsVSAAttnProcessor``, which turns the row split
+        into a head split around the VSA op. A dense student (``topk >= 1.0``) has no ``attn_op``
+        and keeps the parent's dense processor untouched.
         """
-        raise NotImplementedError(
-            "FastVideoMiniMaxH3Transformer3DModel (VSA) supports single-GPU inference only; the "
-            "VSA kernel has no ulysses/ring sequence-parallel path. Set ulysses_degree=ring_degree=1.")
+        super().enable_multi_gpus_inference()
+        from ..dist import MiniMaxH3MultiGPUsVSAAttnProcessor
+        for block in self.transformer_blocks:
+            if hasattr(block.attn, "attn_op"):
+                block.attn.set_processor(MiniMaxH3MultiGPUsVSAAttnProcessor())
 
     def sparse_trainable_params(self):
         """Names of the extra VSA-branch parameters, for logging / freezing."""

videox_fun/models/fastvideo_wan_transformer3d.py

@@ -67,26 +67,45 @@ class FastVideoWanSelfAttention(WanSelfAttention):
 
         q, k = rope_apply_qk(q, k, grid_sizes, freqs)
 
+        out = self._vsa_dispatch(q.to(dtype), k.to(dtype), v.to(dtype), x.to(dtype),
+                                 seq_lens, grid_sizes)
+        x = out.to(dtype).flatten(2)
+        return self.o(x)
+
+    def _vsa_dispatch(self, q, k, v, x, seq_lens, grid_sizes, gate=None):
+        """Run the VSA kernel over POST-RoPE ``q/k/v`` ([B, L, H, D]) -> [B, L, H, D].
+
+        ``x`` drives ``attn_op.to_gate_compress`` unless a pre-computed ``gate`` ([B, L, H, D])
+        is handed in: that Linear mixes heads, so the sequence-parallel path
+        (``usp_attn_fastvideo_forward``) evaluates it on its own sequence shard and passes the
+        head-sharded result, with ``x`` left as None.
+
+        The segment-aware tile geometry needs the sequence length to equal the true video token
+        count. Under sequence parallelism the gathered length is padded to a multiple of the SP
+        degree and can exceed it, so that case takes the per-sample branch below, which attends
+        over the valid grid tokens only and pads the output back to ``L``.
+        """
+        s = q.shape[1]
         g0 = grid_sizes[0]
         if (bool((grid_sizes == g0).all()) and s == int(g0.prod())
                 and bool((seq_lens == s).all())):
             segments = ((int(g0[0]), int(g0[1]), int(g0[2])),)
-            out = self.attn_op(q.to(dtype), k.to(dtype), v.to(dtype), x.to(dtype),
-                               segments=segments)
-        else:
-            outputs = []
-            for i, (grid, length) in enumerate(zip(grid_sizes.tolist(), seq_lens.tolist())):
-                grid = tuple(int(axis) for axis in grid)
-                length = int(length)
-                if min(grid) <= 0 or length != grid[0] * grid[1] * grid[2] or not 0 < length <= s:
-                    raise ValueError("Wan VSA 的 seq_lens 必须等于有效网格 token 数且不超过序列长度。")
-                oi = self.attn_op(q[i:i + 1, :length].to(dtype), k[i:i + 1, :length].to(dtype),
-                                  v[i:i + 1, :length].to(dtype), x[i:i + 1, :length].to(dtype),
-                                  segments=(grid,))
-                outputs.append(torch.nn.functional.pad(oi, (0, 0, 0, 0, 0, s - length)))
-            out = torch.cat(outputs, dim=0)
-        x = out.to(dtype).flatten(2)
-        return self.o(x)
+            return self.attn_op(q, k, v, x, segments=segments, gate=gate)
+
+        outputs = []
+        for i, (grid, length) in enumerate(zip(grid_sizes.tolist(), seq_lens.tolist())):
+            grid = tuple(int(axis) for axis in grid)
+            length = int(length)
+            if min(grid) <= 0 or length != grid[0] * grid[1] * grid[2] or not 0 < length <= s:
+                raise ValueError("Wan VSA 的 seq_lens 必须等于有效网格 token 数且不超过序列长度。")
+            gate_i = None if gate is None else gate[i:i + 1, :length]
+            # x is only consumed by to_gate_compress, so it is dropped once gate is supplied.
+            x_i = None if gate is not None else x[i:i + 1, :length]
+            oi = self.attn_op(q[i:i + 1, :length], k[i:i + 1, :length],
+                              v[i:i + 1, :length], x_i,
+                              segments=(grid,), gate=gate_i)
+            outputs.append(torch.nn.functional.pad(oi, (0, 0, 0, 0, 0, s - length)))
+        return torch.cat(outputs, dim=0)
 
 
 class FastVideoWanTransformer3DModel(WanTransformer3DModel):
@@ -187,17 +206,45 @@ class FastVideoWanTransformer3DModel(WanTransformer3DModel):
             self._apply_sparse_self_attention(topk)
 
     def enable_multi_gpus_inference(self,):
-        """VSA has no sequence-parallel path in this repo.
+        """Bind the VSA-aware USP forward instead of the parent's dense one.
 
-        The in-repo VSA kernel only implements the single-GPU coarse-to-fine block-sparse
-        forward; there is no ulysses head-parallel or ring-attention binding for it (unlike
-        the RoLa/SLA ``usp_attn_*`` forwards). The 1.3B student fits comfortably on one GPU,
-        so run inference single-device. Fail loudly rather than silently falling back to the
-        parent's DENSE ``usp_attn_forward`` (which would bypass VSA and drop the gate params).
+        ``WanTransformer3DModel.enable_multi_gpus_inference`` monkey-patches every
+        ``blocks[i].self_attn.forward`` to ``usp_attn_forward`` (dense full-softmax via
+        ``xFuserLongContextAttention``), which never reaches ``FastVideoWanSelfAttention.attn_op``
+        and would silently drop the VSA kernel plus the trained ``to_gate_compress`` gate.
+        Rebind to ``usp_attn_fastvideo_forward``: Ulysses head-parallel, with the cross-head
+        compress gate evaluated on each rank's sequence shard (where the full hidden dim still
+        lives) and all-to-all'd next to q/k/v.
+
+        Constraint: the VSA kernel has no ring-attention path (no partial-output/logsumexp
+        interface to rescale across KV chunks), so this is pure Ulysses -- ``ring_degree`` must
+        be 1 and ``num_heads`` must divide the SP degree.
         """
-        raise NotImplementedError(
-            "FastVideoWanTransformer3DModel (VSA) supports single-GPU inference only; the VSA "
-            "kernel has no ulysses/ring sequence-parallel path. Set ulysses_degree=ring_degree=1.")
+        import types
+
+        from ..dist import (get_sequence_parallel_rank,
+                            get_sequence_parallel_world_size, get_sp_group,
+                            usp_attn_fastvideo_forward)
+        self.sp_world_size = get_sequence_parallel_world_size()
+        self.sp_world_rank = get_sequence_parallel_rank()
+        self.all_gather = get_sp_group().all_gather
+
+        # Ulysses head-split needs num_heads divisible by the SP degree.
+        n_heads = self.blocks[0].self_attn.num_heads
+        if self.sp_world_size > 1 and n_heads % self.sp_world_size != 0:
+            raise ValueError(
+                f"VSA multi-GPU (ulysses) needs num_heads ({n_heads}) divisible by the "
+                f"sequence-parallel degree ({self.sp_world_size}); set ulysses_degree "
+                f"accordingly and keep ring_degree=1 (the VSA kernel has no ring path).")
+
+        for block in self.blocks:
+            block.self_attn.forward = types.MethodType(
+                usp_attn_fastvideo_forward, block.self_attn)
+
+        if hasattr(self, 'vace_blocks'):
+            for block in self.vace_blocks:
+                block.self_attn.forward = types.MethodType(
+                    usp_attn_fastvideo_forward, block.self_attn)
 
     def sparse_trainable_params(self):
         """Names of the extra VSA-branch parameters, for logging / freezing."""

videox_fun/models/minimax_h3_conversion.py

@@ -111,10 +111,13 @@ class MiniMaxH3MixedPrecisionLoaderMixin:
             kwargs["device_map"] = device_map
         if max_memory is not None:
             kwargs["max_memory"] = max_memory
-        model = super().from_pretrained(pretrained_model_name_or_path, torch_dtype=torch_dtype, **kwargs)
+        result = super().from_pretrained(pretrained_model_name_or_path, torch_dtype=torch_dtype, **kwargs)
+        # ModelMixin returns (model, loading_info) when output_loading_info=True.
+        # Restore precision on the model while preserving that return contract.
+        model = result[0] if isinstance(result, tuple) else result
         if torch_dtype not in (None, torch.float32) and cls._keep_in_fp32_modules:
             restore_fp32_modules(model, pretrained_model_name_or_path, kwargs.get("subfolder", None))
-        return model
+        return result
 
     @classmethod
     def _from_pretrained_with_missing_keys(cls, pretrained_model_name_or_path, torch_dtype, **kwargs):

videox_fun/models/minimax_h3_transformer3d.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Tuple, Union
 
@@ -264,7 +264,11 @@ class MiniMaxH3Attention(nn.Module):
         rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         attention_mask: Optional[torch.Tensor] = None,
         valid_length: Optional[int] = None,
+        vsa_segments=None,
     ) -> torch.Tensor:
+        if vsa_segments is not None:
+            return self.processor(self, hidden_states, rotary_emb, attention_mask, valid_length,
+                                  vsa_segments=vsa_segments)
         return self.processor(self, hidden_states, rotary_emb, attention_mask, valid_length)
 
 
@@ -374,6 +378,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
         rotary_emb: Tuple[torch.Tensor, torch.Tensor],
         attention_mask: Optional[torch.Tensor] = None,
         valid_length: Optional[int] = None,
+        vsa_segments=None,
     ) -> torch.Tensor:
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(temb)
 
@@ -382,7 +387,8 @@ class MiniMaxH3TransformerBlock(nn.Module):
         norm_hidden_states = norm_hidden_states * (
             1.0 + _index_select_rows(scale_msa, adaln_indices)
         ) + _index_select_rows(shift_msa, adaln_indices)
-        attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask, valid_length)
+        attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask, valid_length,
+                                vsa_segments=vsa_segments)
         hidden_states = residual + _index_select_rows(gate_msa, adaln_indices) * attn_output
 
         residual = hidden_states
@@ -546,6 +552,9 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin
 
         self.gradient_checkpointing = False
         self.gradient_checkpointing_save_on_cpu = False
+        # Optional non-reentrant checkpoint `context_fn` for the main blocks; the streaming KV cache installs one so
+        # backward recomputes see the cache view of the original forward.
+        self.gradient_checkpointing_context_fn = None
         self.sp_world_size = 1
         self.sp_world_rank = 0
         self.all_gather = None
@@ -747,10 +756,14 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin
             adaln_indices = adaln_indices[rows]
             rotary_emb = (rotary_emb[0][rows], rotary_emb[1][rows])
 
+        vsa_segments = attention_kwargs.get("vsa_segments") if attention_kwargs else None
+        block_kwargs = {"vsa_segments": vsa_segments} if vsa_segments is not None else {}
         for block in self.transformer_blocks:
             with torch.autograd.graph.save_on_cpu() if self.gradient_checkpointing_save_on_cpu else nullcontext():
                 if torch.is_grad_enabled() and self.gradient_checkpointing:
                     ckpt_kwargs = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    if self.gradient_checkpointing_context_fn is not None:
+                        ckpt_kwargs["context_fn"] = self.gradient_checkpointing_context_fn
                     hidden_states = torch.utils.checkpoint.checkpoint(
                         block,
                         hidden_states,
@@ -759,10 +772,12 @@ class MiniMaxH3Transformer3DModel(MiniMaxH3MixedPrecisionLoaderMixin, ModelMixin
                         rotary_emb,
                         attention_mask,
                         valid_length,
+                        **block_kwargs,
                         **ckpt_kwargs,
                     )
                 else:
-                    hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb, attention_mask, valid_length)
+                    hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb, attention_mask,
+                                          valid_length, **block_kwargs)
 
         if self.sp_world_size > 1:
             hidden_states = self.all_gather(hidden_states.contiguous(), dim=1)[:, :sequence_length]
@@ -822,11 +837,20 @@ class MiniMaxH3StreamingKVCache:
     Args:
         num_layers (`int`): The number of main transformer blocks (50 for the released H3).
         dtype (`torch.dtype`, defaults to `torch.bfloat16`): The dtype the clean K/V is stored at.
+        offload_to_cpu (`bool`, defaults to `False`): Whether the clean K/V is parked on the CPU.
+        commit_attend_history (`bool`, defaults to `False`): Whether a clean commit pass runs the
+            live attention pattern (media rows attend `[text K/V; persistent history; current chunk
+            K/V]`, text rows self-attend) instead of the legacy standalone pass over
+            `[text | current chunk]`. Enable it when the history must condition the committed K/V,
+            as in the Self-Forcing clean refresh pass.
     """
 
-    def __init__(self, num_layers: int, dtype: torch.dtype = torch.bfloat16):
+    def __init__(self, num_layers: int, dtype: torch.dtype = torch.bfloat16, offload_to_cpu: bool = False,
+                 commit_attend_history: bool = False):
         self.layer_names = tuple(f"transformer_blocks.{index}.attn" for index in range(num_layers))
         self.dtype = dtype
+        self.offload_to_cpu = offload_to_cpu
+        self.commit_attend_history = commit_attend_history
         self._history: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
         self._staged: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
         self._commit_active = False
@@ -838,6 +862,10 @@ class MiniMaxH3StreamingKVCache:
         # Per-forward attention mode: None ("plain"), "live" or "commit".
         self.mode: Optional[str] = None
         self.token_tags: Optional[torch.Tensor] = None
+        # Activation-checkpoint replay: while set, `history_pair` exposes only this many leading history rows (the
+        # history a checkpointed forward saw); `_epoch` changes whenever rows are dropped, invalidating pending replays.
+        self._history_view_rows: Optional[int] = None
+        self._epoch = 0
 
     # ------------------------------------------------------------------
     # Per-forward setup consumed by the streaming processor.
@@ -902,11 +930,51 @@ class MiniMaxH3StreamingKVCache:
             self._staged_row_count = int(key.shape[0])
         elif int(key.shape[0]) != self._staged_row_count:
             raise RuntimeError(f"{layer_name}: staged row count differs across layers")
-        self._staged[layer_name] = (key, value)
+        device = torch.device("cpu") if self.offload_to_cpu else key.device
+        self._staged[layer_name] = (
+            key.detach().to(device=device, dtype=self.dtype).contiguous(),
+            value.detach().to(device=device, dtype=self.dtype).contiguous(),
+        )
 
     def history_pair(self, layer_name: str) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
         r"""Return the persistent `(key, value)` of one layer, or `None` when empty."""
-        return self._history.get(layer_name)
+        pair = self._history.get(layer_name)
+        if pair is None or self._history_view_rows is None:
+            return pair
+        if self._history_view_rows == 0:
+            return None
+        return pair[0][: self._history_view_rows], pair[1][: self._history_view_rows]
+
+    # ------------------------------------------------------------------
+    # Activation checkpointing.
+    # ------------------------------------------------------------------
+    def checkpoint_context(self):
+        r"""
+        `context_fn` for non-reentrant `torch.utils.checkpoint` over a main block.
+
+        Called when the checkpointed forward runs, it snapshots the cache view that forward sees — the mode, the
+        token tags and the committed history length — and returns `(forward_context, recompute_context)`; the
+        recompute context restores that view while backward replays the block, after `end_forward()` cleared the
+        mode and later commits appended more history. The history is append-only (`commit` concatenates), so the
+        snapshot is a row count, not a tensor copy: the replay reads a prefix view of the current history.
+        """
+        history_rows = self.history_tokens if self._history_view_rows is None else self._history_view_rows
+        snapshot = (self.mode, self.token_tags, history_rows, self._epoch)
+        return nullcontext(), self._replay_view(*snapshot)
+
+    @contextmanager
+    def _replay_view(self, mode, token_tags, history_rows, epoch):
+        if epoch != self._epoch:
+            raise RuntimeError(
+                "the streaming KV history was cleared or trimmed between a checkpointed forward and its backward "
+                "recompute; clear the cache only after backward"
+            )
+        saved = (self.mode, self.token_tags, self._history_view_rows)
+        self.mode, self.token_tags, self._history_view_rows = mode, token_tags, history_rows
+        try:
+            yield
+        finally:
+            self.mode, self.token_tags, self._history_view_rows = saved
 
     # ------------------------------------------------------------------
     # Transaction control.
@@ -931,6 +999,7 @@ class MiniMaxH3StreamingKVCache:
                     torch.cat((previous[1], current[1]), dim=0),
                 )
             self._history[layer_name] = combined
+            del current, previous, combined
         self._commit_token_counts.append(len(self._commit_tags))
         self._commit_token_tags.append(self._commit_tags)
         self._commit_active = False
@@ -1005,6 +1074,7 @@ class MiniMaxH3StreamingKVCache:
     def clear(self) -> None:
         r"""Clear persistent and staged KV and reset the logical chunk count."""
         self._history.clear()
+        self._epoch += 1
         self._commit_token_counts.clear()
         self._commit_token_tags.clear()
         self._staged.clear()
@@ -1014,6 +1084,7 @@ class MiniMaxH3StreamingKVCache:
         self.end_forward()
 
     def _retain_commit_rows(self, selection) -> None:
+        self._epoch += 1
         offsets = [0]
         for count in self._commit_token_counts:
             offsets.append(offsets[-1] + count)
@@ -1068,9 +1139,11 @@ class MiniMaxH3StreamingKVCacheAttnProcessor:
       [`MiniMaxH3AttnProcessor`];
     * *live*: text rows self-attend, while media rows attend to
       `[text K/V; persistent clean history; current chunk K/V]`;
-    * *commit*: a plain pass whose output is discarded, staging the media rows' K/V of
-      every layer into the cache (the K/V of a row depends only on that row's input, so
-      the commit pass's attention output is irrelevant).
+    * *commit*: the pass whose output is discarded, staging the media rows' K/V of every
+      layer into the cache. With `commit_attend_history` off (legacy) it is a plain pass over
+      `[text | current chunk]`; with it on, the commit pass runs the live attention pattern so
+      the staged K/V is conditioned on the persistent history, matching the Self-Forcing clean
+      refresh pass.
 
     Live and commit passes must be padless (`attention_mask is None`); the streaming
     packed sequences are single documents by construction.
@@ -1107,19 +1180,19 @@ class MiniMaxH3StreamingKVCacheAttnProcessor:
             query = apply_minimax_h3_rotary_emb(query, *rotary_emb)
             key = apply_minimax_h3_rotary_emb(key, *rotary_emb)
 
-        if cache.mode == "commit":
-            # The committed rows' K/V depend only on their own inputs; the attention output
-            # of a commit pass is discarded by the caller.
+        if cache.mode == "commit" and not cache.commit_attend_history:
+            # Legacy commit: a plain pass over [text | current chunk]; the attention output is
+            # discarded by the caller and only the staged media K/V rows matter.
             output = attention(query, key, value, causal=False)
             rows = cache.commit_row_indices.to(key.device)
             staged_key = key.index_select(1, rows).squeeze(0)
             staged_value = value.index_select(1, rows).squeeze(0)
             cache.stage(
                 self.layer_name,
-                staged_key.detach().to(cache.dtype).contiguous(),
-                staged_value.detach().to(cache.dtype).contiguous(),
+                staged_key,
+                staged_value,
             )
-        else:  # live
+        else:  # live, or a commit with commit_attend_history: media rows attend the history.
             tags = cache.token_tags
             text_rows = tags == MINIMAX_H3_STREAMING_TEXT_TAG
             media_rows = ~text_rows
@@ -1153,6 +1226,16 @@ class MiniMaxH3StreamingKVCacheAttnProcessor:
                     values.transpose(1, 2),
                 ).transpose(1, 2)
 
+        if cache.mode == "commit" and cache.commit_attend_history:
+            # The history-conditioned commit pass output is still discarded; only the staged K/V
+            # rows matter, now computed with the history in context.
+            rows = cache.commit_row_indices.to(key.device)
+            cache.stage(
+                self.layer_name,
+                key.index_select(1, rows).squeeze(0),
+                value.index_select(1, rows).squeeze(0),
+            )
+
         output = output.flatten(2, 3).type_as(query)
         output = attn.to_out[0](output)
         output = attn.to_out[1](output)
@@ -1201,21 +1284,26 @@ def install_minimax_h3_streaming_kv_cache(
     transformer: "MiniMaxH3Transformer3DModel",
     *,
     dtype: torch.dtype = torch.bfloat16,
+    offload_to_cpu: bool = False,
+    commit_attend_history: bool = False,
 ) -> MiniMaxH3StreamingKVCache:
     r"""
     Swap every main block's attention processor for the streaming one and return the shared
     cache.
 
     The token refiner keeps its standard processor: the streaming contract only concerns the
-    main block stack over the media sequence.
+    main block stack over the media sequence. `commit_attend_history` is forwarded to
+    [`MiniMaxH3StreamingKVCache`].
     """
     cache = MiniMaxH3StreamingKVCache(
-        num_layers=len(transformer.transformer_blocks), dtype=dtype
+        num_layers=len(transformer.transformer_blocks), dtype=dtype, offload_to_cpu=offload_to_cpu,
+        commit_attend_history=commit_attend_history,
     )
     for index, block in enumerate(transformer.transformer_blocks):
         block.attn.set_processor(
             MiniMaxH3StreamingKVCacheAttnProcessor(cache, cache.layer_names[index])
         )
+    transformer.gradient_checkpointing_context_fn = cache.checkpoint_context
     return cache
 
 
@@ -1223,3 +1311,4 @@ def uninstall_minimax_h3_streaming_kv_cache(transformer: "MiniMaxH3Transformer3D
     r"""Restore the standard attention processor on every main block."""
     for block in transformer.transformer_blocks:
         block.attn.set_processor(MiniMaxH3AttnProcessor())
+    transformer.gradient_checkpointing_context_fn = None

videox_fun/models/turbowan_transformer3d.py

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

videox_fun/models/wan_tae.py

@@ -26,6 +26,29 @@ DETERMINISTIC_LOGVAR = -30.0
 TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))
 
 
+def _load_raw_state_dict(path):
+    """Load a flat tensor state_dict from a .safetensors/.pth/.bin checkpoint.
+
+    Mirrors the normalization done in AutoencoderTinyWan.from_pretrained so that a
+    companion encoder checkpoint can be borrowed the same way.
+    """
+    path = os.fspath(path)
+    if path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+
+        sd = load_file(path)
+    else:
+        sd = torch.load(path, map_location="cpu", weights_only=True)
+    if isinstance(sd, dict):
+        if isinstance(sd.get("state_dict"), dict):
+            sd = sd["state_dict"]
+        elif isinstance(sd.get("decoder"), dict):
+            sd = sd["decoder"]
+    if isinstance(sd, dict) and all(k.startswith("model.") for k in sd):
+        sd = {k[len("model.") :]: v for k, v in sd.items()}
+    return sd
+
+
 def conv(n_in, n_out, **kwargs):
     return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
 
@@ -566,6 +589,10 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
         if spatial_compression_ratio is None:
             spatial_compression_ratio = 8 * patch_size
         self.spatial_compression_ratio = spatial_compression_ratio
+        # Persist the derived value into the frozen config as well: some pipelines
+        # read `vae.config.spatial_compression_ratio` (not the attribute), and the
+        # register_to_config decorator would otherwise keep the `None` default.
+        self.register_to_config(spatial_compression_ratio=spatial_compression_ratio)
 
         self.model = TAEWan(
             encoder_time_downscale=tuple(encoder_time_downscale),
@@ -696,6 +723,33 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
         if all(k.startswith("model.") for k in state_dict):
             state_dict = {k[len("model.") :]: v for k, v in state_dict.items()}
 
+        # Borrow an encoder when the checkpoint is decoder-only. Some taew2_2_super
+        # releases ship only the Super decoder; without an encoder, encode() would
+        # raise. We merge encoder.* weights from a companion base checkpoint so the
+        # model becomes a full encoder(base)+decoder(super) TAE.
+        # Priority: explicit `encoder_path` (via additional_kwargs / vae_kwargs or
+        # kwargs) > auto sibling found by stripping "_super" from the filename.
+        encoder_path = None
+        if isinstance(additional_kwargs, dict):
+            encoder_path = additional_kwargs.get("encoder_path")
+        if encoder_path is None:
+            encoder_path = kwargs.get("encoder_path")
+        if not any(k.startswith("encoder.") for k in state_dict):
+            if encoder_path is None:
+                root, ext = os.path.splitext(pretrained_model_path)
+                sibling = root.replace("_super", "") + ext
+                if sibling != pretrained_model_path and os.path.exists(sibling):
+                    encoder_path = sibling
+            if encoder_path is not None and os.path.exists(encoder_path):
+                enc_sd = {k: v for k, v in _load_raw_state_dict(encoder_path).items() if k.startswith("encoder.")}
+                if enc_sd:
+                    state_dict = {**enc_sd, **state_dict}
+                    print(f"[TAE] decoder-only 权重：已从 {encoder_path} 借用 {len(enc_sd)} 个 encoder 键以支持 encode。")
+            if not any(k.startswith("encoder.") for k in state_dict):
+                print("[TAE][warn] 未找到 encoder 权重，将以 decoder_only 加载，encode() 不可用。"
+                      "可通过 vae_kwargs.encoder_path=<base taew2_2 权重路径> 指定，"
+                      "或把 base 权重放在与 super 同名（去掉 _super）的路径下。")
+
         if "decoder.0.weight" in state_dict and "decoder.1.weight" not in state_dict:
             arch_variant, first_key, block_key = "super", "decoder.0.weight", "decoder.1.conv.0.weight"
             first_channels, first_kernel, block_shape = 512, (1, 1), (1024, 1, 7, 7)

videox_fun/pipeline/pipeline_wan.py

@@ -1,5 +1,6 @@
 import inspect
 import math
+import time
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -21,6 +22,17 @@ from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
+def stochastic_sampling_timesteps(num_inference_steps, shift, device, num_timesteps=1000):
+    """Official FlashHead timestep schedule with shift transform."""
+    if num_inference_steps == 4:
+        timesteps = [1000, 750, 500, 250]
+    else:
+        timesteps = np.linspace(num_timesteps, 1, num_inference_steps, dtype=np.float32).tolist()
+    timesteps = torch.tensor(timesteps + [0.0], dtype=torch.float32, device=device)
+    t = timesteps / num_timesteps
+    return shift * t / (1 + (shift - 1) * t) * num_timesteps
+
+
 EXAMPLE_DOC_STRING = """
     Examples:
         ```python
@@ -409,6 +421,7 @@ class WanPipeline(DiffusionPipeline):
         max_sequence_length: int = 512,
         comfyui_progressbar: bool = False,
         shift: int = 5,
+        stochastic_sampling: bool = False,
     ) -> Union[WanPipelineOutput, Tuple]:
         """
         Function invoked when calling the pipeline for generation.
@@ -471,7 +484,9 @@ class WanPipeline(DiffusionPipeline):
             in_prompt_embeds = prompt_embeds
 
         # 4. Prepare timesteps
-        if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+        if stochastic_sampling:
+            timesteps = stochastic_sampling_timesteps(num_inference_steps, shift, device)
+        elif isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
             timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, mu=1)
         elif isinstance(self.scheduler, FlowUniPCMultistepScheduler):
             self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
@@ -513,8 +528,9 @@ class WanPipeline(DiffusionPipeline):
         # 7. Denoising loop
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
         self.transformer.num_inference_steps = num_inference_steps
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
+        denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps
+        with self.progress_bar(total=len(denoise_timesteps)) as progress_bar:
+            for i, t in enumerate(denoise_timesteps):
                 self.transformer.current_steps = i
 
                 if self.interrupt:
@@ -528,6 +544,8 @@ class WanPipeline(DiffusionPipeline):
                 timestep = t.expand(latent_model_input.shape[0])
                 
                 # predict noise model_output
+                torch.cuda.synchronize()
+                forward_start_time = time.time()
                 with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=device):
                     noise_pred = self.transformer(
                         x=latent_model_input,
@@ -535,6 +553,9 @@ class WanPipeline(DiffusionPipeline):
                         t=timestep,
                         seq_len=seq_len,
                     )
+                torch.cuda.synchronize()
+                forward_time = time.time() - forward_start_time
+                print(f"Step {i}/{len(timesteps)} - Transformer forward time: {forward_time:.4f}s")
 
                 # perform guidance
                 if do_classifier_free_guidance:
@@ -542,7 +563,15 @@ class WanPipeline(DiffusionPipeline):
                     noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
 
                 # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                if stochastic_sampling:
+                    t_i = (timesteps[i] / 1000).to(weight_dtype)
+                    t_i_1 = (timesteps[i + 1] / 1000).to(weight_dtype)
+                    x_0 = latents - noise_pred * t_i
+                    latents = (1 - t_i_1) * x_0 + t_i_1 * torch.randn(
+                        latents.shape, dtype=latents.dtype, device=device, generator=generator
+                    )
+                else:
+                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
 
                 if callback_on_step_end is not None:
                     callback_kwargs = {}
@@ -554,7 +583,7 @@ class WanPipeline(DiffusionPipeline):
                     prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
                     negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
 
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                if i == len(denoise_timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                     progress_bar.update()
                 if comfyui_progressbar:
                     pbar.update(1)

videox_fun/pipeline/pipeline_wan_self_forcing.py

@@ -23,15 +23,19 @@ from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-def stochastic_sampling_timesteps(num_inference_steps, shift, device, num_timesteps=1000):
+def stochastic_sampling_timesteps(num_inference_steps, shift, device, num_timesteps=1000, timestep_list=None):
     """Official FlashHead timestep schedule with shift transform.
 
     Hardcoded schedules match the CF reference for distilled few-step inference:
       - 4-step: [1000, 750, 500, 250] — Stage 2 CCD (`causal_cd_framewise.yaml`)
       - 2-step: [1000, 500]           — Stage 3 DMD (`causal_forcing_dmd_framewise_2step.yaml`)
-    Anything else falls back to linspace.
+    Anything else falls back to linspace. An explicit `timestep_list` overrides
+    all of the above (per-block schedules such as the official One-Forcing
+    first-frame list).
     """
-    if num_inference_steps == 4:
+    if timestep_list is not None:
+        timesteps = list(timestep_list)
+    elif num_inference_steps == 4:
         timesteps = [1000, 750, 500, 250]
     elif num_inference_steps == 2:
         timesteps = [1000, 500]
@@ -589,6 +593,7 @@ class WanSelfForcingPipeline(DiffusionPipeline):
         independent_first_frame: bool = True,
         context_noise: int = 0,
         stochastic_sampling: bool = True,
+        first_frame_denoising_step_list: Optional[List[int]] = None,
         streaming: bool = False,
         decode_callback: Optional[Callable[[torch.Tensor, int], None]] = None,
         forcing_kv_enable: Optional[bool] = None,
@@ -611,6 +616,12 @@ class WanSelfForcingPipeline(DiffusionPipeline):
             num_frame_per_block: Number of frames to generate per block.
             independent_first_frame: Whether to generate the first frame independently (T2V mode).
             context_noise: Context noise level for KV cache update (matches training config).
+            first_frame_denoising_step_list: Optional explicit raw timestep list
+                (e.g. [1000, 750, 500, 250]) used only for the block that starts at
+                frame 0, mirroring the official One-Forcing
+                `first_frame_denoising_step_list`: that block is the only one with
+                an empty KV context and benefits from a multi-step schedule while
+                all other blocks keep the global one. None disables the override.
             streaming: If True, decode each causal block into pixels right after it is
                 generated (instead of accumulating all latents and decoding once at the end).
                 The VAE's causal temporal cache is threaded across blocks so the result is
@@ -696,6 +707,7 @@ class WanSelfForcingPipeline(DiffusionPipeline):
             in_prompt_embeds = prompt_embeds
 
         # 4. Prepare timesteps
+        custom_timesteps = timesteps
         if stochastic_sampling:
             timesteps = stochastic_sampling_timesteps(num_inference_steps, shift, device)
         elif isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
@@ -868,16 +880,40 @@ class WanSelfForcingPipeline(DiffusionPipeline):
             end_idx = start_idx + current_num_frames
             noisy_input = noise[:, :, start_idx:end_idx]
 
-            if not stochastic_sampling and isinstance(self.scheduler, FlowUniPCMultistepScheduler):
-                self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
-                timesteps = self.scheduler.timesteps
-            elif not stochastic_sampling and isinstance(self.scheduler, FlowDPMSolverMultistepScheduler):
-                self.scheduler.set_timesteps(device=device, sigmas=sampling_sigmas)
-                timesteps = self.scheduler.timesteps
-            else:
+            # Official One-Forcing gives the block starting at frame 0 (the only
+            # block with an empty KV context) its own multi-step schedule; all
+            # other blocks keep the global one.
+            if current_start_frame == 0 and first_frame_denoising_step_list is not None:
+                if stochastic_sampling:
+                    block_timesteps = stochastic_sampling_timesteps(
+                        len(first_frame_denoising_step_list), shift, device,
+                        timestep_list=first_frame_denoising_step_list)
+                elif isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+                    self.scheduler.set_timesteps(
+                        timesteps=list(first_frame_denoising_step_list), device=device)
+                    block_timesteps = self.scheduler.timesteps
+                else:
+                    raise ValueError(
+                        "first_frame_denoising_step_list requires stochastic sampling or "
+                        f"FlowMatchEulerDiscreteScheduler, got {type(self.scheduler)}")
                 self.scheduler._step_index = None
-            
-            denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps
+            else:
+                if not stochastic_sampling and isinstance(self.scheduler, FlowUniPCMultistepScheduler):
+                    self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
+                    timesteps = self.scheduler.timesteps
+                elif not stochastic_sampling and isinstance(self.scheduler, FlowDPMSolverMultistepScheduler):
+                    self.scheduler.set_timesteps(device=device, sigmas=sampling_sigmas)
+                    timesteps = self.scheduler.timesteps
+                elif (not stochastic_sampling and first_frame_denoising_step_list is not None
+                        and isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler)):
+                    # Restore the global schedule overridden by the first-frame block
+                    timesteps, num_inference_steps = retrieve_timesteps(
+                        self.scheduler, num_inference_steps, device, custom_timesteps)
+                else:
+                    self.scheduler._step_index = None
+                block_timesteps = timesteps
+
+            denoise_timesteps = block_timesteps[:-1] if stochastic_sampling else block_timesteps
             with self.progress_bar(total=len(denoise_timesteps)) as progress_bar:
                 for step_idx, t in enumerate(denoise_timesteps):
                     # Per-frame timesteps for causal generation
@@ -947,8 +983,8 @@ class WanSelfForcingPipeline(DiffusionPipeline):
                     
                     # compute the previous noisy sample x_t -> x_t-1
                     if stochastic_sampling:
-                        t_i = (timesteps[step_idx] / 1000).to(weight_dtype)
-                        t_i_1 = (timesteps[step_idx + 1] / 1000).to(weight_dtype)
+                        t_i = (block_timesteps[step_idx] / 1000).to(weight_dtype)
+                        t_i_1 = (block_timesteps[step_idx + 1] / 1000).to(weight_dtype)
                         denoised_pred = noisy_input - flow_pred * t_i
                         noisy_input = (1 - t_i_1) * denoised_pred + t_i_1 * torch.randn(
                             denoised_pred.shape, dtype=denoised_pred.dtype, device=device, generator=generator

videox_fun/utils/fsdp_ema.py

@@ -109,6 +109,22 @@ class FSDPEMA:
             b_shadow.copy_(b_source.to(device=b_shadow.device, dtype=b_shadow.dtype))
         self.optimization_step += 1
 
+    @torch.no_grad()
+    def reinit_from(self, parameters):
+        r"""Reset the shadow to the current live weights.
+
+        Official One-Forcing builds its EMA at ``ema_start_step`` from the live weights of
+        that step; this reproduces that lazy creation for a shadow instantiated at step 0,
+        dropping whatever warm-start mass it had accumulated before the start step.
+        """
+        source_params, shadow_params = self._paired_params(parameters)
+        for p_shadow, p_source in zip(shadow_params, source_params):
+            if p_shadow.numel() == 0:
+                continue
+            p_shadow.copy_(p_source.detach().to(device=p_shadow.device, dtype=p_shadow.dtype))
+        for b_shadow, b_source in zip(self.module.buffers(), self._source_module.buffers()):
+            b_shadow.copy_(b_source.to(device=b_shadow.device, dtype=b_shadow.dtype))
+
     @torch.no_grad()
     def store(self, parameters):
         # Snapshots the live local shards so `copy_to` can be undone with

videox_fun/utils/fsdp_utils.py

@@ -91,7 +91,13 @@ def checkpoint_fsdp_forward(model, accumulated, *, context_fn=None, **kwargs):
                     else:
                         parameter.grad.add_(previous)
                 if parameter.grad is not None:
-                    accumulated[parameter] = parameter.grad.detach().clone()
+                    if previous is None:
+                        accumulated[parameter] = parameter.grad.detach().clone()
+                    else:
+                        # Keep one independent accumulator per shard across blocks.
+                        # Replacing it with a clone retains both full copies until
+                        # previous_grads leaves this backward call.
+                        previous.copy_(parameter.grad)
             result = [None] * len(keys)
             for index, tensor in zip(ctx.tensor_indices, tensors):
                 result[index] = tensor.grad