FastGen

Note: The current codebase offers early access to new, experimental features. Results reported in the README files were obtained using FastGen v0.1.0.

FastGen is a PyTorch-based framework for building fast generative models using various distillation and acceleration techniques. It supports:

• large-scale training with ≥20B parameters.
• different tasks and modalities, including T2I, I2I, I2V, V2V, multimodal video+audio generation, and causal generation.
• various distillation methods, including flow maps, distribution matching, parallel decoding distillation, self-forcing, and more.

Repository Structure

fastgen/
├── fastgen/
│   ├── callbacks/           # Training callbacks (EMA, profiling, etc.)
│   ├── configs/             # Configuration system
│   │   ├── experiments/     # Experiment configs
│   │   └── methods/         # Method-specific configs
│   ├── datasets/            # Dataset loaders
│   ├── methods/             # Trajectory matching, distribution matching, and fine-tuning
│   ├── networks/            # Neural network architectures
│   ├── third_party/         # Third-party dependencies
│   ├── trainer.py           # Main training loop
│   └── utils/               # Utilities (distributed, checkpointing)
├── comfyui/                 # ComfyUI custom nodes for interactive inference
├── scripts/                 # Inference and evaluation scripts
├── tests/                   # Unit tests
├── Makefile                 # Development commands (lint, format, test)
└── train.py                 # Main training entry point

Setup

Recommended: Use the provided Docker container for a consistent environment. See CONTRIBUTING.md for Docker setup instructions. Otherwise, create a new conda environment with conda create -y -n fastgen python=3.12.3 pip; conda activate fastgen.

Installation

git clone https://github.com/NVlabs/FastGen.git
cd FastGen
pip install -e .

On ARM64, PyPI only provides CPU builds of PyTorch, so install the CUDA build first:

pip install "torch>=2.10,<2.11" "torchvision>=0.25,<0.26" --index-url https://download.pytorch.org/whl/cu130
pip install -e .

Credentials (Optional)

For W&B logging, get your API key and save it to credentials/wandb_api.txt or set the WANDB_API_KEY environment variable. Without either of these, W&B will prompt for your API key interactively. For more details, including S3 storage and other environment variables, see fastgen/configs/README.md.

Quick Start

Before running the following commands, download the CIFAR-10 dataset and pretrained EDM models:

python scripts/download_data.py --dataset cifar10

For other datasets and models, see fastgen/networks/README.md and fastgen/datasets/README.md.

Basic Training

python train.py --config=fastgen/configs/experiments/EDM/config_dmd2_fast.py

If you run out-of-memory, try a smaller batch-size, e.g., dataloader_train.batch_size=32, which automatically uses gradient accumulation to match the global batch-size.

Expected Output: See the training log for a link to the run on wandb.ai. Training outputs go to $FASTGEN_OUTPUT_ROOT/{project}/{group}/{name}/, where name defaults to the config file's name without .py (see Run Directory). With default settings, outputs are organized as follows:

FASTGEN_OUTPUT/fastgen/cifar10/config_dmd2_fast/
├── checkpoints/    # Model checkpoints in the format {iteration:07d}.pth
│   ├── 0001000.pth
│   └── ...
├── config.yaml     # Resolved configuration for reproducibility
├── wandb_id.txt    # W&B run ID for resuming
└── ...

DDP/FSDP2 Training

For multi-GPU training, use DDP:

torchrun --nproc_per_node=8 train.py \
    --config=fastgen/configs/experiments/EDM/config_dmd2_fast.py \
    - trainer.ddp=True log_config.name=test_ddp

For large models, use FSDP2 for model sharding by replacing trainer.ddp=True with trainer.fsdp=True.

Inference

python scripts/inference/image_model_inference.py --config fastgen/configs/experiments/EDM/config_dmd2_fast.py \
  --classes=10 --prompt_file=scripts/inference/prompts/classes.txt --ckpt_path=FASTGEN_OUTPUT/fastgen/cifar10/config_dmd2_fast/checkpoints/0002000.pth - log_config.name=test_inference

For other inference modes and FID evaluations, see scripts/README.md. FastGen-distilled models can also be served with:

• FlashDreams for optimized streaming inference of Self-Forcing models.
• vLLM-Omni for DMD2 model inference.

Command-Line Overrides

Override any config parameter using Hydra-style syntax (note the - separator):

python train.py --config=path/to/config.py - key=value nested.key=value

Use +key=value for a key the config does not define yet; see fastgen/configs/README.md.

Documentation

The documentation is published at docs.nvidia.com/fastgen. Its pages are generated from the READMEs in this repository, plus the authored training and checkpoints, troubleshooting, and release notes pages. See Contributing to the docs for how to edit them. The main components and their source READMEs are:

Component	README	Description
Methods	fastgen/methods/README.md	Trajectory matching, distribution matching, and fine-tuning methods
Networks	fastgen/networks/README.md	Network architectures (Flux, Qwen-Image, Wan, LingBot-world, Cosmos, LTX-2, etc.) and pretrained models
Configs	fastgen/configs/README.md	Configuration system, environment variables, and creating custom configs
Datasets	fastgen/datasets/README.md	Dataset preparation and WebDataset loaders (image, video, video+audio)
Callbacks	fastgen/callbacks/README.md	Training callbacks (EMA, logging, gradient clipping, etc.)
Inference	scripts/README.md	Inference modes (T2I, I2I, T2V, I2V, V2V, etc.) and FID evaluation
ComfyUI	comfyui/README.md	ComfyUI custom nodes for interactive teacher/student inference
Skills	skills/README.md	Agent Skills for coding agents: setup, configuration, data, training, inference, validation, and extending FastGen
Third Party	fastgen/third_party/README.md	Third-party dependencies (Depth Anything V2, etc.)

Supported Methods

Category	Methods
Trajectory Matching	CM, ECT, sCM, TCM, MeanFlow, AYF, KD, CausalKD, PDD
Distribution Matching	DMD2, f-Distill, Phased DMD, LADD, AnyFlow, CausVid, Self-Forcing, Causal-Forcing
Fine-Tuning	SFT, CausalSFT (diffusion/teacher forcing)

See fastgen/methods/README.md for details.

Supported Networks and Data

FastGen is designed to be agnostic to the network and data and you can add your own architectures and datasets (see fastgen/networks/README.md and fastgen/datasets/README.md). For reference, we provide the following implementations:

Data	Networks
Image	EDM, EDM2, DiT, SiT, SD 1.5, SDXL, FLUX.1, FLUX.2, Qwen-Image
Video	Wan2.1, Wan2.2, Wan Vace, LingBot-world, CogVideoX, Cosmos Predict2, Cosmos 3
Video+Audio	LTX-2, LTX-2.3, LTX-2.5, MiniMax H3

See fastgen/networks/README.md for details. Not all combinations of methods and networks are currently supported. We provide typical use-cases in our predefined configs in fastgen/configs/experiments.

Adoption, Collaboration, and Usage

FastGen is used across NVIDIA and the community to build fast generative models.

Models Distilled Using FastGen

• FastGen-AYF (distilled to 4-8 steps using AYF from FastGen)
• Qwen-Image-Flash (distilled to 4 steps using DMD2 from FastGen, NVIDIA ModelOpt, and NVIDIA AutoModel)
• Baseten FLUX.2 (distilled to 8 steps using DMD2)
• Reka Real-Time Video Generation (distilled to a real-time video generator using FastGen)
• CorrDiff (distilled to 1 step using ECT)
• cBottle Super-Resolution (distilled to 2 steps using ECT)

Projects Using FastGen

• Simplifying and Improving On-Policy Distillation for Few-Step Diffusion (built upon FastGen)
• Parametric Trajectory Distillation for Few-Step Video Generation (built upon FastGen)
• Data-Forcing Distillation: Restoring Diversity and Fidelity in Few-Step Video Generation (built upon FastGen)
• Parallel Decoding Distillation for Image and Video Generation (built upon FastGen)
• Transition Matching Distillation for Fast Video Generation (built upon FastGen)
• Mode Seeking meets Mean Seeking for Fast Long Video Generation (built upon FastGen)
• StreamingEffect: Real-Time Human-Centric Video Effect Generation (built upon FastGen)
• CARV: Variance Reduction for Expectations with Diffusion Teachers (used FastGen for the DMD2 experiments)
• Flex-Forcing: Towards a Unified Autoregressive and Bidirectional Video Diffusion Model (built upon FastGen)

FastGen Collaborations

• FastH3 Omni (Minimax H3 Ref2VA distillation with FastVideo and Nuva Lab)
• FastH3 Preview (MiniMax H3 distillation with FastVideo and Nuva Lab)
• FlashDreams (optimized serving for FastGen-distilled autoregressive video models)
• NVIDIA Model Optimizer (ModelOpt) (integration of FastGen step-distillation methods)

Contributing

We welcome contributions! Please see CONTRIBUTING.md for details.

We thank everyone who has helped design, build, and test FastGen!

• Core contributors: Julius Berner, Chao Liu
• Other contributors: Weili Nie, Amirmojtaba Sabour, James Lucas, David Pankratz, Sushil Koundinyan, Sihyun Yu, Willis Ma, Yilun Xu, Shengqu Cai, Xinyin Ma, Yanke Song
• Collaborators: Sophia Zalewski, Wei Xiong, Christian Laforte, Sajad Norouzi, Kaiwen Zheng, Miloš Hašan, Saeed Hadadan, Gene Liu, David Dynerman, Alicia Sui, Grace Lam, Pooya Jannaty, Jan Kautz, and many more.
• Project lead: Arash Vahdat

License

This project is licensed under the Apache License 2.0 - see LICENSE for details. Third-party licenses are documented in licenses/README.md.

Reference

@misc{fastgen2026,
  title={NVIDIA FastGen: Fast Generation from Diffusion Models},
  author={Berner, Julius and Liu, Chao and Vahdat, Arash},
  url={https://github.com/NVlabs/FastGen},
  year={2026},
}