2.08× faster Self-Forcing DMD training iterations on Hopper GPUs, optimized by AI agents.
| Implementation | Time / training iteration (ms) ↓ | Speedup ↑ |
|---|---|---|
| Baseline | 10,811.0 | 1.00× |
| Faster-SelfForcing | 5,191.9 | 2.08× |
Measured in a short benchmark on 2 H200-class GPUs (Hopper), one video per GPU as in the upstream recipe, BF16 mixed precision. The timer runs the repository's real train.py; baseline and optimized runs alternate in independent processes.
Prepare the environment, public weights and prompts. The measured configuration is a set of environment switches on the upstream trainer; source them and launch as upstream does (64 GPUs, one video per GPU):
source scripts/configs/faster.env
torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
--rdzv_backend=c10d --rdzv_endpoint $MASTER_ADDR \
train.py --config_path configs/self_forcing_dmd.yaml \
--logdir logs/faster_self_forcing_dmd --disable-wandbTested with PyTorch 2.7.1 + CUDA 12.8 on Hopper GPUs. Attention falls back to FA2 when FA3 is unavailable or the input is unsupported.
Measured on two H200-class GPUs, one video per GPU.
- Attention: FlashAttention-3 installed from the official repository replaces FA2 in the dense attention calls. It is enabled by default when available for supported inputs;
SF_ATTN=fa2restores the upstream path. - Rollout and FSDP: the generator is all-gathered once per rollout, cross-attention K/V is computed once, the no-grad KV-cache forwards replay as CUDA graphs, and host synchronizations are removed.
- Fused kernels: Triton kernels for RMSNorm, LayerNorm + modulate, QK-norm + RoPE and the modulate / gated-residual chains.
- Memory and overhead: selective activation checkpointing, frozen teacher and text encoder held in BF16, a resident text encoder, and a fast asynchronous EMA.
In the recorded tests with FA2 attention, each of the 17 optimization settings matched the baseline bit for bit on the recorded tensors. The final configuration also enabled FA3 and passed the short-trajectory statistical check described below. See the switch table and the measured configuration.
The recorded FA3 configuration passed a statistical comparison against the baseline using three random seeds and 100 training steps per seed. Critic and generator losses and gradient norms were within the preset thresholds.
| Training metric | Mean difference / baseline standard deviation | Within threshold |
|---|---|---|
| Critic loss | −0.145 | ✓ |
| Critic gradient norm | +0.073 | ✓ |
| Generator loss | −0.039 | ✓ |
| Generator gradient norm | +0.105 | ✓ |
See the validation method for details. Full 600-iteration training and generated-video quality have not yet been evaluated; the 2.08× result measures training-iteration speed.
Based on guandeh17/Self-Forcing at 33593df, by Xun Huang, Zhengqi Li, Guande He, Mingyuan Zhou and Eli Shechtman. Paper · Project page · Apache-2.0 license. The original README and authorship are retained below.
Original Self Forcing README: installation, inference and training
Xun Huang1
·
Zhengqi Li1
·
Guande He2
·
Mingyuan Zhou2
·
Eli Shechtman1
1Adobe Research 2UT Austin
Self Forcing trains autoregressive video diffusion models by simulating the inference process during training, performing autoregressive rollout with KV caching. It resolves the train-test distribution mismatch and enables real-time, streaming video generation on a single RTX 4090 while matching the quality of state-of-the-art diffusion models.
demo.mp4
We tested this repo on the following setup:
- Nvidia GPU with at least 24 GB memory (RTX 4090, A100, and H100 are tested).
- Linux operating system.
- 64 GB RAM.
Other hardware setup could also work but hasn't been tested.
Create a conda environment and install dependencies:
conda create -n self_forcing python=3.10 -y
conda activate self_forcing
pip install -r requirements.txt
pip install flash-attn --no-build-isolation
python setup.py develop
huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B --local-dir-use-symlinks False --local-dir wan_models/Wan2.1-T2V-1.3B
huggingface-cli download gdhe17/Self-Forcing checkpoints/self_forcing_dmd.pt --local-dir .
python demo.py
Note:
- Our model works better with long, detailed prompts since it's trained with such prompts. We will integrate prompt extension into the codebase (similar to Wan2.1) in the future. For now, it is recommended to use third-party LLMs (such as GPT-4o) to extend your prompt before providing to the model.
- You may want to adjust FPS so it plays smoothly on your device.
- The speed can be improved by enabling
torch.compile, TAEHV-VAE, or using FP8 Linear layers, although the latter two options may sacrifice quality. It is recommended to usetorch.compileif possible and enable TAEHV-VAE if further speedup is needed.
Example inference script using the chunk-wise autoregressive checkpoint trained with DMD:
python inference.py \
--config_path configs/self_forcing_dmd.yaml \
--output_folder videos/self_forcing_dmd \
--checkpoint_path checkpoints/self_forcing_dmd.pt \
--data_path prompts/MovieGenVideoBench_extended.txt \
--use_ema
Other config files and corresponding checkpoints can be found in configs folder and our huggingface repo.
huggingface-cli download gdhe17/Self-Forcing checkpoints/ode_init.pt --local-dir .
huggingface-cli download gdhe17/Self-Forcing vidprom_filtered_extended.txt --local-dir prompts
Note: Our training algorithm (except for the GAN version) is data-free (no video data is needed). For now, we directly provide the ODE initialization checkpoint and will add more instructions on how to perform ODE initialization in the future (which is identical to the process described in the CausVid repo).
torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
--rdzv_backend=c10d \
--rdzv_endpoint $MASTER_ADDR \
train.py \
--config_path configs/self_forcing_dmd.yaml \
--logdir logs/self_forcing_dmd \
--disable-wandb
Our training run uses 600 iterations and completes in under 2 hours using 64 H100 GPUs. By implementing gradient accumulation, it should be possible to reproduce the results in less than 16 hours using 8 H100 GPUs.
This codebase is built on top of the open-source implementation of CausVid by Tianwei Yin and the Wan2.1 repo.
If you find this codebase useful for your research, please kindly cite our paper:
@article{huang2025selfforcing,
title={Self Forcing: Bridging the Train-Test Gap in Autoregressive Video Diffusion},
author={Huang, Xun and Li, Zhengqi and He, Guande and Zhou, Mingyuan and Shechtman, Eli},
journal={arXiv preprint arXiv:2506.08009},
year={2025}
}