Update self forcing gan and causal forcing precision

2026-10-23 13:10:15 +08:00
parent 42b33f4998
commit e7afc604e1
30 changed files with 7265 additions and 213 deletions
@@ -0,0 +1,648 @@
+# FastVideo VSA Distill Training Guide
+
+This document provides a complete workflow for distilling Wan2.1 into FastWan2.1 with VSA (Video Sparse Attention), including environment configuration, data preparation, distributed training, and inference testing.
+
+> **Note**: FastVideo VSA combines a block-sparse video attention (a learned gate keeps ~20% of key blocks per query) with DMD few-step distillation, so the student runs CFG-free in ~3 steps while attending sparsely. The training recipe, data pipeline, and distributed backends are identical to the TurboDiffusion distill guide; only the sparse student model class differs.
+
+---
+
+## Table of Contents
+- [1. Environment Configuration](#1-environment-configuration)
+- [2. Data Preparation](#2-data-preparation)
+  - [2.1 Quick Test Dataset](#21-quick-test-dataset)
+  - [2.2 Dataset Structure](#22-dataset-structure)
+  - [2.3 metadata.json Format](#23-metadatajson-format)
+  - [2.4 Relative vs Absolute Path Usage](#24-relative-vs-absolute-path-usage)
+- [3. Distill Training](#3-distill-training)
+  - [3.1 Download Pretrained Model](#31-download-pretrained-model)
+  - [3.2 Quick Start (DeepSpeed-Zero-2)](#32-quick-start-deepspeed-zero-2)
+  - [3.3 Common Training Parameters](#33-common-training-parameters)
+  - [3.4 Training Validation](#34-training-validation)
+  - [3.5 Training with FSDP](#35-training-with-fsdp)
+  - [3.6 Other Backends](#36-other-backends)
+  - [3.7 Multi-Machine Distributed Training](#37-multi-machine-distributed-training)
+- [4. Inference Testing](#4-inference-testing)
+  - [4.1 Inference Parameters](#41-inference-parameters)
+  - [4.2 Text-to-Video (T2V) Inference](#42-text-to-video-t2v-inference)
+  - [4.3 Single-GPU Only](#43-single-gpu-only)
+- [5. Additional Resources](#5-additional-resources)
+
+---
+
+## 1. Environment Configuration
+
+**Method 1: Using requirements.txt**
+
+```bash
+pip install -r requirements.txt
+```
+
+**Method 2: Manual Dependency Installation**
+
+```bash
+pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
+pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
+pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
+pip install yunchang xfuser modelscope openpyxl
+pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
+pip install opencv-python-headless
+pip install deepspeed==0.17.0 numpy==1.26.4
+```
+
+> 💡 The VSA block-sparse attention runs on a Triton kernel; Triton ships with PyTorch (>=2.1) on Linux, so no extra install is needed.
+
+**Method 3: Using Docker**
+
+When using Docker, please ensure that the GPU driver and CUDA environment are correctly installed on your machine, then execute the following commands:
+
+```
+# pull image
+docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+
+# enter image
+docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+```
+
+---
+
+## 2. Data Preparation
+
+### 2.1 Quick Test Dataset
+
+We provide a test dataset containing several training samples.
+
+```bash
+# Download official demo dataset
+modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
+```
+
+### 2.2 Dataset Structure
+
+```
+📦 datasets/
+├── 📂 my_dataset/
+│   ├── 📂 train/
+│   │   ├── 📄 video001.mp4
+│   │   ├── 📄 video002.mp4
+│   │   └── 📄 ...
+│   └── 📄 metadata.json
+```
+
+### 2.3 metadata.json Format
+
+**Relative Path Format** (example):
+```json
+[
+  {
+    "file_path": "train/video001.mp4",
+    "text": "A beautiful sunset over the ocean, golden hour lighting",
+    "type": "video",
+    "width": 1024,
+    "height": 1024
+  },
+  {
+    "file_path": "train/video002.mp4",
+    "text": "A person walking through a forest, cinematic view",
+    "type": "video",
+    "width": 1328,
+    "height": 1328
+  }
+]
+```
+
+**Absolute Path Format**:
+```json
+[
+  {
+    "file_path": "/mnt/data/videos/sunset.mp4",
+    "text": "A beautiful sunset over the ocean",
+    "type": "video",
+    "width": 1024,
+    "height": 1024
+  }
+]
+```
+
+**Key Field Descriptions**:
+- `file_path`: Video path (relative or absolute path)
+- `text`: Video description (English prompt)
+- `type`: Data type, fixed as `"video"`
+- `width` / `height`: Video dimensions (**recommended** to provide for bucket training. If not provided, they will be automatically read during training, which may slow down training when data is stored on slow systems like OSS).
+  - You can use `scripts/process_json_add_width_and_height.py` to extract width and height fields for JSON files without them, supporting both images and videos.
+  - Usage: `python scripts/process_json_add_width_and_height.py --input_file datasets/X-Fun-Videos-Demo/metadata.json --output_file datasets/X-Fun-Videos-Demo/metadata_add_width_height.json`.
+
+### 2.4 Relative vs Absolute Path Usage
+
+**Relative Path**:
+
+If your data uses relative paths, configure in the training script:
+
+```bash
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+```
+
+**Absolute Path**:
+
+If your data uses absolute paths, configure in the training script:
+
+```bash
+export DATASET_NAME=""
+export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
+```
+
+> 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths.
+
+---
+
+## 3. Distill Training
+
+The distillation is a DMD (distribution matching distillation) run with three networks: a sparse **generator** (the FastWan student, `FastVideoWanTransformer3DModel`), a sparse **fake_score** critic (same class), and a frozen **dense** `real_score` teacher (`WanTransformer3DModel`). The student and critic carry the VSA block-sparse self-attention; the teacher keeps full attention.
+
+### 3.1 Download Pretrained Model
+
+```bash
+# Create model directory
+mkdir -p models/Diffusion_Transformer
+
+# Download Wan2.1 official weights
+modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
+```
+
+### 3.2 Quick Start (DeepSpeed-Zero-2)
+
+After downloading the data as per **2.1 Quick Test Dataset** and the weights as per **3.1 Download Pretrained Model**, you can directly copy and run the quick start command.
+
+DeepSpeed-Zero-2 and FSDP are recommended for training. Here we use DeepSpeed-Zero-2 as an example.
+
+The difference between DeepSpeed-Zero-2 and FSDP is whether the model weights are sharded. **If you experience insufficient GPU memory when using multiple GPUs with DeepSpeed-Zero-2**, you can switch to FSDP for training.
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
+```
+
+### 3.3 Common Training Parameters
+
+**Key Parameter Descriptions**:
+
+| Parameter | Description | Example Value |
+|-----|------|-------|
+| `--config_path` | Config file path | `config/wan2.1/wan_civitai.yaml` |
+| `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` |
+| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` |
+| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` |
+| `--train_batch_size` | Batch size per GPU | 1 |
+| `--image_sample_size` | Maximum image training resolution | 640 |
+| `--video_sample_size` | Maximum video training resolution | 640 |
+| `--token_sample_size` | Token sample size | 640 |
+| `--video_sample_stride` | Video sampling stride | 2 |
+| `--video_sample_n_frames` | Video sampling frames | 81 |
+| `--gradient_accumulation_steps` | Gradient accumulation steps (effectively increases batch) | 1 |
+| `--dataloader_num_workers` | DataLoader worker processes | 8 |
+| `--num_train_epochs` | Number of training epochs | 100 |
+| `--checkpointing_steps` | Save checkpoint every N steps | 50 |
+| `--learning_rate` | Initial learning rate (generator) | 2e-06 |
+| `--learning_rate_critic` | Initial learning rate (critic) | 2e-06 |
+| `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` |
+| `--lr_warmup_steps` | Learning rate warmup steps | 100 |
+| `--seed` | Random seed | 42 |
+| `--output_dir` | Output directory | `output_dir_fastvideo_wan2.1_1.3b` |
+| `--gradient_checkpointing` | Enable gradient checkpointing | - |
+| `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
+| `--adam_weight_decay` | AdamW weight decay | 3e-2 |
+| `--adam_epsilon` | AdamW epsilon | 1e-10 |
+| `--vae_mini_batch` | VAE encoding mini-batch size | 1 |
+| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
+| `--enable_bucket` | Enable bucket training, no cropping, group by resolution | - |
+| `--random_hw_adapt` | Auto-scale images/videos to random size in `[min_size, max_size]` range | - |
+| `--training_with_video_token_length` | Train based on token length, supports arbitrary resolutions | - |
+| `--uniform_sampling` | Uniform timestep sampling | - |
+| `--low_vram` | Low VRAM mode | - |
+| `--train_mode` | Training mode: `normal` (standard) or `i2v` (image-to-video) | `normal` |
+| `--resume_from_checkpoint` | Resume training path, use `"latest"` to auto-select latest checkpoint | None |
+| `--validation_steps` | Run validation every N steps | 2000 |
+| `--validation_epochs` | Run validation every N epochs | 5 |
+| `--validation_prompts` | Prompts for video generation validation | `"A dog shaking head..."` |
+| `--trainable_modules` | Trainable modules (`"."` means all modules) | `"."` |
+
+**Distill-Specific Parameters**:
+
+| Parameter | Description | Example Value |
+|-----|------|-------|
+| `--denoising_step_indices_list` | Denoising step list (core distill parameter) | `1000 750 500 250` |
+| `--real_guidance_scale` | Real guidance scale for scoring | 6.0 |
+| `--fake_guidance_scale` | Fake guidance scale for scoring | 0.0 |
+| `--gen_update_interval` | Generator update interval | 5 |
+| `--negative_prompt` | Negative prompt for distillation | Chinese negative prompt |
+| `--validation_paths` | Validation image paths for I2V mode | Image path list |
+| `--train_sampling_steps` | Training sampling steps | 1000 |
+
+**VSA Sparsity**:
+
+- The student/critic self-attention is VSA block-sparse. The keep-ratio `student_topk` defaults to `0.20` (i.e. 80% sparsity) inside `FastVideoWanTransformer3DModel`; it is a model-class default, not a CLI flag.
+- The learned compression gate `to_gate_compress` is created zero-initialized and becomes trainable automatically under `--trainable_modules "."` (its parameter names contain a dot, like every other module).
+- The `real_score` teacher stays dense (`WanTransformer3DModel`) and frozen, so the sparsity is distilled into the student rather than copied from the teacher.
+- Keep the inference-time `student_topk` equal to the value the checkpoint was trained at; changing it rebuilds the attention op and resets the learned gate.
+
+**Sample Size Configuration Guide**:
+- `video_sample_size` represents the resolution size of videos; when `random_hw_adapt` is True, it represents the minimum value between video and image resolutions.
+- `image_sample_size` represents the resolution size of images; when `random_hw_adapt` is True, it represents the maximum value between video and image resolutions.
+- `token_sample_size` represents the resolution corresponding to the maximum token length when `training_with_video_token_length` is True.
+- Due to potential confusion in configuration, **if you don't require arbitrary resolution for finetuning**, it is recommended to set `video_sample_size`, `image_sample_size`, and `token_sample_size` to the same fixed value, such as **(320, 480, 512, 640, 960)**.
+  - **All set to 320** represents **240P**.
+  - **All set to 480** represents **320P**.
+  - **All set to 640** represents **480P**.
+  - **All set to 960** represents **720P**.
+
+**Other Parameter Explanations**:
+- `enable_bucket` is used to enable bucket training. When enabled, the model does not crop the images and videos at the center, but instead, it trains the entire images and videos after grouping them into buckets based on resolution.
+- `random_hw_adapt` is used to enable automatic height and width scaling for images and videos. When `random_hw_adapt` is enabled, the training images will have their height and width set to `image_sample_size` as the maximum and `min(video_sample_size, 512)` as the minimum.
+- `training_with_video_token_length` specifies training the model according to token length. For training images and videos, the height and width will be set to `image_sample_size` as the maximum and `video_sample_size` as the minimum.
+- `train_mode` is used to specify the training mode, which can be either normal or i2v. If you only wish to achieve text-to-video generation, keep `normal`.
+- `resume_from_checkpoint` is used to set the training should be resumed from a previous checkpoint. Use a path or `"latest"` to automatically select the last available checkpoint.
+
+### 3.4 Training Validation
+
+You can configure validation parameters to periodically generate test videos during training, allowing you to monitor training progress and model quality.
+
+**Validation Parameter Descriptions**:
+
+| Parameter | Description | Recommended Value |
+|------|------|--------|
+| `--validation_steps` | Run validation every N steps | 2000 |
+| `--validation_epochs` | Run validation every N epochs | 5 |
+| `--validation_prompts` | Prompts for video generation validation | English prompts |
+| `--validation_paths` | Validation image paths for I2V mode (i2v mode only) | `"asset/1.png"` |
+
+**Normal Mode Example** (T2V Validation):
+
+```bash
+  --validation_steps=2000 \
+  --validation_epochs=5 \
+  --validation_prompts="A dog shaking head. The video is of high quality, and the view is very clear."
+```
+
+**Notes**:
+- Validation videos will be saved to the `output_dir` directory.
+- Multiple prompts validation format: `--validation_prompts "prompt1" "prompt2" "prompt3"`.
+- Distill model validation uses the few steps defined in `denoising_step_indices_list` and runs CFG-free (`guidance_scale=1.0`), matching inference.
+
+### 3.5 Training with FSDP
+
+**If you experience insufficient GPU memory when using multiple GPUs with DeepSpeed-Zero-2**, you can switch to FSDP for training. `scripts/fastvideo/train_distill_wan2.1.sh` is a ready-to-run FSDP launcher (it wraps `WanAttentionBlock` with `FULL_SHARD`).
+
+```sh
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
+```
+
+### 3.6 Other Backends
+
+#### 3.6.1 Training with DeepSpeed-Zero-3
+
+DeepSpeed Zero-3 is not highly recommended at the moment. In this repository, using FSDP has fewer errors and is more stable.
+
+DeepSpeed Zero-3 is suitable for 14B Wan at high resolutions. After training, you can use the following command to get the final model:
+```bash
+python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/checkpoint-{your-num-steps}-outputs --max_shard_size 80GB --safe_serialization
+```
+
+Training shell command is as follows:
+```sh
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+NCCL_DEBUG=INFO
+
+accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
+```
+
+#### 3.6.2 Training Without DeepSpeed and FSDP
+
+**This approach is not recommended because without memory-saving backends, it easily causes insufficient GPU memory**. This is only provided as a reference shell for training.
+
+```sh
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
+```
+
+### 3.7 Multi-Machine Distributed Training
+
+**Suitable for**: Ultra-large datasets, faster training speed
+
+#### 3.7.1 Environment Configuration
+
+Assuming 2 machines, each with 8 GPUs:
+
+**Machine 0 (Master)**:
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export MASTER_ADDR="192.168.1.100"  # Master machine IP
+export MASTER_PORT=10086
+export WORLD_SIZE=2                  # Total number of machines
+export NUM_PROCESS=16                # Total processes = machines × 8
+export RANK=0                        # Current machine rank (0 or 1)
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
+```
+
+**Machine 1 (Worker)**:
+```bash
+export MASTER_ADDR="192.168.1.100"  # Same as Master
+export MASTER_PORT=10086
+export WORLD_SIZE=2
+export NUM_PROCESS=16
+export RANK=1  # Note this is 1
+NCCL_DEBUG=INFO
+
+# Use the same accelerate launch command as Machine 0
+```
+
+#### 3.7.2 Multi-Machine Training Notes
+
+- **Network Requirements**:
+   - RDMA/InfiniBand recommended (high performance)
+   - Without RDMA, add environment variables:
+     ```bash
+     export NCCL_IB_DISABLE=1
+     export NCCL_P2P_DISABLE=1
+     ```
+- **Data Synchronization**: All machines must be able to access the same data paths (NFS/shared storage)
+
+---
+
+## 4. Inference Testing
+
+### 4.1 Inference Parameters
+
+**Key Parameter Descriptions** (from `examples/fastvideo/predict_t2v_wan2.1.py`):
+
+| Parameter | Description | Example Value |
+|------|------|-------|
+| `GPU_memory_mode` | GPU memory mode, see table below for options | `model_cpu_offload` |
+| `ulysses_degree` | Ulysses parallelism degree (VSA is single-GPU only, keep 1) | 1 |
+| `ring_degree` | Ring parallelism degree (VSA is single-GPU only, keep 1) | 1 |
+| `config_path` | Config file path | `config/wan2.1/wan_civitai.yaml` |
+| `model_name` | Base model path (config + dense backbone + VAE + T5) | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
+| `student_topk` | VSA keep-ratio; must equal the trained value | `0.20` |
+| `sampler_name` | Sampler type: `Flow`, `Flow_Unipc`, `Flow_DPM++` | `Flow_Unipc` |
+| `shift` | Noise schedule shift (official FastWan flow_shift) | `3.0` |
+| `transformer_path` | Path to the FastWan student Transformer weights | trained checkpoint / official FastWan |
+| `vae_path` | Path to trained VAE weights | `None` |
+| `lora_path` | LoRA weights path | `None` |
+| `sample_size` | Generated video resolution `[height, width]` | `[480, 848]` |
+| `video_length` | Number of generated frames | `81` |
+| `fps` | Frames per second | `16` |
+| `weight_dtype` | Model weight precision, use `torch.float16` for GPUs without bf16 support | `torch.bfloat16` |
+| `prompt` | Positive prompt describing generated content | `"A cinematic aerial shot..."` |
+| `negative_prompt` | Negative prompt to avoid certain content | Chinese negative prompt |
+| `guidance_scale` | Guidance strength (the DMD student is CFG-free) | 1.0 |
+| `seed` | Random seed for reproducibility | 1024 |
+| `num_inference_steps` | Number of inference steps (FastWan uses 3) | 3 |
+| `save_path` | Path to save generated videos | `samples/fastvideo-wan2.1-official` |
+
+**GPU Memory Mode Descriptions**:
+
+| Mode | Description | Memory Usage |
+|------|------|---------|
+| `model_full_load` | Full model loaded to GPU | Highest |
+| `model_full_load_and_qfloat8` | Full load + FP8 quantization | High |
+| `model_cpu_offload` | Offload model to CPU after use | Medium |
+| `model_cpu_offload_and_qfloat8` | CPU offload + FP8 quantization | Medium-Low |
+| `model_group_offload` | Layer groups switch between CPU/CUDA | Low |
+| `sequential_cpu_offload` | Layer-by-layer offload (slowest) | Lowest |
+
+### 4.2 Text-to-Video (T2V) Inference
+
+Run single-GPU inference:
+
+```bash
+python examples/fastvideo/predict_t2v_wan2.1.py
+```
+
+Edit `examples/fastvideo/predict_t2v_wan2.1.py` according to your needs. For first-time inference, focus on the following key parameters. For other parameters, refer to the inference parameter descriptions above.
+
+```python
+# Choose based on GPU memory
+GPU_memory_mode = "model_cpu_offload"
+# Base model path (config + dense backbone + VAE + T5)
+model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+# VSA keep-ratio; must equal the value the checkpoint was trained at
+student_topk = 0.20
+# Path to the FastWan student weights (your trained checkpoint, or the official FastWan release)
+transformer_path = "models/Personalized_Model/FastWan2.1-T2V-1.3B.pth"
+# FastWan sampler + shift (CFG-free few-step student)
+sampler_name = "Flow_Unipc"
+shift = 3.0
+guidance_scale = 1.0
+num_inference_steps = 3
+# Write based on your generation content
+prompt = "A cinematic aerial shot of a snow-covered mountain village at dawn..."
+# ...
+```
+
+> 💡 The official `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` checkpoint ships in diffusers key naming; `examples/fastvideo/predict_t2v_wan2.1.py` remaps it to the VideoX-Fun civitai layout (including the VSA `to_gate_compress` gate) before loading. A checkpoint produced by this training guide is already in the native layout.
+
+### 4.3 Single-GPU Only
+
+The VSA block-sparse kernel has no sequence-parallel path, so a FastWan student runs on a **single GPU**. Keep `ulysses_degree = ring_degree = 1`; calling `enable_multi_gpus_inference()` raises otherwise. To fit larger resolutions on one card, lower the memory footprint via `GPU_memory_mode` (e.g. `model_cpu_offload` or `sequential_cpu_offload`) rather than splitting across GPUs.
+
+---
+
+## 5. Additional Resources
+
+- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
+- **FastVideo GitHub**: https://github.com/hao-ai-lab/FastVideo
+- **VSA Paper (Video Sparse Attention)**: https://arxiv.org/abs/2505.13389
@@ -0,0 +1,648 @@
+# FastVideo VSA 蒸馏训练指南
+
+本文档提供了将 Wan2.1 蒸馏为 FastWan2.1（采用 VSA，Video Sparse Attention）的完整工作流，包括环境配置、数据准备、分布式训练和推理测试。
+
+> **注意**: FastVideo VSA 将块稀疏视频注意力（由可学习门控为每个 query 保留约 20% 的 key 块）与 DMD 少步蒸馏结合，使学生模型在约 3 步、无 CFG 的条件下以稀疏注意力推理。训练配方、数据管线与分布式后端与 TurboDiffusion 蒸馏指南完全一致，仅稀疏学生模型类不同。
+
+---
+
+## 目录
+- [一、环境配置](#一环境配置)
+- [二、数据准备](#二数据准备)
+  - [2.1 快速测试数据集](#21-快速测试数据集)
+  - [2.2 数据集结构](#22-数据集结构)
+  - [2.3 metadata.json 格式](#23-metadatajson-格式)
+  - [2.4 相对路径与绝对路径使用方案](#24-相对路径与绝对路径使用方案)
+- [三、蒸馏训练](#三蒸馏训练)
+  - [3.1 下载预训练模型](#31-下载预训练模型)
+  - [3.2 快速开始（DeepSpeed-Zero-2）](#32-快速开始deepspeed-zero-2)
+  - [3.3 常见训练参数](#33-常见训练参数)
+  - [3.4 训练验证](#34-训练验证)
+  - [3.5 使用 FSDP 训练](#35-使用-fsdp-训练)
+  - [3.6 其他后端](#36-其他后端)
+  - [3.7 多机分布式训练](#37-多机分布式训练)
+- [四、推理测试](#四推理测试)
+  - [4.1 推理参数解析](#41-推理参数解析)
+  - [4.2 文本生成视频（T2V）推理](#42-文本生成视频t2v推理)
+  - [4.3 仅支持单卡](#43-仅支持单卡)
+- [五、更多资源](#五更多资源)
+
+---
+
+## 一、环境配置
+
+**方式一：使用 requirements.txt**
+
+```bash
+pip install -r requirements.txt
+```
+
+**方式二：手动安装依赖**
+
+```bash
+pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
+pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
+pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
+pip install yunchang xfuser modelscope openpyxl
+pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
+pip install opencv-python-headless
+pip install deepspeed==0.17.0 numpy==1.26.4
+```
+
+> 💡 VSA 块稀疏注意力运行在 Triton 内核上；Triton 随 PyTorch（>=2.1）在 Linux 上一并提供，无需额外安装。
+
+**方式三：使用 Docker**
+
+使用 Docker 时，请确保机器上已正确安装 GPU 驱动与 CUDA 环境，然后执行以下命令：
+
+```
+# 拉取镜像
+docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+
+# 进入容器
+docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+```
+
+---
+
+## 二、数据准备
+
+### 2.1 快速测试数据集
+
+我们提供了一个包含若干训练样本的测试数据集。
+
+```bash
+# 下载官方演示数据集
+modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
+```
+
+### 2.2 数据集结构
+
+```
+📦 datasets/
+├── 📂 my_dataset/
+│   ├── 📂 train/
+│   │   ├── 📄 video001.mp4
+│   │   ├── 📄 video002.mp4
+│   │   └── 📄 ...
+│   └── 📄 metadata.json
+```
+
+### 2.3 metadata.json 格式
+
+**相对路径格式**（示例）：
+```json
+[
+  {
+    "file_path": "train/video001.mp4",
+    "text": "A beautiful sunset over the ocean, golden hour lighting",
+    "type": "video",
+    "width": 1024,
+    "height": 1024
+  },
+  {
+    "file_path": "train/video002.mp4",
+    "text": "A person walking through a forest, cinematic view",
+    "type": "video",
+    "width": 1328,
+    "height": 1328
+  }
+]
+```
+
+**绝对路径格式**：
+```json
+[
+  {
+    "file_path": "/mnt/data/videos/sunset.mp4",
+    "text": "A beautiful sunset over the ocean",
+    "type": "video",
+    "width": 1024,
+    "height": 1024
+  }
+]
+```
+
+**关键字段说明**：
+- `file_path`：视频路径（相对或绝对路径）
+- `text`：视频描述（英文 prompt）
+- `type`：数据类型，固定为 `"video"`
+- `width` / `height`：视频尺寸（**建议**提供以便分桶训练。若不提供，训练时会自动读取，当数据存放在 OSS 等较慢系统上时会拖慢训练）。
+  - 可使用 `scripts/process_json_add_width_and_height.py` 为缺少宽高字段的 JSON 文件提取 width 和 height，同时支持图片与视频。
+  - 用法：`python scripts/process_json_add_width_and_height.py --input_file datasets/X-Fun-Videos-Demo/metadata.json --output_file datasets/X-Fun-Videos-Demo/metadata_add_width_height.json`。
+
+### 2.4 相对路径与绝对路径使用方案
+
+**相对路径**：
+
+若数据使用相对路径，在训练脚本中配置：
+
+```bash
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+```
+
+**绝对路径**：
+
+若数据使用绝对路径，在训练脚本中配置：
+
+```bash
+export DATASET_NAME=""
+export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
+```
+
+> 💡 **建议**：数据集较小且存放在本地时使用相对路径；数据集存放在外部存储（如 NAS、OSS）或跨多台机器共享时使用绝对路径。
+
+---
+
+## 三、蒸馏训练
+
+蒸馏是一次 DMD（distribution matching distillation）训练，包含三个网络：稀疏的 **generator**（FastWan 学生，`FastVideoWanTransformer3DModel`）、稀疏的 **fake_score** critic（同一个类），以及冻结的 **dense** `real_score` 教师（`WanTransformer3DModel`）。学生与 critic 携带 VSA 块稀疏自注意力，教师保持全注意力。
+
+### 3.1 下载预训练模型
+
+```bash
+# 创建模型目录
+mkdir -p models/Diffusion_Transformer
+
+# 下载 Wan2.1 官方权重
+modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
+```
+
+### 3.2 快速开始（DeepSpeed-Zero-2）
+
+按照 **2.1 快速测试数据集** 下载数据、按照 **3.1 下载预训练模型** 下载权重后，即可直接复制运行快速开始命令。
+
+推荐使用 DeepSpeed-Zero-2 或 FSDP 训练，这里以 DeepSpeed-Zero-2 为例。
+
+DeepSpeed-Zero-2 与 FSDP 的区别在于模型权重是否分片。**若使用 DeepSpeed-Zero-2 多卡训练时显存不足**，可切换到 FSDP 训练。
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
+```
+
+### 3.3 常见训练参数
+
+**关键参数说明**：
+
+| 参数 | 描述 | 示例值 |
+|-----|------|-------|
+| `--config_path` | 配置文件路径 | `config/wan2.1/wan_civitai.yaml` |
+| `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` |
+| `--train_data_dir` | 训练数据目录 | `datasets/X-Fun-Videos-Demo/` |
+| `--train_data_meta` | 训练数据元信息文件 | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` |
+| `--train_batch_size` | 每张卡的 batch size | 1 |
+| `--image_sample_size` | 图片训练最大分辨率 | 640 |
+| `--video_sample_size` | 视频训练最大分辨率 | 640 |
+| `--token_sample_size` | Token 采样尺寸 | 640 |
+| `--video_sample_stride` | 视频采样步长 | 2 |
+| `--video_sample_n_frames` | 视频采样帧数 | 81 |
+| `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 |
+| `--dataloader_num_workers` | DataLoader 工作进程数 | 8 |
+| `--num_train_epochs` | 训练轮数 | 100 |
+| `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 |
+| `--learning_rate` | 初始学习率（generator） | 2e-06 |
+| `--learning_rate_critic` | 初始学习率（critic） | 2e-06 |
+| `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
+| `--lr_warmup_steps` | 学习率 warmup 步数 | 100 |
+| `--seed` | 随机种子 | 42 |
+| `--output_dir` | 输出目录 | `output_dir_fastvideo_wan2.1_1.3b` |
+| `--gradient_checkpointing` | 启用梯度检查点 | - |
+| `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
+| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
+| `--adam_epsilon` | AdamW epsilon | 1e-10 |
+| `--vae_mini_batch` | VAE 编码 mini-batch 大小 | 1 |
+| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
+| `--enable_bucket` | 启用分桶训练，不裁剪，按分辨率分组 | - |
+| `--random_hw_adapt` | 将图片/视频自动缩放到 `[min_size, max_size]` 范围内的随机尺寸 | - |
+| `--training_with_video_token_length` | 基于 token 长度训练，支持任意分辨率 | - |
+| `--uniform_sampling` | 均匀时间步采样 | - |
+| `--low_vram` | 低显存模式 | - |
+| `--train_mode` | 训练模式：`normal`（标准）或 `i2v`（图生视频） | `normal` |
+| `--resume_from_checkpoint` | 续训路径，使用 `"latest"` 自动选择最新 checkpoint | None |
+| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
+| `--validation_epochs` | 每 N 轮运行一次验证 | 5 |
+| `--validation_prompts` | 用于视频生成验证的 prompt | `"A dog shaking head..."` |
+| `--trainable_modules` | 可训练模块（`"."` 表示全部模块） | `"."` |
+
+**蒸馏专用参数**：
+
+| 参数 | 描述 | 示例值 |
+|-----|------|-------|
+| `--denoising_step_indices_list` | 去噪步列表（蒸馏核心参数） | `1000 750 500 250` |
+| `--real_guidance_scale` | 打分用的 real guidance scale | 6.0 |
+| `--fake_guidance_scale` | 打分用的 fake guidance scale | 0.0 |
+| `--gen_update_interval` | generator 更新间隔 | 5 |
+| `--negative_prompt` | 蒸馏用的负向 prompt | 中文负向 prompt |
+| `--validation_paths` | I2V 模式的验证图片路径 | 图片路径列表 |
+| `--train_sampling_steps` | 训练采样步数 | 1000 |
+
+**VSA 稀疏度**：
+
+- 学生/critic 的自注意力为 VSA 块稀疏。保留比 `student_topk` 在 `FastVideoWanTransformer3DModel` 内默认为 `0.20`（即 80% 稀疏度）；它是模型类默认值，而非命令行旗标。
+- 可学习的压缩门控 `to_gate_compress` 以零初始化创建，并在 `--trainable_modules "."` 下自动变为可训练（其参数名与其它模块一样包含点号）。
+- `real_score` 教师保持 dense（`WanTransformer3DModel`）且冻结，因此稀疏性是被蒸馏进学生，而非从教师拷贝。
+- 推理时的 `student_topk` 须与 checkpoint 训练时的取值一致；修改它会重建注意力算子并重置已训练好的门控。
+
+**采样尺寸配置指南**：
+- `video_sample_size` 表示视频的分辨率尺寸；当 `random_hw_adapt` 为 True 时，表示视频与图片分辨率之间的最小值。
+- `image_sample_size` 表示图片的分辨率尺寸；当 `random_hw_adapt` 为 True 时，表示视频与图片分辨率之间的最大值。
+- `token_sample_size` 表示 `training_with_video_token_length` 为 True 时最大 token 长度对应的分辨率。
+- 由于配置容易混淆，**若微调不需要任意分辨率**，建议将 `video_sample_size`、`image_sample_size`、`token_sample_size` 设为同一固定值，例如 **(320, 480, 512, 640, 960)**。
+  - **全部设为 320** 代表 **240P**。
+  - **全部设为 480** 代表 **320P**。
+  - **全部设为 640** 代表 **480P**。
+  - **全部设为 960** 代表 **720P**。
+
+**其它参数说明**：
+- `enable_bucket` 用于启用分桶训练。启用后模型不在中心裁剪图片与视频，而是按分辨率分桶后对完整图片与视频进行训练。
+- `random_hw_adapt` 用于对图片与视频自动高宽缩放。启用后，训练图片的高宽以 `image_sample_size` 为最大、`min(video_sample_size, 512)` 为最小。
+- `training_with_video_token_length` 指定按 token 长度训练。训练图片与视频的高宽以 `image_sample_size` 为最大、`video_sample_size` 为最小。
+- `train_mode` 用于指定训练模式，可为 normal 或 i2v。若只需文本生成视频，保持 `normal`。
+- `resume_from_checkpoint` 用于设置从之前的 checkpoint 续训。使用路径或 `"latest"` 自动选择最近可用的 checkpoint。
+
+### 3.4 训练验证
+
+你可以配置验证参数，在训练期间周期性生成测试视频，以便监控训练进度与模型质量。
+
+**验证参数说明**：
+
+| 参数 | 描述 | 推荐值 |
+|------|------|--------|
+| `--validation_steps` | 每 N 步运行一次验证 | 2000 |
+| `--validation_epochs` | 每 N 轮运行一次验证 | 5 |
+| `--validation_prompts` | 用于视频生成验证的 prompt | 英文 prompt |
+| `--validation_paths` | I2V 模式的验证图片路径（仅 i2v 模式） | `"asset/1.png"` |
+
+**Normal 模式示例**（T2V 验证）：
+
+```bash
+  --validation_steps=2000 \
+  --validation_epochs=5 \
+  --validation_prompts="A dog shaking head. The video is of high quality, and the view is very clear."
+```
+
+**说明**：
+- 验证视频会保存到 `output_dir` 目录。
+- 多 prompt 验证格式：`--validation_prompts "prompt1" "prompt2" "prompt3"`。
+- 蒸馏模型验证使用 `denoising_step_indices_list` 定义的少步、并以无 CFG（`guidance_scale=1.0`）运行，与推理保持一致。
+
+### 3.5 使用 FSDP 训练
+
+**若使用 DeepSpeed-Zero-2 多卡训练时显存不足**，可切换到 FSDP 训练。`scripts/fastvideo/train_distill_wan2.1.sh` 是一个开箱即用的 FSDP launcher（它以 `FULL_SHARD` 包装 `WanAttentionBlock`）。
+
+```sh
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
+```
+
+### 3.6 其他后端
+
+#### 3.6.1 使用 DeepSpeed-Zero-3 进行训练
+
+目前不强烈推荐 DeepSpeed Zero-3。在本仓库中，使用 FSDP 报错更少、更稳定。
+
+DeepSpeed Zero-3 适合高分辨率下的 14B Wan。训练后可用以下命令得到最终模型：
+```bash
+python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/checkpoint-{your-num-steps}-outputs --max_shard_size 80GB --safe_serialization
+```
+
+训练 shell 命令如下：
+```sh
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+NCCL_DEBUG=INFO
+
+accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
+```
+
+#### 3.6.2 不使用 DeepSpeed 与 FSDP 训练
+
+**不推荐这种方式，因为没有省显存后端时极易导致显存不足**。这里仅作为训练参考 shell 提供。
+
+```sh
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
+```
+
+### 3.7 多机分布式训练
+
+**适用于**：超大数据集、更快的训练速度
+
+#### 3.7.1 环境配置
+
+假设有 2 台机器，每台 8 张 GPU：
+
+**机器 0（Master）**：
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export MASTER_ADDR="192.168.1.100"  # Master 机器 IP
+export MASTER_PORT=10086
+export WORLD_SIZE=2                  # 机器总数
+export NUM_PROCESS=16                # 总进程数 = 机器数 × 8
+export RANK=0                        # 当前机器 rank（0 或 1）
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
+```
+
+**机器 1（Worker）**：
+```bash
+export MASTER_ADDR="192.168.1.100"  # 与 Master 相同
+export MASTER_PORT=10086
+export WORLD_SIZE=2
+export NUM_PROCESS=16
+export RANK=1  # 注意这里是 1
+NCCL_DEBUG=INFO
+
+# 使用与机器 0 相同的 accelerate launch 命令
+```
+
+#### 3.7.2 多机训练注意事项
+
+- **网络要求**：
+   - 推荐 RDMA/InfiniBand（高性能）
+   - 没有 RDMA 时，添加环境变量：
+     ```bash
+     export NCCL_IB_DISABLE=1
+     export NCCL_P2P_DISABLE=1
+     ```
+- **数据同步**：所有机器必须能访问相同的数据路径（NFS/共享存储）
+
+---
+
+## 四、推理测试
+
+### 4.1 推理参数解析
+
+**关键参数说明**（取自 `examples/fastvideo/predict_t2v_wan2.1.py`）：
+
+| 参数 | 描述 | 示例值 |
+|------|------|-------|
+| `GPU_memory_mode` | GPU 显存模式，选项见下表 | `model_cpu_offload` |
+| `ulysses_degree` | Ulysses 并行度（VSA 仅单卡，保持 1） | 1 |
+| `ring_degree` | Ring 并行度（VSA 仅单卡，保持 1） | 1 |
+| `config_path` | 配置文件路径 | `config/wan2.1/wan_civitai.yaml` |
+| `model_name` | 基座模型路径（config + dense 主干 + VAE + T5） | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` |
+| `student_topk` | VSA 保留比；须与训练时取值一致 | `0.20` |
+| `sampler_name` | 采样器类型：`Flow`、`Flow_Unipc`、`Flow_DPM++` | `Flow_Unipc` |
+| `shift` | 噪声调度 shift（官方 FastWan flow_shift） | `3.0` |
+| `transformer_path` | FastWan 学生 Transformer 权重路径 | 训练 checkpoint / 官方 FastWan |
+| `vae_path` | 训练好的 VAE 权重路径 | `None` |
+| `lora_path` | LoRA 权重路径 | `None` |
+| `sample_size` | 生成视频分辨率 `[height, width]` | `[480, 848]` |
+| `video_length` | 生成帧数 | `81` |
+| `fps` | 每秒帧数 | `16` |
+| `weight_dtype` | 模型权重精度，不支持 bf16 的 GPU 用 `torch.float16` | `torch.bfloat16` |
+| `prompt` | 描述生成内容的正向 prompt | `"A cinematic aerial shot..."` |
+| `negative_prompt` | 用于规避特定内容的负向 prompt | 中文负向 prompt |
+| `guidance_scale` | 引导强度（DMD 学生无 CFG） | 1.0 |
+| `seed` | 复现用随机种子 | 1024 |
+| `num_inference_steps` | 推理步数（FastWan 用 3 步） | 3 |
+| `save_path` | 生成视频保存路径 | `samples/fastvideo-wan2.1-official` |
+
+**GPU 显存模式说明**：
+
+| 模式 | 描述 | 显存占用 |
+|------|------|---------|
+| `model_full_load` | 整个模型加载到 GPU | 最高 |
+| `model_full_load_and_qfloat8` | 全量加载 + FP8 量化 | 高 |
+| `model_cpu_offload` | 用后将模型卸载到 CPU | 中 |
+| `model_cpu_offload_and_qfloat8` | CPU 卸载 + FP8 量化 | 中低 |
+| `model_group_offload` | 层组在 CPU/CUDA 间切换 | 低 |
+| `sequential_cpu_offload` | 逐层卸载（最慢） | 最低 |
+
+### 4.2 文本生成视频（T2V）推理
+
+运行单卡推理：
+
+```bash
+python examples/fastvideo/predict_t2v_wan2.1.py
+```
+
+根据需求编辑 `examples/fastvideo/predict_t2v_wan2.1.py`。首次推理重点关注以下关键参数，其它参数参见上文推理参数说明。
+
+```python
+# 根据 GPU 显存选择
+GPU_memory_mode = "model_cpu_offload"
+# 基座模型路径（config + dense 主干 + VAE + T5）
+model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+# VSA 保留比；须与 checkpoint 训练时取值一致
+student_topk = 0.20
+# FastWan 学生权重路径（你训练的 checkpoint，或官方 FastWan 发布权重）
+transformer_path = "models/Personalized_Model/FastWan2.1-T2V-1.3B.pth"
+# FastWan 采样器 + shift（无 CFG 少步学生）
+sampler_name = "Flow_Unipc"
+shift = 3.0
+guidance_scale = 1.0
+num_inference_steps = 3
+# 根据你的生成内容编写
+prompt = "A cinematic aerial shot of a snow-covered mountain village at dawn..."
+# ...
+```
+
+> 💡 官方 `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` checkpoint 采用 diffusers 键名；`examples/fastvideo/predict_t2v_wan2.1.py` 会在加载前将其 remap 到 VideoX-Fun civitai 布局（包含 VSA `to_gate_compress` 门控）。本训练指南产出的 checkpoint 已是原生布局。
+
+### 4.3 仅支持单卡
+
+VSA 块稀疏内核没有序列并行路径，因此 FastWan 学生只能在**单卡**上运行。请保持 `ulysses_degree = ring_degree = 1`；否则调用 `enable_multi_gpus_inference()` 会 raise。要在单卡上容纳更大分辨率，请通过 `GPU_memory_mode`（如 `model_cpu_offload` 或 `sequential_cpu_offload`）降低显存占用，而不是跨卡拆分。
+
+---
+
+## 五、更多资源
+
+- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
+- **FastVideo GitHub**：https://github.com/hao-ai-lab/FastVideo
+- **VSA 论文（Video Sparse Attention）**：https://arxiv.org/abs/2505.13389
@@ -0,0 +1,48 @@
+# This shell script is used to distill a regular Wan2.1-1.3B into FastWan2.1-1.3B (VSA sparse student)
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+
+accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \
+    --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \
+    --fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --video_sample_stride=2 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --video_repeat=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --learning_rate_critic=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_fastvideo_wan2.1_1.3b" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --uniform_sampling \
+  --train_mode="normal" \
+  --trainable_modules "." \
+  --low_vram
@@ -172,6 +172,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -211,6 +212,7 @@ bash scripts/wan2.1_causal_forcing/train_ar_diffusion.sh
 | `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_ar_diffusion` |
 | `--gradient_checkpointing` | Enable activation checkpointing | - |
 | `--mixed_precision` | `fp16` / `bf16` | `bf16` |
+| `--fsdp_reduce_fp32` | FP32 gradient reduction under FSDP mixed precision; forward and parameter gathering stay low precision. CLI default is off; enabled in the launcher | enabled |
 | `--adam_weight_decay` | AdamW weight decay | 3e-2 |
 | `--adam_epsilon` | AdamW epsilon | 1e-10 |
 | `--max_grad_norm` | Gradient clipping threshold | 0.05 |
@@ -172,6 +172,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -211,6 +212,7 @@ bash scripts/wan2.1_causal_forcing/train_ar_diffusion.sh
 | `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_ar_diffusion` |
 | `--gradient_checkpointing` | 启用激活重计算 | - |
 | `--mixed_precision` | `fp16` / `bf16` | `bf16` |
+| `--fsdp_reduce_fp32` | FSDP 混合精度下使用 FP32 梯度归约，前向与参数聚合仍保持低精度。CLI 默认关闭，launcher 已启用 | 启用 |
 | `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
 | `--adam_epsilon` | AdamW epsilon | 1e-10 |
 | `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
@@ -183,6 +183,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --output_dir="output_dir_wan2.1_causal_forcing_ccd" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=0.0 \
  --adam_beta1=0.0 \
  --adam_beta2=0.999 \
@@ -228,6 +229,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh
 | `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_ccd` |
 | `--gradient_checkpointing` | Enable activation checkpointing | - |
 | `--mixed_precision` | `fp16` / `bf16` | `bf16` |
+| `--fsdp_reduce_fp32` | FP32 gradient reduction under FSDP mixed precision; forward and parameter gathering stay low precision. CLI default is off; enabled in the launcher | enabled |
 | `--adam_weight_decay` | AdamW weight decay | 0.0 |
 | `--adam_beta1` | AdamW beta1 (CCD uses 0.0) | 0.0 |
 | `--adam_beta2` | AdamW beta2 | 0.999 |
@@ -183,6 +183,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --output_dir="output_dir_wan2.1_causal_forcing_ccd" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=0.0 \
  --adam_beta1=0.0 \
  --adam_beta2=0.999 \
@@ -228,6 +229,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh
 | `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_ccd` |
 | `--gradient_checkpointing` | 启用激活重计算 | - |
 | `--mixed_precision` | `fp16` / `bf16` | `bf16` |
+| `--fsdp_reduce_fp32` | FSDP 混合精度下使用 FP32 梯度归约，前向与参数聚合仍保持低精度。CLI 默认关闭，launcher 已启用 | 启用 |
 | `--adam_weight_decay` | AdamW 权重衰减 | 0.0 |
 | `--adam_beta1` | AdamW beta1（CCD 使用 0.0） | 0.0 |
 | `--adam_beta2` | AdamW beta2 | 0.999 |
@@ -26,6 +26,7 @@ This document provides the complete workflow for **Causal-Forcing Stage 3 — Di
  - [5.1 Quick Start](#51-quick-start)
  - [5.2 Key Parameters](#52-key-parameters)
  - [5.3 DMD-Specific Parameters](#53-dmd-specific-parameters)
+  - [5.4 Optional One-Forcing GAN and Precision Controls](#54-optional-one-forcing-gan-and-precision-controls)
 - [6. Use the Trained Checkpoint](#6-use-the-trained-checkpoint)
 - [7. Additional Resources](#7-additional-resources)

@@ -103,7 +104,7 @@ The Stage 2 CCD checkpoint (generator/critic init) is loaded via `--ode_transfor

 ## 4. Prepare Training Data

-DMD uses a **TextDataset** (`train_mode="normal"`) — it only needs prompts, not video data, because the generator creates its own training samples via autoregressive rollout. However, a `metadata.json` with prompts is still required.
+With `train_mode="normal"`, teacher forcing disabled, and all GAN/R1/R2 weights set to zero, DMD uses a **TextDataset**: only prompts are required because the generator creates training samples via autoregressive rollout. Enabling any GAN/R1/R2 weight switches to **ImageVideoDataset** and requires captioned real videos; see Section 5.4.

 ### 4.1 Quick Demo Dataset

@@ -118,7 +119,7 @@ Same as Stage 1/2. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](./README_TRA

 ### 4.3 metadata.json Format

-DMD only uses the `"text"` field from each entry. The `file_path` and other fields are ignored when `--train_mode="normal"` (TextDataset mode).
+In prompt-only TextDataset mode, DMD only uses the `"text"` field from each entry. GAN/R1/R2 or teacher forcing requires video metadata with `file_path`, `type="video"`, and `text`, as in Stage 1/2.

 ```json
 [
@@ -131,7 +132,7 @@ DMD only uses the `"text"` field from each entry. The `file_path` and other fiel
 ]
 ```

-> **Note**: You can reuse any video metadata.json — only the `text` field matters for DMD.
+> **Note**: Existing video metadata can be reused. Only `text` is read in prompt-only mode; GAN training also reads the video files.

 ---

@@ -179,6 +180,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --output_dir="output_dir_wan2.1_causal_forcing_dmd" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=0.0 \
  --adam_beta1=0.0 \
  --adam_beta2=0.999 \
@@ -212,7 +214,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--real_score_pretrained_model_name_or_path` | 14B real-score teacher for DMD | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
 | `--config_path` | Model config YAML | `config/wan2.1/wan_civitai.yaml` |
 | `--train_data_dir` | Data root directory | `""` |
-| `--train_data_meta` | Path to `metadata.json` (prompts only) | `datasets/X-Fun-Videos-Demo/metadata.json` |
+| `--train_data_meta` | Metadata path: prompts for plain DMD, captioned videos for GAN/R1/R2 | `datasets/X-Fun-Videos-Demo/metadata.json` |
 | `--ode_transformer_path` | Stage 2 CCD checkpoint for generator/critic init | `$STAGE2_CKPT` |
 | `--train_batch_size` | Per-GPU batch size | 1 |
 | `--gradient_accumulation_steps` | Gradient accumulation steps | 1 |
@@ -264,11 +266,48 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--context_noise` | Context noise level for KV cache update | 0 |
 | `--use_teacher_forcing` | Enable teacher forcing (pass clean_x to transformer) | - |
 | `--teacher_forcing_prob` | Probability of applying teacher forcing per step (1.0 = always) | 1.0 |
-| `--train_mode` | Training mode: `normal` (TextDataset, prompt-only) or `i2v` | `normal` |
+| `--train_mode` | `normal` (prompt-only unless teacher forcing or GAN/R1/R2 is enabled) or `i2v` | `normal` |
 | `--resume_from_checkpoint` | Resume from checkpoint. Use `"latest"` to auto-select | `"latest"` |

 ---

+### 5.4 Optional One-Forcing GAN and Precision Controls
+
+These options reuse the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms in `scripts/wan2.1_self_forcing/train_distill.py`. All GAN/R1/R2 weights default to **0**, and both precision flags default to **off** in argparse. The AR diffusion, CCD, and DMD launchers enable `--fsdp_reduce_fp32`; GAN/R1/R2 and `--score_timestep_clamp` remain off. AR diffusion and CCD do not gain GAN losses.
+
+The discriminator shares the **fake-score backbone**, with one learnable register token per selected feature layer. Its head is optimized by the critic optimizer, not the generator optimizer. Generator loss is `DMD + weighted GAN`; critic loss is `denoising + weighted GAN + R1 + R2`.
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--gan_g_weight` | Non-saturating generator GAN loss weight | `0.0` |
+| `--gan_d_weight` | Discriminator GAN loss weight | `0.0` |
+| `--gan_feature_layers` | Fake-score block indices used by the head | `21 29` |
+| `--gan_block_ffn_dim` | Head cross-attention FFN dimension | `2048` |
+| `--gan_block_num_heads` | Head count; must divide fake-score hidden dimension | `12` |
+| `--relativistic_discriminator` | Use paired real/fake logit differences | off |
+| `--r1_weight` | APT-style finite-difference penalty on noisy real latents | `0.0` |
+| `--r1_sigma` | R1 perturbation standard deviation | `0.01` |
+| `--r2_weight` | Normalized finite-difference penalty on noisy generated latents | `0.0` |
+| `--r2_sigma` | R2 perturbation standard deviation | `0.01` |
+| `--fsdp_reduce_fp32` | FSDP FP32 gradient reduction; parameter gathering stays low precision | off (enabled in the launcher) |
+| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off |
+
+For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` for the generator and `softplus(-D_real) + softplus(D_fake)` for the discriminator. Relativistic losses instead use `softplus(-(D_fake - D_real))` and `softplus(-(D_real - D_fake))`.
+
+R1 follows the source's APT form: `r1_weight * mean((D_real_perturbed - D_real)^2)`, **without dividing by sigma**. R2 follows `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`. These are finite-difference penalties, not input-gradient double-backward regularizers. Real/fake samples share timestep and noise; perturbations use independent noise.
+
+To enable GAN, pass positive generator/discriminator weights and, optionally, R1/R2 weights to `train_causal_dmd.py`. Requirements:
+- Captioned real videos, not a prompt-only JSON. Use the Stage 1/2 video metadata schema and enough frames for `score_num_frames` and the selected causal block geometry.
+- `--gradient_accumulation_steps=1`; DeepSpeed is not supported by the GAN branch.
+- Valid feature-layer indices for the selected fake-score model. The default `21 29` layers and `12` heads fit Wan2.1-T2V-1.3B; other backbones may need different settings.
+- Enabling only `gan_g_weight` does not supply a discriminator classification-learning objective. Choose the discriminator weight deliberately; nonzero values are not a guaranteed quality improvement.
+
+Trainable fake-score and head master weights stay **FP32**, while forward uses the configured autocast/FSDP precision. Logits, softplus losses, and finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient communication precision, not the forward dtype or model architecture. The clamp changes score-training timesteps, not the generator rollout or validation schedule.
+
+Resume GAN checkpoints with the same GAN/head configuration. The head is saved with fake-score, including FP32 head tensors, and is recreated before loading. It is not part of the exported generator, so inference needs no discriminator or new GAN flags.
+
+---
+
 ## 6. Use the Trained Checkpoint

 The Stage 3 DMD checkpoint is the **final model** in the Causal-Forcing pipeline. Use it for inference:
@@ -300,4 +339,5 @@ python examples/wan2.1_causal_forcing/predict_t2v.py
 ## 7. Additional Resources

 - **Causal-Forcing Paper**: https://github.com/thu-ml/Causal-Forcing
+- **One-Forcing (optional GAN mechanism)**: https://github.com/Aurora-edu/One-Forcing
 - **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
@@ -26,6 +26,7 @@
  - [5.1 快速开始](#51-快速开始)
  - [5.2 关键参数](#52-关键参数)
  - [5.3 DMD 特有参数](#53-dmd-特有参数)
+  - [5.4 可选 One-Forcing GAN 与精度控制](#54-可选-one-forcing-gan-与精度控制)
 - [六、使用训练好的 Checkpoint](#六使用训练好的-checkpoint)
 - [七、更多资源](#七更多资源)

@@ -39,7 +40,7 @@ Stage 3 需要：
 2. **Wan2.1-T2V-14B** 模型：作为 DMD real-score teacher。

 ```bash
-# 示例：Stage 2 CCD 训练的 checkpoint
+# Example: Stage 2 CCD checkpoint
 export STAGE2_CKPT="output_dir_wan2.1_causal_forcing_ccd/checkpoint-5000/transformer/diffusion_pytorch_model.safetensors"
 ```

@@ -70,10 +71,10 @@ pip install deepspeed==0.17.0 numpy==1.26.4
 **方式 3：使用 Docker**

 ```bash
-# 拉取镜像
+# Pull the image
 docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun

-# 进入容器
+# Enter the container
 docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
 ```

@@ -87,13 +88,13 @@ DMD 需要**两个**预训练模型：
 - **Wan2.1-T2V-14B**：非因果的 real-score teacher，用于计算真实分布得分。

 ```bash
-# 创建模型目录
+# Create the model directory
 mkdir -p models/Diffusion_Transformer

-# 下载 Wan2.1 T2V 1.3B（student 基础模型）
+# Download the Wan2.1 T2V 1.3B student base model
 modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B

-# 下载 Wan2.1 T2V 14B（DMD real-score teacher）
+# Download the Wan2.1 T2V 14B real-score teacher
 modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B
 ```

@@ -103,12 +104,12 @@ Stage 2 CCD checkpoint（生成器/判别器初始化）通过 `--ode_transforme

 ## 四、数据准备

-DMD 使用 **TextDataset**（`train_mode="normal"`）— 只需要提示词，不需要视频数据，因为生成器通过自回归 rollout 自行创建训练样本。但仍需提供包含提示词的 `metadata.json`。
+当 `train_mode="normal"`、未启用 teacher forcing 且 GAN/R1/R2 权重均为零时，DMD 使用 **TextDataset**：只需要提示词，生成器通过自回归 rollout 自行创建训练样本。任一 GAN/R1/R2 权重为正时会切换为 **ImageVideoDataset**，需要带描述的真实视频，详见 5.4。

 ### 4.1 快速测试数据集

 ```bash
-# 下载官方示例数据集
+# Download the official demo dataset
 modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
 ```

@@ -118,7 +119,7 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun

 ### 4.3 metadata.json 格式

-DMD 仅使用每条记录的 `"text"` 字段。当 `--train_mode="normal"`（TextDataset 模式）时，`file_path` 和其他字段会被忽略。
+仅提示词的 TextDataset 模式只读取每条记录的 `"text"` 字段。GAN/R1/R2 或 teacher forcing 需要与 Stage 1/2 相同的视频元数据，包括 `file_path`、`type="video"` 和 `text`。

 ```json
 [
@@ -131,7 +132,7 @@ DMD 仅使用每条记录的 `"text"` 字段。当 `--train_mode="normal"`（Tex
 ]
 ```

-> **说明**：你可以复用任何视频 metadata.json — DMD 只关心 `text` 字段。
+> **说明**：可以复用现有视频元数据。仅提示词模式只读取 `text`，GAN 训练还会读取视频文件。

 ---

@@ -179,6 +180,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --output_dir="output_dir_wan2.1_causal_forcing_dmd" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=0.0 \
  --adam_beta1=0.0 \
  --adam_beta2=0.999 \
@@ -212,7 +214,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--real_score_pretrained_model_name_or_path` | 14B real-score teacher | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
 | `--config_path` | 模型配置 YAML | `config/wan2.1/wan_civitai.yaml` |
 | `--train_data_dir` | 数据根目录 | `""` |
-| `--train_data_meta` | `metadata.json` 路径（仅含提示词） | `datasets/X-Fun-Videos-Demo/metadata.json` |
+| `--train_data_meta` | 元数据路径：普通 DMD 使用提示词，GAN/R1/R2 使用带描述的视频 | `datasets/X-Fun-Videos-Demo/metadata.json` |
 | `--ode_transformer_path` | Stage 2 CCD checkpoint（生成器/判别器初始化） | `$STAGE2_CKPT` |
 | `--train_batch_size` | 每 GPU batch 大小 | 1 |
 | `--gradient_accumulation_steps` | 梯度累积步数 | 1 |
@@ -264,23 +266,60 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--context_noise` | KV 缓存更新的上下文噪声级别 | 0 |
 | `--use_teacher_forcing` | 启用 teacher forcing（将 clean_x 传给 transformer） | - |
 | `--teacher_forcing_prob` | 每步应用 teacher forcing 的概率（1.0 = 始终） | 1.0 |
-| `--train_mode` | 训练模式：`normal`（TextDataset，仅提示词）或 `i2v` | `normal` |
+| `--train_mode` | `normal`（未启用 teacher forcing 或 GAN/R1/R2 时仅提示词）或 `i2v` | `normal` |
 | `--resume_from_checkpoint` | 恢复训练。使用 `"latest"` 自动选择最新 checkpoint | `"latest"` |
 ---

+### 5.4 可选 One-Forcing GAN 与精度控制
+
+这些选项复用 `scripts/wan2.1_self_forcing/train_distill.py` 中的 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 机制。GAN/R1/R2 权重默认均为 **0**，两个精度开关在 argparse 中默认均**关闭**。AR diffusion、CCD 和 DMD 的 launcher 均启用 `--fsdp_reduce_fp32`；GAN/R1/R2 和 `--score_timestep_clamp` 仍关闭。AR diffusion 和 CCD 不添加 GAN 损失。
+
+判别器共享 **fake-score backbone**，每个选定特征层对应一个可学习 register token。判别器头由 critic optimizer 更新，不加入 generator optimizer。生成器损失为 `DMD + 加权 GAN`，critic 损失为 `denoising + 加权 GAN + R1 + R2`。
+
+| 参数 | 说明 | 默认值 |
+|------|------|--------|
+| `--gan_g_weight` | 非饱和生成器 GAN 损失权重 | `0.0` |
+| `--gan_d_weight` | 判别器 GAN 损失权重 | `0.0` |
+| `--gan_feature_layers` | 判别器头使用的 fake-score block 索引 | `21 29` |
+| `--gan_block_ffn_dim` | 判别器交叉注意力 FFN 维度 | `2048` |
+| `--gan_block_num_heads` | 判别器头数，须整除 fake-score 隐藏维度 | `12` |
+| `--relativistic_discriminator` | 使用配对真实/生成样本 logit 差值 | 关闭 |
+| `--r1_weight` | 噪声真实 latent 上的 APT 形式有限差分正则权重 | `0.0` |
+| `--r1_sigma` | R1 扰动标准差 | `0.01` |
+| `--r2_weight` | 噪声生成 latent 上的归一化有限差分正则权重 | `0.0` |
+| `--r2_sigma` | R2 扰动标准差 | `0.01` |
+| `--fsdp_reduce_fp32` | FSDP 使用 FP32 梯度归约，参数仍以低精度聚合 | 关闭（launcher 已启用） |
+| `--score_timestep_clamp` | 将 DMD、生成器 GAN 和 critic 时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关闭 |
+
+设 logits 为 `D_fake` 和 `D_real`，普通 GAN 的生成器损失为 `softplus(-D_fake)`，判别器损失为 `softplus(-D_real) + softplus(D_fake)`。相对论模式分别使用 `softplus(-(D_fake - D_real))` 和 `softplus(-(D_real - D_fake))`。
+
+R1 保持源实现的 APT 形式：`r1_weight * mean((D_real_perturbed - D_real)^2)`，**不除以 sigma**。R2 为 `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`。二者都是有限差分正则，不是需要 double backward 的输入梯度惩罚。真实/生成样本共用时间步和噪声，扰动噪声独立抽样。
+
+开启 GAN 时，向 `train_causal_dmd.py` 传入正的生成器/判别器权重，并按需配置 R1/R2。要求：
+- 使用带描述的真实视频，不能只提供提示词 JSON。元数据采用 Stage 1/2 视频格式，帧数须满足 `score_num_frames` 与因果块几何。
+- `--gradient_accumulation_steps=1`；GAN 分支不支持 DeepSpeed。
+- 特征层索引必须存在。默认 `21 29` 与 `12` 个头适用于 Wan2.1-T2V-1.3B，其他 backbone 可能需要调整。
+- 仅开启 `gan_g_weight` 不会提供判别器的分类学习目标。须明确选择判别器权重，非零权重不保证质量提升。
+
+可训练 fake-score 与判别器头的 master weights 保持 **FP32**，前向按配置使用 autocast/FSDP 混合精度。Logits、softplus 损失及有限差分归约使用 FP32。`--fsdp_reduce_fp32` 只改变梯度通信精度，不改变前向 dtype 或模型结构。时间步 clamp 只作用于 score 训练，不改变生成器 rollout 或 validation schedule。
+
+恢复 GAN checkpoint 时应使用相同的 GAN/判别器头配置。判别器头随 fake-score 保存，包括 FP32 判别器头张量，并在加载前重建。导出的 generator 不包含该头，推理无需判别器或新增 GAN 参数。
+
+---
+
 ## 六、使用训练好的 Checkpoint

 Stage 3 DMD checkpoint 是 Causal-Forcing 流水线的**最终模型**，直接用于推理：

 ```python
-# 在 examples/wan2.1_causal_forcing/predict_t2v.py 中
+# In examples/wan2.1_causal_forcing/predict_t2v.py
 transformer_path = "output_dir_wan2.1_causal_forcing_dmd/checkpoint-{N}/diffusion_pytorch_model.safetensors"

-# DMD Stage 3 推理配置
-guidance_scale      = 1.0    # CFG 已烘焦到蒸馏权重中
-num_inference_steps = 4      # 4 步 DMD
+# DMD Stage 3 inference config
+guidance_scale      = 1.0    # CFG is baked into distilled weights
+num_inference_steps = 4      # 4-step DMD
 stochastic_sampling = True
-num_frame_per_block = 3      # 分块生成
+num_frame_per_block = 3      # Chunk-wise generation
 ```

 或运行：
@@ -299,4 +338,5 @@ python examples/wan2.1_causal_forcing/predict_t2v.py
 ## 七、更多资源

 - **Causal-Forcing 论文**：https://github.com/thu-ml/Causal-Forcing
+- **One-Forcing（可选 GAN 机制）**：https://github.com/Aurora-edu/One-Forcing
 - **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
@@ -18,6 +18,7 @@ scripts/wan2.1_self_forcing/train_distill.py and Causal-Forcing
 # See the License for the specific language governing permissions and

 import argparse
+import dataclasses
 import gc
 import logging
 import math
@@ -314,6 +315,10 @@ def parse_args():
    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument(
+        "--fsdp_reduce_fp32", action="store_true",
+        help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
+    )
    parser.add_argument(
        "--logging_dir",
        type=str,
@@ -596,6 +601,9 @@ def main():
        else:
            fsdp_stage = 0
        print(f"Using FSDP stage: {fsdp_stage}")
+        if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
+            fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
+            print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")

        args.use_fsdp = True
        if fsdp_stage == 3:
@@ -1254,6 +1262,8 @@ def main():
    training_weights_table = training_weights_table.to(accelerator.device)

    for epoch in range(first_epoch, args.num_train_epochs):
+        if global_step >= args.max_train_steps:
+            break
        train_loss = 0.0
        batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
        for step, batch in enumerate(train_dataloader):
@@ -1487,6 +1497,9 @@ def main():
            if global_step >= args.max_train_steps:
                break

+        if global_step >= args.max_train_steps:
+            break
+
        if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
            with progress_bar.paused():
                log_validation(
@@ -18,6 +18,7 @@ scripts/wan2.1_self_forcing/train_distill.py and Causal-Forcing
 # See the License for the specific language governing permissions and

 import argparse
+import dataclasses
 import gc
 import logging
 import math
@@ -317,6 +318,10 @@ def parse_args():
    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument(
+        "--fsdp_reduce_fp32", action="store_true",
+        help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
+    )
    parser.add_argument(
        "--logging_dir",
        type=str,
@@ -610,6 +615,9 @@ def main():
        else:
            fsdp_stage = 0
        print(f"Using FSDP stage: {fsdp_stage}")
+        if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
+            fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
+            print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")

        args.use_fsdp = True
        if fsdp_stage == 3:
@@ -769,7 +777,7 @@ def main():
            os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
            transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
            low_cpu_mem_usage=True,
-            torch_dtype=weight_dtype,
+            torch_dtype=torch.float32,
        )
        ema_transformer3d.num_frame_per_block = args.num_frame_per_block
        ema_transformer3d.independent_first_frame = args.independent_first_frame
@@ -1437,6 +1445,8 @@ def main():
            torch.cuda.empty_cache()

    for epoch in range(first_epoch, args.num_train_epochs):
+        if global_step >= args.max_train_steps:
+            break
        train_loss = 0.0
        batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
        for step, batch in enumerate(train_dataloader):
@@ -1737,6 +1747,9 @@ def main():
            if global_step >= args.max_train_steps:
                break

+        if global_step >= args.max_train_steps:
+            break
+
        if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
            with progress_bar.paused():
                log_validation(
@@ -17,6 +17,7 @@

 import argparse
 import contextlib
+import dataclasses
 import gc
 import json
 import logging
@@ -74,10 +75,12 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512,
                             get_random_mask)
 from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
                               WanTransformer3DModel,
-                               WanTransformer3DModel_SelfForcing)
+                               WanTransformer3DModel_SelfForcing,
+                               add_wan_gan_discriminator)
 from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
                                 WanSelfForcingPipeline)
 from videox_fun.utils.discrete_sampler import DiscreteSampling
+from videox_fun.utils.fsdp_ema import FSDPEMA
 from videox_fun.utils.tqdm_bar import PauseAwareTqdm
 from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
                                    get_image_to_video_latent,
@@ -158,6 +161,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi
    return pred_for_score, score_num_frames, True


+def repeat_batch(value, repeats):
+    """Repeat a batched tensor or per-sample list along the batch dimension."""
+    if value is None or repeats == 1:
+        return value
+    if torch.is_tensor(value):
+        return torch.cat([value] * repeats)
+    return list(value) * repeats
+
+
 def filter_kwargs(cls, kwargs):
    import inspect
    sig = inspect.signature(cls.__init__)
@@ -494,6 +506,14 @@ def parse_args():
        ),
    )
    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.")
+    parser.add_argument(
+        "--fsdp_reduce_fp32", action="store_true",
+        help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.",
+    )
+    parser.add_argument(
+        "--score_timestep_clamp", action="store_true",
+        help="Clamp DMD, GAN and critic timesteps to [0.02, 0.98] * num_train_timesteps.",
+    )
    parser.add_argument(
        "--dataloader_num_workers",
        type=int,
@@ -838,6 +858,47 @@ def parse_args():
        help="Probability of applying teacher forcing per step (1.0 = always)"
    )

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

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

    if args.report_to == "wandb" and args.hub_token is not None:
        raise ValueError(
@@ -875,6 +944,8 @@ def main():

    deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None
    fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None
+    if use_gan and deepspeed_plugin is not None:
+        raise ValueError("The GAN branch does not support DeepSpeed.")
    if deepspeed_plugin is not None:
        zero_stage = int(deepspeed_plugin.zero_stage)
        fsdp_stage = 0
@@ -896,6 +967,9 @@ def main():
        else:
            fsdp_stage = 0
        print(f"Using FSDP stage: {fsdp_stage}")
+        if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
+            fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
+            print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")

        args.use_fsdp = True
        if fsdp_stage == 3:
@@ -1111,28 +1185,55 @@ def main():
                param.requires_grad = True
                break

+    if use_gan:
+        if (not args.gan_feature_layers or any(i < 0 or i >= len(fake_score_transformer3d.blocks) for i in args.gan_feature_layers)
+                or args.gan_block_ffn_dim <= 0 or args.gan_block_num_heads <= 0
+                or fake_score_transformer3d.dim % args.gan_block_num_heads != 0):
+            raise ValueError("GAN feature layers must exist and the positive head count must divide the fake_score hidden dimension.")
+        # Keep trainable discriminator master weights in fp32; forward uses autocast.
+        add_wan_gan_discriminator(
+            fake_score_transformer3d,
+            args.gan_feature_layers,
+            args.gan_block_ffn_dim,
+            args.gan_block_num_heads,
+            torch.float32,
+            seed=args.seed or 0,
+        )
+
    # Create EMA for the generator_transformer3d.
    if args.use_ema:
        if zero_stage == 3:
            raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.")

-        ema_generator_transformer3d = EMAModel(
-            generator_transformer3d.parameters(),
-            model_cls=WanTransformer3DModel_SelfForcing,
-            model_config=generator_transformer3d.config,
-        )
+        if fsdp_plugin is not None:
+            if getattr(fsdp_plugin, "fsdp_version", 1) != 1:
+                raise NotImplementedError("FSDPEMA 仅支持 FSDP1。")
+            ema_module = WanTransformer3DModel_SelfForcing.from_config(generator_transformer3d.config)
+            ema_generator_transformer3d = FSDPEMA(
+                ema_module, source=generator_transformer3d,
+                accelerator=accelerator, fsdp_plugin=fsdp_plugin,
+            )
+        else:
+            ema_generator_transformer3d = EMAModel(
+                generator_transformer3d.parameters(),
+                model_cls=WanTransformer3DModel_SelfForcing,
+                model_config=generator_transformer3d.config,
+            )

    # `accelerate` 0.16.0 will have better support for customized saving
    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
        # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
-        if fsdp_stage != 0 or zero_stage == 3:
+        if fsdp_plugin is not None or zero_stage == 3:
            def save_model_hook(models, weights, output_dir):
                accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True)
                if accelerator.is_main_process:
                    from safetensors.torch import save_file

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

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

@@ -1160,34 +1265,23 @@ def main():
                        pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file)

            def load_model_hook(models, input_dir):
-                if args.use_ema and models[0] is generator_transformer3d:
-                    ema_path = os.path.join(input_dir, "transformer_ema")
-                    _, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True)
-                    load_model = WanTransformer3DModel_SelfForcing.from_pretrained(
-                        input_dir, subfolder="transformer_ema",
-                        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
-                        low_cpu_mem_usage=True,
-                    )
-                    load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config)
-                    load_model.load_state_dict(ema_kwargs)
-
-                    ema_generator_transformer3d.load_state_dict(load_model.state_dict())
-                    ema_generator_transformer3d.to(accelerator.device)
-                    del load_model
-
                for i in range(len(models)):
                    # pop models so that they are not loaded again
                    model = models.pop()

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

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

+        if args.use_ema:
+            def save_ema_hook(models, weights, output_dir):
+                if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process:
+                    ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
+
+            def load_ema_hook(models, input_dir):
+                ema_path = os.path.join(input_dir, "transformer_ema")
+                if isinstance(ema_generator_transformer3d, FSDPEMA):
+                    ema_generator_transformer3d.load_pretrained(ema_path)
+                else:
+                    _, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True)
+                    load_model = WanTransformer3DModel_SelfForcing.from_pretrained(
+                        input_dir, subfolder="transformer_ema",
+                        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+                        low_cpu_mem_usage=True,
+                        torch_dtype=torch.float32,
+                    )
+                    load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config)
+                    load_model.load_state_dict(ema_kwargs)
+                    ema_generator_transformer3d.load_state_dict(load_model.state_dict())
+                    ema_generator_transformer3d.to(accelerator.device)
+                    del load_model
+
+            accelerator.register_save_state_pre_hook(save_ema_hook)
+            accelerator.register_load_state_pre_hook(load_ema_hook)
+
    if args.gradient_checkpointing:
        generator_transformer3d.enable_gradient_checkpointing()
        fake_score_transformer3d.enable_gradient_checkpointing()
@@ -1267,8 +1387,8 @@ def main():

    fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters()))
    fake_trainable_params_optim = [
-        {'params': [], 'lr': args.learning_rate},
-        {'params': [], 'lr': args.learning_rate / 2},
+        {'params': [], 'lr': args.learning_rate_critic},
+        {'params': [], 'lr': args.learning_rate_critic / 2},
    ]
    in_already = []
    for name, param in fake_score_transformer3d.named_parameters():
@@ -1281,7 +1401,7 @@ def main():
                high_lr_flag = True
                fake_trainable_params_optim[0]['params'].append(param)
                if accelerator.is_main_process:
-                    print(f"Set {name} to lr : {args.learning_rate}")
+                    print(f"Set {name} to lr : {args.learning_rate_critic}")
                break
        if high_lr_flag:
            continue
@@ -1290,9 +1410,15 @@ def main():
                in_already.append(name)
                fake_trainable_params_optim[1]['params'].append(param)
                if accelerator.is_main_process:
-                    print(f"Set {name} to lr : {args.learning_rate / 2}")
+                    print(f"Set {name} to lr : {args.learning_rate_critic / 2}")
                break

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

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

        return length_to_frame_num

-    if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing:
+    if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan:
        aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
        batch_sampler_generator = torch.Generator().manual_seed(args.seed)
        batch_sampler = AspectRatioBatchImageVideoSampler(
@@ -1672,7 +1798,7 @@ def main():
    )
    fake_score_lr_scheduler = get_scheduler(
        args.lr_scheduler,
-        optimizer=optimizer,
+        optimizer=critic_optimizer,
        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
        num_training_steps=args.max_train_steps * accelerator.num_processes,
    )
@@ -1798,6 +1924,8 @@ def main():
        vae_stream_2 = None

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

    def randomize_denoising_step_indices(
        denoising_step_indices_list,
@@ -1852,12 +1980,16 @@ def main():
        return result

    for epoch in range(first_epoch, args.num_train_epochs):
+        if global_step >= args.max_train_steps:
+            break
        train_dmd_loss = 0.0
        # Number of generator backward contributions since the last log flush; the
        # generator only backprops every gen_update_interval batches, so its metrics
        # must be averaged by contribution count, not by gradient_accumulation_steps.
        train_gen_log_count = 0
        train_denoising_loss = 0.0
+        train_gan_g_loss = 0.0
+        train_gan_d_loss = 0.0
        batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
        for step, batch in enumerate(train_dataloader):
            # Data batch sanity check
@@ -1878,7 +2010,7 @@ def main():
                    save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.mp4", rescale=True)

            with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
-                if args.train_mode != "normal" or args.use_teacher_forcing:
+                if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
                    # Convert images to latent space
                    pixel_values = batch["pixel_values"].to(weight_dtype)

@@ -1942,8 +2074,9 @@ def main():
                        temp_n_frames = (temp_n_frames - 1) // sample_n_frames_bucket_interval + 1

                        pixel_values = pixel_values[:, :temp_n_frames, :, :]
-                        mask_pixel_values = mask_pixel_values[:, :temp_n_frames, :, :]
-                        mask = mask[:, :temp_n_frames, :, :]
+                        if args.train_mode != "normal":
+                            mask_pixel_values = mask_pixel_values[:, :temp_n_frames, :, :]
+                            mask = mask[:, :temp_n_frames, :, :]
                        
                    # Keep all node same token length to accelerate the traning when resolution grows.
                    if args.keep_all_node_same_token_length:
@@ -1966,8 +2099,9 @@ def main():
                        actual_video_length = (actual_video_length - 1) // sample_n_frames_bucket_interval + 1

                        pixel_values = pixel_values[:, :actual_video_length, :, :]
-                        mask_pixel_values = mask_pixel_values[:, :actual_video_length, :, :]
-                        mask = mask[:, :actual_video_length, :, :]
+                        if args.train_mode != "normal":
+                            mask_pixel_values = mask_pixel_values[:, :actual_video_length, :, :]
+                            mask = mask[:, :actual_video_length, :, :]

                    if args.low_vram:
                        torch.cuda.empty_cache()
@@ -1990,7 +2124,7 @@ def main():
                                pixel_values_bs = pixel_values_bs.sample()
                                new_pixel_values.append(pixel_values_bs)
                            return torch.cat(new_pixel_values, dim = 0)
-                        if args.use_teacher_forcing:
+                        if args.use_teacher_forcing or use_gan:
                            clean_latents = _batch_encode_vae(pixel_values)
                        else:
                            clean_latents = None
@@ -2022,7 +2156,7 @@ def main():
                                clip_context.append(_clip_context)
                            clip_context = torch.cat(clip_context)

-                    if args.use_teacher_forcing:
+                    if args.use_teacher_forcing or use_gan:
                        target_shape = clean_latents.size()
                    else:
                        target_shape = mask_latents.size()
@@ -2426,17 +2560,7 @@ def main():
                        # Final output — slice generated frames (may be < num_frames for variable-length)
                        generator_pred_full = output_pred[:, :, :num_generated_frames]

-                        # Gradient mask: first block gets no gradient when generating > min frames
-                        # (matches Self-Forcing reference: model/base.py L182-L190)
-                        min_num_frames_score = args.score_num_frames
-                        need_gradient_mask = (num_generated_frames != min_num_frames_score)
-                        gradient_mask = None
-                        if need_gradient_mask:
-                            gradient_mask = torch.ones_like(generator_pred_full, dtype=torch.bool)
-                            if args.independent_first_frame:
-                                gradient_mask[:, :, :1] = False
-                            else:
-                                gradient_mask[:, :, :args.num_frame_per_block] = False
+                        need_gradient_mask = (num_generated_frames != args.score_num_frames)

                        # Slice for score computation: last score_num_frames frames
                        if num_generated_frames > args.score_num_frames:
@@ -2450,11 +2574,13 @@ def main():
                            generator_pred_for_score = generator_pred_full
                            score_num_frames = num_generated_frames

-                        # Compute score_mask for DMD loss (matches Self-Forcing: dmd.py L199-204)
                        score_mask = None
-                        if gradient_mask is not None:
-                            mask_offset = num_generated_frames - score_num_frames
-                            score_mask = gradient_mask[:, :, mask_offset:mask_offset + score_num_frames]
+                        if need_gradient_mask:
+                            score_mask = torch.ones_like(generator_pred_for_score, dtype=torch.bool)
+                            if args.independent_first_frame:
+                                score_mask[:, :, :1] = False
+                            else:
+                                score_mask[:, :, :args.num_frame_per_block] = False

                        # generator_pred = the sliced version for DMD loss
                        generator_pred = generator_pred_for_score
@@ -2530,7 +2656,7 @@ def main():
                                    generator_noise_list = [generator_noise[i] for i in range(bsz)]
                                    clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None
                                    
-                                    # Use block_mask for causal training (一次性处理整个视频)
+                                    # Use block_mask for causal training
                                    generator_pred = generator_transformer3d(
                                        x=generator_noise_list,
                                        context=prompt_embeds,
@@ -2563,6 +2689,8 @@ def main():
                    # Common code for both KV cache and block mask training
                    indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
                    generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
+                    if args.score_timestep_clamp:
+                        generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep)
                    generator_denoised_input = add_noise(
                        generator_pred,
                        torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
@@ -2668,17 +2796,53 @@ def main():
                    train_dmd_loss += avg_dmd_loss.item()
                    train_gen_log_count += 1

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

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

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

                fake_score_denoised_input = add_noise(
@@ -2975,15 +3141,25 @@ def main():
                    critic_timestep
                )

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

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

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

                if global_step % args.checkpointing_steps == 0:
                    if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
@@ -3092,6 +3398,9 @@ def main():
            if global_step >= args.max_train_steps:
                break

+        if global_step >= args.max_train_steps:
+            break
+
        if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
            with progress_bar.paused():
                if args.use_ema:
@@ -3121,12 +3430,6 @@ def main():

    # Create the pipeline using the trained modules and save it.
    accelerator.wait_for_everyone()
-    if args.use_ema:
-        # Copy EMA weights to generator on ALL processes before save_state (collective op).
-        ema_generator_transformer3d.copy_to(generator_transformer3d.parameters())
-    if accelerator.is_main_process:
-        generator_transformer3d = unwrap_model(generator_transformer3d)
-
    if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
        gc.collect()
        torch.cuda.empty_cache()
@@ -3135,8 +3438,6 @@ def main():
        fake_score_save_path = os.path.join(save_path, "fake_score")
        accelerator.save_state(save_path)
        accelerator_fake_score_transformer3d.save_state(fake_score_save_path)
-        if args.use_ema and accelerator.is_main_process:
-            ema_generator_transformer3d.save_pretrained(os.path.join(save_path, "transformer_ema"))
        logger.info(f"Saved state to {save_path}")

    accelerator.end_training()
@@ -258,6 +258,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | `--adam_epsilon` | AdamW epsilon | 1e-10 |
 | `--vae_mini_batch` | VAE encoding mini-batch size | 1 |
 | `--max_grad_norm` | Gradient clipping threshold | 0.05 |
+| `--fsdp_reduce_fp32` | Optional: reduce gradients in fp32 under FSDP mixed precision (default reduces in bf16) | - |
 | `--enable_bucket` | Enable bucket training, no cropping, group by resolution | - |
 | `--random_hw_adapt` | Auto-scale images/videos to random sizes in `[min_size, max_size]` range | - |
 | `--training_with_video_token_length` | Train based on token length, supports arbitrary resolutions | - |
@@ -275,12 +276,17 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | Parameter | Description | Example Value |
 |-----|------|-------|
 | `--denoising_step_indices_list` | Denoising step indices list (core distillation parameter) | `1000 750 500 250` |
+| `--randomize_step_indices` | Randomize the denoising timestep indices during training | - |
 | `--real_guidance_scale` | Real guidance scale for scoring | 6.0 |
+| `--score_timestep_clamp` | Optional: clamp DMD/critic timesteps to `[0.02, 0.98] * num_train_timesteps` | - |
 | `--fake_guidance_scale` | Fake guidance scale for scoring | 0.0 |
 | `--gen_update_interval` | Generator update interval | 5 |
 | `--negative_prompt` | Negative prompt for distillation | Chinese negative prompt |
 | `--train_sampling_steps` | Training sampling steps | 1000 |
 | `--ode_transformer_path` | Path to ODE-trained weights to load into generator transformer3d | `models/Diffusion_Transformer/Self-Forcing/checkpoints/ode_init.pt` |
+| `--real_score_model_name_or_path` | Optional: separate model used as real score (defaults to `pretrained_model_name_or_path`) | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
+| `--score_num_frames` | Latent frames in the score-computation window; `fake_score`/`real_score` always receive this many frames | 21 |
+| `--min_length_prob_bias` | Probability bias toward sampling the minimum length (`score_num_frames`); `0.0` = uniform sampling | 0.0 |

 **Self-Forcing-Specific Parameters**:

@@ -294,6 +300,23 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | `--use_teacher_forcing` | Enable teacher forcing training (pass clean_x to transformer) | - |
 | `--teacher_forcing_prob` | Probability of applying teacher forcing per step | 1.0 |

+**GAN / Discriminator Parameters (One-Forcing)**:
+
+> All default to `0` / off, so a plain DMD run is unaffected. Setting **any** positive `--gan_g_weight` / `--gan_d_weight` / `--r1_weight` / `--r2_weight` attaches a discriminator head to `fake_score` and **requires real videos** in the training data.
+
+| Parameter | Description | Example Value |
+|-----|------|-------|
+| `--gan_g_weight` | One-Forcing: weight of the non-saturating GAN loss added to the DMD generator loss (`0` disables) | 0.0 |
+| `--gan_d_weight` | One-Forcing: weight of the GAN discriminator loss added to the `fake_score` denoising loss (`0` disables) | 0.0 |
+| `--gan_feature_layers` | `fake_score` block indices whose outputs feed the discriminator head (one register token per layer) | `21 29` |
+| `--gan_block_ffn_dim` | FFN dim of the discriminator cross-attention blocks | 2048 |
+| `--gan_block_num_heads` | Number of heads of the discriminator cross-attention blocks | 12 |
+| `--relativistic_discriminator` | Use the relativistic GAN loss (fake vs. real logit difference) | - |
+| `--r1_weight` | Finite-difference R1 penalty on noisy real latents (APT form, not divided by sigma) | 0.0 |
+| `--r1_sigma` | Perturbation std of the R1 penalty | 0.01 |
+| `--r2_weight` | Finite-difference R2 penalty on noisy generated latents | 0.0 |
+| `--r2_sigma` | Perturbation std of the R2 penalty | 0.01 |
+
 **Sample Size Configuration Guide**:
 - `video_sample_size` represents the resolution size of videos; when `random_hw_adapt` is True, it represents the minimum value between video and image resolutions.
 - `image_sample_size` represents the resolution size of images; when `random_hw_adapt` is True, it represents the maximum value between video and image resolutions.
@@ -259,6 +259,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | `--adam_epsilon` | AdamW epsilon 值 | 1e-10 |
 | `--vae_mini_batch` | VAE 编码时的迷你批次大小 | 1 |
 | `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
+| `--fsdp_reduce_fp32` | 可选：FSDP 混精下梯度以 fp32 做跨卡 reduce（默认 bf16） | - |
 | `--enable_bucket` | 启用分桶训练，不裁剪图片/视频，按分辨率分组训练 | - |
 | `--random_hw_adapt` | 自动缩放图片/视频到 `[min_size, max_size]` 范围内的随机尺寸 | - |
 | `--training_with_video_token_length` | 根据 token 长度训练，支持任意分辨率 | - |
@@ -276,12 +277,17 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | 参数 | 说明 | 示例值 |
 |-----|------|-------|
 | `--denoising_step_indices_list` | 去噪步骤列表（蒸馏核心参数） | `1000 750 500 250` |
+| `--randomize_step_indices` | 训练时随机化去噪 timestep 索引 | - |
 | `--real_guidance_scale` | 用于评分的真实 guidance scale | 6.0 |
+| `--score_timestep_clamp` | 可选：将 DMD/critic 的 timestep 截断到 `[0.02, 0.98] * num_train_timesteps` | - |
 | `--fake_guidance_scale` | 用于评分的虚拟 guidance scale | 0.0 |
 | `--gen_update_interval` | 生成器更新间隔 | 5 |
 | `--negative_prompt` | 用于蒸馏的负向提示词 | 中文负向提示词 |
 | `--train_sampling_steps` | 训练采样步数 | 1000 |
 | `--ode_transformer_path` | ODE 训练权重路径，加载到 generator transformer3d 中 | `models/Diffusion_Transformer/Self-Forcing/checkpoints/ode_init.pt` |
+| `--real_score_model_name_or_path` | 可选：单独指定 real score 模型（默认与 `pretrained_model_name_or_path` 相同） | `models/Diffusion_Transformer/Wan2.1-T2V-14B` |
+| `--score_num_frames` | 评分窗口的 latent 帧数；`fake_score`/`real_score` 始终接收这么多帧 | 21 |
+| `--min_length_prob_bias` | 采样最小长度（`score_num_frames`）的概率偏置；`0.0` 表示均匀采样 | 0.0 |

 **Self-Forcing 特有参数**：

@@ -295,6 +301,23 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | `--use_teacher_forcing` | 启用 teacher forcing 训练（将 clean_x 传给 transformer） | - |
 | `--teacher_forcing_prob` | 每步应用 teacher forcing 的概率 | 1.0 |

+**GAN / 判别器参数（One-Forcing）**：
+
+> 全部默认为 `0` / 关闭，因此普通 DMD 训练不受影响。只要 `--gan_g_weight` / `--gan_d_weight` / `--r1_weight` / `--r2_weight` 中**任一**为正值，就会在 `fake_score` 上挂一个判别器头，并且训练数据**必须包含真实视频**。
+
+| 参数 | 说明 | 示例值 |
+|-----|------|-------|
+| `--gan_g_weight` | One-Forcing：加到 DMD 生成器损失上的非饱和 GAN 损失权重（`0` 关闭） | 0.0 |
+| `--gan_d_weight` | One-Forcing：加到 `fake_score` 去噪损失上的 GAN 判别器损失权重（`0` 关闭） | 0.0 |
+| `--gan_feature_layers` | 输出喂给判别器头的 `fake_score` block 索引（每层一个 register token） | `21 29` |
+| `--gan_block_ffn_dim` | 判别器 cross-attention block 的 FFN 维度 | 2048 |
+| `--gan_block_num_heads` | 判别器 cross-attention block 的头数 | 12 |
+| `--relativistic_discriminator` | 使用 relativistic GAN 损失（fake 与 real logit 之差） | - |
+| `--r1_weight` | 对含噪真实 latent 的有限差分 R1 惩罚（APT 形式，不除以 sigma） | 0.0 |
+| `--r1_sigma` | R1 惩罚的扰动标准差 | 0.01 |
+| `--r2_weight` | 对含噪生成 latent 的有限差分 R2 惩罚 | 0.0 |
+| `--r2_sigma` | R2 惩罚的扰动标准差 | 0.01 |
+
 **Sample Size 配置指南**：
 - `video_sample_size` 表示视频的分辨率大小；当 `random_hw_adapt` 为 True 时，表示视频和图像分辨率的最小值。
 - `image_sample_size` 表示图像的分辨率大小；当 `random_hw_adapt` 为 True 时，表示视频和图像分辨率的最大值。
@@ -277,6 +277,9 @@ bash scripts/wan2.1_self_forcing/train_ode.sh
 | `--max_grad_norm` | Gradient clipping threshold | 0.05 |
 | `--trainable_modules` | Trainable modules (`"."` = all) | `"."` |
 | `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` |
+| `--save_state` | Save the full accelerator state on checkpoint (required to resume training) | - |
+| `--use_kv_cache_training` | Run block-by-block KV cache training that matches `pipeline_wan_self_forcing` inference; otherwise fall back to one-shot causal-mask ODE regression (baseline) | - |
+| `--prob_full_zero_start` | Per-sample probability of forcing all frames in all blocks to timestep index `0` (pure-noise start) to bridge the train-inference gap; `0.0` disables | 0.0 |

 **ODE-specific parameters** (must match Step 1 unless you understand the consequences):

@@ -277,6 +277,9 @@ bash scripts/wan2.1_self_forcing/train_ode.sh
 | `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
 | `--trainable_modules` | 可训练模块（`"."` 表示全量） | `"."` |
 | `--resume_from_checkpoint` | 恢复训练路径或 `"latest"` | `latest` |
+| `--save_state` | 在 checkpoint 中保存完整的 accelerator 状态（恢复训练所需） | - |
+| `--use_kv_cache_training` | 运行逐块 KV 缓存训练，完全匹配 `pipeline_wan_self_forcing` 推理行为；否则回退到一次性的 causal-mask ODE 回归（baseline） | - |
+| `--prob_full_zero_start` | 每个样本强制所有块的所有帧使用 timestep 索引 `0`（纯噪声起点）的概率，用于弥合训练-推理差距；`0.0` 表示关闭 | 0.0 |

 **ODE 特有参数**（除非清楚后果，否则需与第一步保持一致）：

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

 import argparse
 import contextlib
+import dataclasses
 import gc
 import json
 import logging
@@ -74,10 +75,12 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512,
                             get_random_mask)
 from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
                               WanTransformer3DModel,
-                               WanTransformer3DModel_SelfForcing)
+                               WanTransformer3DModel_SelfForcing,
+                               add_wan_gan_discriminator)
 from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
                                 WanSelfForcingPipeline)
 from videox_fun.utils.discrete_sampler import DiscreteSampling
+from videox_fun.utils.fsdp_ema import FSDPEMA
 from videox_fun.utils.tqdm_bar import PauseAwareTqdm
 from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
                                    get_image_to_video_latent,
@@ -165,6 +168,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi
    return pred_for_score, score_num_frames, True


+def repeat_batch(value, repeats):
+    """Repeat a batched tensor or per-sample list along the batch dimension."""
+    if value is None or repeats == 1:
+        return value
+    if torch.is_tensor(value):
+        return torch.cat([value] * repeats)
+    return list(value) * repeats
+
+
 def filter_kwargs(cls, kwargs):
    import inspect
    sig = inspect.signature(cls.__init__)
@@ -492,6 +504,8 @@ def parse_args():
    parser.add_argument(
        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
    )
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.")
+    parser.add_argument("--ema_decay", type=float, default=0.9999, help="EMA decay applied per generator update (Causal-Forcing uses 0.99).")
    parser.add_argument(
        "--allow_tf32",
        action="store_true",
@@ -523,6 +537,9 @@ def parse_args():
    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument(
+        "--fsdp_reduce_fp32", action="store_true", help="Reduce gradients in fp32 under FSDP mixed precision (params are still gathered in low precision)."
+    )
    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
    parser.add_argument(
@@ -722,6 +739,12 @@ def parse_args():
        default=None,
        help=("If you want to load the ode-trained weight into generator transformer3d, input its path."),
    )
+    parser.add_argument(
+        "--real_score_model_name_or_path",
+        type=str,
+        default=None,
+        help=("Path to a separate model (e.g. Wan2.1-T2V-14B) used as real score. Defaults to pretrained_model_name_or_path."),
+    )
    parser.add_argument(
        "--vae_path",
        type=str,
@@ -782,6 +805,11 @@ def parse_args():
        default=4.5,
        help="The cfg scale for real score.",
    )
+    parser.add_argument(
+        "--score_timestep_clamp",
+        action="store_true",
+        help="Clamp DMD and critic timesteps to [0.02, 0.98] * num_train_timesteps.",
+    )
    parser.add_argument(
        '--denoising_step_indices_list', 
        nargs='+', 
@@ -842,6 +870,67 @@ def parse_args():
        default=1.0,
        help="Probability of applying teacher forcing per step (1.0 = always)"
    )
+    parser.add_argument(
+        "--gan_g_weight",
+        type=float,
+        default=0.0,
+        help="One-Forcing: weight of the non-saturating GAN loss added to the DMD generator loss (0 disables). "
+             "Any positive GAN/R1/R2 weight attaches a discriminator head to fake_score and requires real videos."
+    )
+    parser.add_argument(
+        "--gan_d_weight",
+        type=float,
+        default=0.0,
+        help="One-Forcing: weight of the GAN discriminator loss added to the fake_score denoising loss (0 disables)."
+    )
+    parser.add_argument(
+        "--gan_feature_layers",
+        nargs="+",
+        type=int,
+        default=[21, 29],
+        help="fake_score block indices whose outputs feed the discriminator head (one register token per layer)."
+    )
+    parser.add_argument(
+        "--gan_block_ffn_dim",
+        type=int,
+        default=2048,
+        help="FFN dim of the discriminator cross-attention blocks."
+    )
+    parser.add_argument(
+        "--gan_block_num_heads",
+        type=int,
+        default=12,
+        help="Number of heads of the discriminator cross-attention blocks."
+    )
+    parser.add_argument(
+        "--relativistic_discriminator",
+        action="store_true",
+        help="Use the relativistic GAN loss (fake vs. real logit difference)."
+    )
+    parser.add_argument(
+        "--r1_weight",
+        type=float,
+        default=0.0,
+        help="Finite-difference R1 penalty on noisy real latents (APT form, not divided by sigma)."
+    )
+    parser.add_argument(
+        "--r1_sigma",
+        type=float,
+        default=0.01,
+        help="Perturbation std of the R1 penalty."
+    )
+    parser.add_argument(
+        "--r2_weight",
+        type=float,
+        default=0.0,
+        help="Finite-difference R2 penalty on noisy generated latents."
+    )
+    parser.add_argument(
+        "--r2_sigma",
+        type=float,
+        default=0.01,
+        help="Perturbation std of the R2 penalty."
+    )

    args = parser.parse_args()
    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
@@ -914,6 +1003,9 @@ def main():
        else:
            fsdp_stage = 0
        print(f"Using FSDP stage: {fsdp_stage}")
+        if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
+            fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32)
+            print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")

        args.use_fsdp = True
        if fsdp_stage == 3:
@@ -970,6 +1062,9 @@ def main():
        args.mixed_precision = accelerator.mixed_precision

    args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list]
+    use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight))
+    if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1):
+        raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.")
    # Load scheduler, tokenizer and models.
    noise_scheduler = FlowMatchEulerDiscreteScheduler(
        **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
@@ -1030,7 +1125,7 @@ def main():
        low_cpu_mem_usage=True,
    ).to(weight_dtype)
    real_score_transformer3d = WanTransformer3DModel.from_pretrained(
-        os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+        os.path.join(args.real_score_model_name_or_path or args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
        low_cpu_mem_usage=True,
    ).to(weight_dtype)
@@ -1062,7 +1157,8 @@ def main():
            state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()}

        m, u = generator_transformer3d.load_state_dict(state_dict, strict=False)
-        m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False)
+        if args.real_score_model_name_or_path is None:
+            m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False)
        m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False)
        print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
        assert len(u) == 0
@@ -1116,6 +1212,33 @@ def main():
            if trainable_module_name in name:
                param.requires_grad = True
                break
+    if use_gan:
+        add_wan_gan_discriminator(
+            fake_score_transformer3d, args.gan_feature_layers, args.gan_block_ffn_dim, args.gan_block_num_heads,
+            weight_dtype, seed=args.seed or 0,
+        )
+
+    # Create EMA for the generator_transformer3d.
+    if args.use_ema:
+        if zero_stage == 3:
+            raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.")
+
+        if fsdp_plugin is not None:
+            if getattr(fsdp_plugin, "fsdp_version", 1) != 1:
+                raise NotImplementedError("FSDPEMA only supports FSDP1.")
+            ema_module = WanTransformer3DModel_SelfForcing.from_config(generator_transformer3d.config)
+            ema_generator_transformer3d = FSDPEMA(
+                ema_module, source=generator_transformer3d,
+                accelerator=accelerator, fsdp_plugin=fsdp_plugin,
+                decay=args.ema_decay,
+            )
+        else:
+            ema_generator_transformer3d = EMAModel(
+                generator_transformer3d.parameters(),
+                decay=args.ema_decay,
+                model_cls=WanTransformer3DModel_SelfForcing,
+                model_config=generator_transformer3d.config,
+            )

    # `accelerate` 0.16.0 will have better support for customized saving
    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
@@ -1178,6 +1301,32 @@ def main():
        accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
        accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)

+        if args.use_ema:
+            def save_ema_hook(models, weights, output_dir):
+                if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process:
+                    ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema"))
+
+            def load_ema_hook(models, input_dir):
+                ema_path = os.path.join(input_dir, "transformer_ema")
+                if isinstance(ema_generator_transformer3d, FSDPEMA):
+                    ema_generator_transformer3d.load_pretrained(ema_path)
+                else:
+                    _, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True)
+                    load_model = WanTransformer3DModel_SelfForcing.from_pretrained(
+                        input_dir, subfolder="transformer_ema",
+                        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+                        low_cpu_mem_usage=True,
+                        torch_dtype=torch.float32,
+                    )
+                    load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config)
+                    load_model.load_state_dict(ema_kwargs)
+                    ema_generator_transformer3d.load_state_dict(load_model.state_dict())
+                    ema_generator_transformer3d.to(accelerator.device)
+                    del load_model
+
+            accelerator.register_save_state_pre_hook(save_ema_hook)
+            accelerator.register_load_state_pre_hook(load_ema_hook)
+
    if args.gradient_checkpointing:
        generator_transformer3d.enable_gradient_checkpointing()
        fake_score_transformer3d.enable_gradient_checkpointing()
@@ -1269,6 +1418,11 @@ def main():
                if accelerator.is_main_process:
                    print(f"Set {name} to lr : {args.learning_rate_critic / 2}")
                break
+    if use_gan:
+        fake_trainable_params_optim[0]['params'].extend(
+            param for name, param in fake_score_transformer3d.named_parameters()
+            if name.startswith("gan_head.") and name not in in_already
+        )

    if args.use_came:
        optimizer = optimizer_cls(
@@ -1311,7 +1465,7 @@ def main():
        args.random_hw_adapt = False

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

        return length_to_frame_num

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

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

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

    def randomize_denoising_step_indices(
        denoising_step_indices_list,
@@ -1832,6 +1991,8 @@ def main():
        # must be averaged by contribution count, not by gradient_accumulation_steps.
        train_gen_log_count = 0
        train_denoising_loss = 0.0
+        train_gan_g_loss = 0.0
+        train_gan_d_loss = 0.0
        batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
        for step, batch in enumerate(train_dataloader):
            # Data batch sanity check
@@ -1852,7 +2013,7 @@ def main():
                    save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.mp4", rescale=True)

            with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
-                if args.train_mode != "normal" or args.use_teacher_forcing:
+                if args.train_mode != "normal" or args.use_teacher_forcing or use_gan:
                    # Convert images to latent space
                    pixel_values = batch["pixel_values"].to(weight_dtype)

@@ -1964,7 +2125,7 @@ def main():
                                pixel_values_bs = pixel_values_bs.sample()
                                new_pixel_values.append(pixel_values_bs)
                            return torch.cat(new_pixel_values, dim = 0)
-                        if args.use_teacher_forcing:
+                        if args.use_teacher_forcing or use_gan:
                            clean_latents = _batch_encode_vae(pixel_values)
                        else:
                            clean_latents = None
@@ -1996,7 +2157,7 @@ def main():
                                clip_context.append(_clip_context)
                            clip_context = torch.cat(clip_context)

-                    if args.use_teacher_forcing:
+                    if args.use_teacher_forcing or use_gan:
                        target_shape = clean_latents.size()
                    else:
                        target_shape = mask_latents.size()
@@ -2504,7 +2665,7 @@ def main():
                                    generator_noise_list = [generator_noise[i] for i in range(bsz)]
                                    clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None
                                    
-                                    # Use block_mask for causal training (一次性处理整个视频)
+                                    # Use block_mask for causal training (process the whole video at once)
                                    generator_pred = generator_transformer3d(
                                        x=generator_noise_list,
                                        context=prompt_embeds,
@@ -2537,6 +2698,8 @@ def main():
                    # Common code for both KV cache and block mask training
                    indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
                    generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
+                    if args.score_timestep_clamp:
+                        generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep)
                    generator_denoised_input = add_noise(
                        generator_pred,
                        torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
@@ -2642,17 +2805,53 @@ def main():
                    train_dmd_loss += avg_dmd_loss.item()
                    train_gen_log_count += 1

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

-                    accelerator.backward(dmd_loss)
+                    accelerator.backward(generator_loss)
                    if accelerator.sync_gradients:
                        accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()
+                    if args.gan_g_weight > 0:
+                        # Drop the fake_score grads produced by the generator GAN loss
+                        critic_optimizer.zero_grad()
                    
                    if args.low_vram:
                        fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device)
@@ -2941,6 +3140,8 @@ def main():

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

                fake_score_denoised_input = add_noise(
@@ -2950,14 +3151,41 @@ def main():
                )

                with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
-                    fake_score_denoised_output = fake_score_transformer3d(
-                        x=fake_score_denoised_input,
-                        context=prompt_embeds,
-                        t=critic_timestep,
-                        seq_len=seq_len,
-                        y=inpaint_latents if args.train_mode != "normal" else None,
-                        clip_fea=clip_context if args.train_mode != "normal" else None,
-                    )
+                    if use_gan:
+                        # Generated and real latents share critic_noise / critic_timestep. The critic forward is
+                        # split into three passes so the peak activation stays at 2x batch instead of 4x:
+                        # A) fake + real for the denoising and GAN loss, B) real + R1 perturbation, C) fake + R2.
+                        real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype)
+                        real_noisy = add_noise(real_latents, critic_noise, critic_timestep)
+                        if args.r1_weight > 0:
+                            r1_perturbation = args.r1_sigma * torch.randn(
+                                real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng
+                            )
+                        if args.r2_weight > 0:
+                            r2_perturbation = args.r2_sigma * torch.randn(
+                                fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype,
+                                device=fake_score_denoised_input.device, generator=torch_rng
+                            )
+                        fake_score_denoised_output, gan_logits = fake_score_transformer3d(
+                            x=torch.cat([fake_score_denoised_input, real_noisy]),
+                            context=repeat_batch(prompt_embeds, 2),
+                            t=critic_timestep.repeat(2),
+                            seq_len=seq_len,
+                            y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
+                            clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                        fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0]
+                        gan_logits = gan_logits.float().chunk(2)
+                    else:
+                        fake_score_denoised_output = fake_score_transformer3d(
+                            x=fake_score_denoised_input,
+                            context=prompt_embeds,
+                            t=critic_timestep,
+                            seq_len=seq_len,
+                            y=inpaint_latents if args.train_mode != "normal" else None,
+                            clip_fea=clip_context if args.train_mode != "normal" else None,
+                        )

                def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
                    noise_pred = noise_pred.float()
@@ -2974,8 +3202,53 @@ def main():
                denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
                avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
                train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
-                
-                accelerator_fake_score_transformer3d.backward(denoising_loss)
+
+                critic_loss = denoising_loss
+                if use_gan:
+                    fake_logit, real_logit = gan_logits[0], gan_logits[1]
+                    if args.relativistic_discriminator:
+                        gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
+                    else:
+                        gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
+                    gan_d_loss = gan_d_loss * args.gan_d_weight
+                    train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
+                    critic_loss = critic_loss + gan_d_loss
+
+                accelerator_fake_score_transformer3d.backward(critic_loss)
+                del critic_loss
+
+                if use_gan and args.r1_weight > 0:
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        _, r1_logits = fake_score_transformer3d(
+                            x=torch.cat([real_noisy + r1_perturbation, real_noisy]),
+                            context=repeat_batch(prompt_embeds, 2),
+                            t=critic_timestep.repeat(2),
+                            seq_len=seq_len,
+                            y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
+                            clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                        r1_logits = r1_logits.float().chunk(2)
+                        r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean()
+                    accelerator_fake_score_transformer3d.backward(r1_loss)
+                    del r1_loss
+
+                if use_gan and args.r2_weight > 0:
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
+                        _, r2_logits = fake_score_transformer3d(
+                            x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]),
+                            context=repeat_batch(prompt_embeds, 2),
+                            t=critic_timestep.repeat(2),
+                            seq_len=seq_len,
+                            y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
+                            clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
+                            gan_classify=True,
+                        )
+                        r2_logits = r2_logits.float().chunk(2)
+                        r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean()
+                    accelerator_fake_score_transformer3d.backward(r2_loss)
+                    del r2_loss
+
                if accelerator_fake_score_transformer3d.sync_gradients:
                    accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm)
                critic_optimizer.step()
@@ -2988,13 +3261,22 @@ def main():
                    
            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
+                # EMA tracks the generator only on rounds where it actually ran optimizer.step();
+                # with gen_update_interval > 1 the critic-only rounds leave the generator unchanged.
+                if args.use_ema and generator_update:
+                    ema_generator_transformer3d.step(generator_transformer3d.parameters())

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

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

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

            logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
+            if use_gan:
+                logs["gan_d_loss"] = gan_d_loss.detach().item()
+                if args.gan_g_weight > 0:
+                    logs["gan_g_loss"] = gan_g_loss.detach().item()
            progress_bar.set_postfix(**logs)

            if global_step >= args.max_train_steps:
@@ -3059,6 +3352,10 @@ def main():

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

    # Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever
    # something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto
@@ -15,6 +15,7 @@
 # See the License for the specific language governing permissions and

 import argparse
+import dataclasses
 import gc
 import logging
 import math
@@ -362,6 +363,10 @@ def parse_args():
    parser.add_argument("--adam_weight_decay", type=float, default=3e-2, help="Weight decay to use.")
    parser.add_argument("--adam_epsilon", type=float, default=1e-10, help="Epsilon value for the Adam optimizer.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument(
+        "--fsdp_reduce_fp32", action="store_true",
+        help="Reduce gradients in fp32 under FSDP mixed precision (params are still gathered in low precision).",
+    )
    parser.add_argument(
        "--mixed_precision",
        type=str,
@@ -581,6 +586,11 @@ def main():
        else:
            fsdp_stage = 0
        print(f"Using FSDP stage: {fsdp_stage}")
+        if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None:
+            fsdp_plugin.mixed_precision_policy = dataclasses.replace(
+                fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32
+            )
+            print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}")

        args.use_fsdp = True
        if fsdp_stage == 3: